-`streamstore` is the Python package that provides an async client for interacting with [s2.dev](https://s2.dev/).
+`s2_sdk` is the Python package that provides an async client for interacting with [s2.dev](https://s2.dev/).
## Project links
-- [PyPI](https://pypi.org/project/streamstore/)
-- [Documentation](https://streamstore.readthedocs.io/)
+- [PyPI](https://pypi.org/project/s2-sdk/)
+- [Documentation](https://s2-sdk.readthedocs.io/)
- [GitHub](https://github.com/s2-streamstore/s2-sdk-python)
## Requirements
@@ -26,10 +26,10 @@ Python >= 3.11
## Installation
-You can install the package from the [Python Package Index](https://pypi.org/project/streamstore) using the package manager of your choice. E.g., with `pip`:
+You can install the package from the [Python Package Index](https://pypi.org/project/s2-sdk) using the package manager of your choice. E.g., with `pip`:
```bash
-pip install streamstore
+pip install s2-sdk
```
## Examples
diff --git a/docs/source/api-reference.md b/docs/source/api-reference.md
index a81775c..569a06c 100644
--- a/docs/source/api-reference.md
+++ b/docs/source/api-reference.md
@@ -1,36 +1,183 @@
# API Reference
```{eval-rst}
-.. module:: streamstore
+.. module:: s2_sdk
.. autoclass:: S2
:members:
- :member-order: bysource
-.. autoclass:: Basin()
+
+.. autoclass:: S2Basin()
+ :members:
+
+
+.. autoclass:: S2Stream()
+ :members:
+
+
+.. autoclass:: AppendSession()
+ :members:
+
+
+.. autoclass:: BatchSubmitTicket()
+ :members:
+
+.. autoclass:: Producer()
+ :members:
+
+
+.. autoclass:: RecordSubmitTicket()
+ :members:
+
+.. autoclass:: Endpoints
+ :members:
+
+
+.. autoclass:: Timeout(request: timedelta = timedelta(seconds=5), connection: timedelta = timedelta(seconds=3))
+ :members:
+
+.. autoclass:: Retry(max_attempts: int = 3, min_base_delay: timedelta = timedelta(milliseconds=100), max_base_delay: timedelta = timedelta(seconds=1), append_retry_policy: AppendRetryPolicy = AppendRetryPolicy.ALL)
+ :members:
+
+.. autoclass:: Batching(max_records: int = 1000, max_bytes: int = 1048576, linger: timedelta = timedelta(milliseconds=5))
+ :members:
+
+.. autoclass:: Record(body: bytes, headers: list[tuple[bytes, bytes]] = [], timestamp: int | None = None)
+ :members:
+
+.. autoclass:: AppendInput
+ :members:
+
+.. autoclass:: AppendAck()
+ :members:
+
+.. autoclass:: IndexedAppendAck()
+ :members:
+
+.. autoclass:: StreamPosition()
+ :members:
+
+.. autoclass:: ReadLimit
+ :members:
+
+.. autoclass:: ReadBatch()
+ :members:
+
+.. autoclass:: SequencedRecord()
+ :members:
+
+.. autoclass:: SeqNum
+ :members:
+
+.. autoclass:: Timestamp
+ :members:
+
+.. autoclass:: TailOffset
+ :members:
+
+.. autoclass:: Page()
+ :members:
+
+.. autoclass:: CommandRecord()
+ :members:
+
+
+.. autofunction:: metered_bytes
+
+.. autofunction:: append_record_batches
+
+.. autofunction:: append_inputs
+
+.. autoenum:: Compression
+
+.. autoenum:: AppendRetryPolicy
+
+.. autoenum:: StorageClass
+
+.. autoenum:: TimestampingMode
+
+.. autoclass:: Timestamping
+ :members:
+
+.. autoclass:: StreamConfig
+ :members:
+
+.. autoclass:: BasinConfig
+ :members:
+
+.. autoenum:: BasinScope
+
+.. autoclass:: BasinInfo()
+ :members:
+
+.. autoclass:: StreamInfo()
+ :members:
+
+.. autoclass:: ExactMatch
+ :members:
+
+.. autoclass:: PrefixMatch
+ :members:
+
+.. autoenum:: Permission
+
+.. autoenum:: Operation
+
+.. autoclass:: OperationGroupPermissions
+ :members:
+
+.. autoclass:: AccessTokenScope(basins: ExactMatch | PrefixMatch | None = None, streams: ExactMatch | PrefixMatch | None = None, access_tokens: ExactMatch | PrefixMatch | None = None, op_groups: OperationGroupPermissions | None = None, ops: list[Operation] = [])
+ :members:
+
+.. autoclass:: AccessTokenInfo()
+ :members:
+
+.. autoenum:: MetricUnit
+
+.. autoenum:: TimeseriesInterval
+
+.. autoenum:: AccountMetricSet
+
+.. autoenum:: BasinMetricSet
+
+.. autoenum:: StreamMetricSet
+
+.. autoclass:: Scalar()
+ :members:
+
+.. autoclass:: Accumulation()
+ :members:
+
+.. autoclass:: Gauge()
+ :members:
+
+.. autoclass:: Label()
+ :members:
+
+.. autoclass:: S2Error()
+ :members:
+
+.. autoclass:: S2ClientError()
:members:
- :member-order: bysource
+ :show-inheritance:
-.. autoclass:: Stream()
+.. autoclass:: S2ServerError()
:members:
- :member-order: bysource
+ :show-inheritance:
-.. module:: streamstore.schemas
- :no-index:
-.. autoclass:: Record(body: bytes, headers: list[tuple[bytes, bytes]] = [])
+.. autoclass:: AppendConditionError()
:members:
+ :show-inheritance:
-.. automodule:: streamstore.schemas
+.. autoclass:: FencingTokenMismatchError()
:members:
- :exclude-members: Record, Endpoints
- :member-order: bysource
+ :show-inheritance:
-.. module:: streamstore.schemas
- :no-index:
-.. autoclass:: Endpoints()
+.. autoclass:: SeqNumMismatchError()
:members:
+ :show-inheritance:
-.. automodule:: streamstore.utils
+.. autoclass:: ReadUnwrittenError()
:members:
- :member-order: bysource
+ :show-inheritance:
```
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 7bb8cbb..112a68b 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -8,9 +8,9 @@
from datetime import date
-project = "streamstore"
+project = "s2-sdk"
copyright = f"{date.today().year}, Bandar Systems Inc"
-release = "5.0.0"
+release = "0.1.0"
# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
@@ -28,6 +28,11 @@
templates_path = ["_templates"]
exclude_patterns = []
+
+autodoc_member_order = "bysource"
+autodoc_typehints_format = "short"
+python_use_unqualified_type_names = True
+
intersphinx_mapping = {"python": ("https://docs.python.org/3", None)}
# -- Options for HTML output -------------------------------------------------
diff --git a/examples/append_session.py b/examples/append_session.py
index 4ce878a..37451bf 100644
--- a/examples/append_session.py
+++ b/examples/append_session.py
@@ -1,36 +1,31 @@
import asyncio
import os
import random
-from typing import AsyncIterable
-from streamstore import S2
-from streamstore.schemas import AppendInput, Record
+from s2_sdk import S2, AppendInput, Record
ACCESS_TOKEN = os.getenv("S2_ACCESS_TOKEN")
MY_BASIN = os.getenv("MY_BASIN")
MY_STREAM = os.getenv("MY_STREAM")
-async def append_inputs_gen() -> AsyncIterable[AppendInput]:
- num_inputs = random.randint(1, 100)
- for _ in range(num_inputs):
- num_records = random.randint(1, 100)
- records = []
- for _ in range(num_records):
- body_size = random.randint(1, 1024)
- records.append(Record(body=os.urandom(body_size)))
- input = AppendInput(records)
- if random.random() < 0.5:
- await asyncio.sleep(random.random() * 2.5)
- yield input
-
-
async def producer():
- async with S2(access_token=ACCESS_TOKEN) as s2:
+ async with S2(ACCESS_TOKEN) as s2:
stream = s2[MY_BASIN][MY_STREAM]
- async for output in stream.append_session(append_inputs_gen()):
- num_appended_records = output.end_seq_num - output.start_seq_num
- print(f"appended {num_appended_records} records")
+ async with stream.append_session() as session:
+ num_inputs = random.randint(1, 100)
+ for _ in range(num_inputs):
+ num_records = random.randint(1, 100)
+ records = []
+ for _ in range(num_records):
+ body_size = random.randint(1, 1024)
+ records.append(Record(body=os.urandom(body_size)))
+ ticket = await session.submit(AppendInput(records))
+ ack = await ticket
+ num_appended_records = ack.end.seq_num - ack.start.seq_num
+ print(f"appended {num_appended_records} records")
+ if random.random() < 0.5:
+ await asyncio.sleep(random.random() * 2.5)
if __name__ == "__main__":
diff --git a/examples/append_session_with_auto_batching.py b/examples/append_session_with_auto_batching.py
index cc0f501..d3988d7 100644
--- a/examples/append_session_with_auto_batching.py
+++ b/examples/append_session_with_auto_batching.py
@@ -4,9 +4,7 @@
from datetime import timedelta
from typing import AsyncIterable
-from streamstore import S2
-from streamstore.schemas import Record
-from streamstore.utils import append_inputs_gen
+from s2_sdk import S2, Batching, Record, append_inputs
ACCESS_TOKEN = os.getenv("S2_ACCESS_TOKEN")
MY_BASIN = os.getenv("MY_BASIN")
@@ -23,17 +21,20 @@ async def records_gen() -> AsyncIterable[Record]:
async def producer():
- async with S2(access_token=ACCESS_TOKEN) as s2:
+ async with S2(ACCESS_TOKEN) as s2:
stream = s2[MY_BASIN][MY_STREAM]
- async for output in stream.append_session(
- append_inputs_gen(
+ async with stream.append_session() as session:
+ async for batch in append_inputs(
records=records_gen(),
- max_records_per_batch=10,
- max_linger_per_batch=timedelta(milliseconds=5),
- )
- ):
- num_appended_records = output.end_seq_num - output.start_seq_num
- print(f"appended {num_appended_records} records")
+ batching=Batching(
+ max_records=10,
+ linger=timedelta(milliseconds=5),
+ ),
+ ):
+ ticket = await session.submit(batch)
+ ack = await ticket
+ num_appended_records = ack.end.seq_num - ack.start.seq_num
+ print(f"appended {num_appended_records} records")
if __name__ == "__main__":
diff --git a/examples/producer.py b/examples/producer.py
new file mode 100644
index 0000000..aeefe4b
--- /dev/null
+++ b/examples/producer.py
@@ -0,0 +1,24 @@
+import asyncio
+import os
+import random
+
+from s2_sdk import S2, Record
+
+ACCESS_TOKEN = os.getenv("S2_ACCESS_TOKEN")
+MY_BASIN = os.getenv("MY_BASIN")
+MY_STREAM = os.getenv("MY_STREAM")
+
+
+async def main():
+ async with S2(ACCESS_TOKEN) as s2:
+ stream = s2[MY_BASIN][MY_STREAM]
+ async with stream.producer() as producer:
+ for i in range(100):
+ body_size = random.randint(1, 1024)
+ ticket = await producer.submit(Record(body=os.urandom(body_size)))
+ ack = await ticket
+ print(f"seq_num={ack.seq_num}")
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/examples/read_session.py b/examples/read_session.py
index 981b3c9..c03c1ab 100644
--- a/examples/read_session.py
+++ b/examples/read_session.py
@@ -1,8 +1,7 @@
import asyncio
import os
-from streamstore import S2
-from streamstore.schemas import SeqNum
+from s2_sdk import S2, SeqNum
ACCESS_TOKEN = os.getenv("S2_ACCESS_TOKEN")
MY_BASIN = os.getenv("MY_BASIN")
@@ -10,20 +9,14 @@
async def consumer():
- async with S2(access_token=ACCESS_TOKEN) as s2:
+ async with S2(ACCESS_TOKEN) as s2:
stream = s2[MY_BASIN][MY_STREAM]
tail = await stream.check_tail()
print(f"reading from tail: {tail}")
total_num_records = 0
- async for output in stream.read_session(start=SeqNum(tail.next_seq_num)):
- match output:
- case list(records):
- total_num_records += len(records)
- print(f"read {len(records)} now, {total_num_records} so far")
- case _:
- raise RuntimeError(
- "Records not received, which is unexpected as we start from the tail of the stream"
- )
+ async for batch in stream.read_session(start=SeqNum(tail.seq_num)):
+ total_num_records += len(batch.records)
+ print(f"read {len(batch.records)} now, {total_num_records} so far")
if __name__ == "__main__":
diff --git a/pyproject.toml b/pyproject.toml
index 21c4719..0ca3d4f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,26 +1,29 @@
[project]
-name = "streamstore"
-version = "5.0.0"
+name = "s2-sdk"
+version = "0.1.0"
description = "Python SDK for s2.dev"
readme = "README.md"
license = "MIT"
license-files = ["LICENSE"]
requires-python = ">=3.11"
dependencies = [
- "grpcio-tools>=1.69.0",
- "grpcio>=1.69.0",
- "types-protobuf>=5.29.1.20241207",
- "grpc-stubs>=1.53.0.5",
- "anyio>=4.8.0",
+ "h2>=4.1.0",
+ "protobuf>=5.29.0",
+ "zstandard>=0.23.0",
]
[dependency-groups]
-dev = ["mypy>=1.14.1", "poethepoet>=0.36.0", "ruff>=0.9.1"]
+dev = [
+ "grpcio-tools>=1.69.0",
+ "mypy>=1.14.1",
+ "poethepoet>=0.36.0",
+ "ruff>=0.9.1",
+ "types-protobuf>=5.29.1.20241207",
+]
test = [
"pytest>=8.0.0",
"pytest-asyncio>=0.23.0",
"pytest-timeout>=2.3.0",
- "pytest-xdist>=3.5.0",
]
docs = [
"enum-tools[sphinx]>=0.12.0",
@@ -28,14 +31,20 @@ docs = [
"myst-parser>=4.0.0",
"sphinx==8.1.3",
"sphinx-autodoc-typehints>=3.0.0",
+ "sphinx-autobuild>=2024.10.3",
]
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
+[tool.hatch.build.targets.wheel]
+packages = ["src/s2_sdk"]
+
[tool.mypy]
-files = ["src/", "tests/", "examples/"]
+files = ["src/s2_sdk/", "tests/"]
+exclude = ["src/s2_sdk/_generated/"]
+check_untyped_defs = true
[tool.ruff]
exclude = [
@@ -59,7 +68,10 @@ ci_linter = "uv run ruff check"
ci_formatter = "uv run ruff format --check"
checker = ["linter", "formatter", "type_checker"]
ci_checker = ["ci_linter", "ci_formatter", "type_checker"]
-e2e_tests = "uv run pytest tests/ -v -s"
+docs_check = "uv run --group docs sphinx-build -W -b html docs/source docs/build/html"
+docs_clean = "rm -rf docs/build"
+docs_live.sequence = ["docs_clean", "docs_check", {cmd = "uv run --group docs sphinx-autobuild docs/source docs/build/html --watch src/s2_sdk"}]
+e2e_tests = "uv run pytest tests/ -v -s -m 'account or basin or stream'"
e2e_account_tests = "uv run pytest tests/ -v -s -m account"
e2e_basin_tests = "uv run pytest tests/ -v -s -m basin"
e2e_stream_tests = "uv run pytest tests/ -v -s -m stream"
diff --git a/pytest.ini b/pytest.ini
index 5a18750..b998b45 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -16,3 +16,5 @@ markers =
account: tests for account operations
basin: tests for basin operations
stream: tests for stream operations
+ metrics: tests for metrics operations
+ cloud_only: tests that require s2-cloud (not supported on s2-lite)
diff --git a/s2-specs b/s2-specs
index 8babbd6..7d9167d 160000
--- a/s2-specs
+++ b/s2-specs
@@ -1 +1 @@
-Subproject commit 8babbd68f20829b8ec92bbe4864a3f020782bd5c
+Subproject commit 7d9167d55609c98830278bd00c66f188fdd0099d
diff --git a/src/s2_sdk/__init__.py b/src/s2_sdk/__init__.py
new file mode 100644
index 0000000..34eba20
--- /dev/null
+++ b/src/s2_sdk/__init__.py
@@ -0,0 +1,122 @@
+from s2_sdk._append_session import AppendSession, BatchSubmitTicket
+from s2_sdk._batching import append_inputs, append_record_batches
+from s2_sdk._exceptions import (
+ AppendConditionError,
+ FencingTokenMismatchError,
+ ReadUnwrittenError,
+ S2ClientError,
+ S2Error,
+ S2ServerError,
+ SeqNumMismatchError,
+)
+from s2_sdk._ops import S2, S2Basin, S2Stream
+from s2_sdk._producer import Producer, RecordSubmitTicket
+from s2_sdk._types import (
+ AccessTokenInfo,
+ AccessTokenScope,
+ AccountMetricSet,
+ Accumulation,
+ AppendAck,
+ AppendInput,
+ AppendRetryPolicy,
+ BasinConfig,
+ BasinInfo,
+ BasinMetricSet,
+ BasinScope,
+ Batching,
+ CommandRecord,
+ Compression,
+ Endpoints,
+ ExactMatch,
+ Gauge,
+ IndexedAppendAck,
+ Label,
+ MetricUnit,
+ Operation,
+ OperationGroupPermissions,
+ Page,
+ Permission,
+ PrefixMatch,
+ ReadBatch,
+ ReadLimit,
+ Record,
+ Retry,
+ Scalar,
+ SeqNum,
+ SequencedRecord,
+ StorageClass,
+ StreamConfig,
+ StreamInfo,
+ StreamMetricSet,
+ StreamPosition,
+ TailOffset,
+ Timeout,
+ TimeseriesInterval,
+ Timestamp,
+ Timestamping,
+ TimestampingMode,
+ metered_bytes,
+)
+
+__all__ = [
+ "S2",
+ "S2Basin",
+ "S2Stream",
+ "AppendSession",
+ "BatchSubmitTicket",
+ "Producer",
+ "RecordSubmitTicket",
+ "append_record_batches",
+ "append_inputs",
+ "Retry",
+ "Timeout",
+ "Compression",
+ "AppendRetryPolicy",
+ "Batching",
+ "Endpoints",
+ "Record",
+ "AppendInput",
+ "AppendAck",
+ "IndexedAppendAck",
+ "StreamPosition",
+ "SeqNum",
+ "Timestamp",
+ "TailOffset",
+ "ReadBatch",
+ "ReadLimit",
+ "SequencedRecord",
+ "Page",
+ "CommandRecord",
+ "metered_bytes",
+ "StorageClass",
+ "TimestampingMode",
+ "Timestamping",
+ "StreamConfig",
+ "BasinConfig",
+ "BasinInfo",
+ "BasinScope",
+ "StreamInfo",
+ "ExactMatch",
+ "PrefixMatch",
+ "Permission",
+ "Operation",
+ "OperationGroupPermissions",
+ "AccessTokenScope",
+ "AccessTokenInfo",
+ "MetricUnit",
+ "TimeseriesInterval",
+ "AccountMetricSet",
+ "BasinMetricSet",
+ "StreamMetricSet",
+ "Scalar",
+ "Accumulation",
+ "Gauge",
+ "Label",
+ "S2Error",
+ "S2ClientError",
+ "S2ServerError",
+ "AppendConditionError",
+ "FencingTokenMismatchError",
+ "SeqNumMismatchError",
+ "ReadUnwrittenError",
+]
diff --git a/src/s2_sdk/_append_session.py b/src/s2_sdk/_append_session.py
new file mode 100644
index 0000000..bd9cd17
--- /dev/null
+++ b/src/s2_sdk/_append_session.py
@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import asyncio
+from collections import deque
+from typing import AsyncIterable, NamedTuple, Self
+
+from s2_sdk._client import HttpClient
+from s2_sdk._exceptions import S2ClientError
+from s2_sdk._s2s._append_session import run_append_session
+from s2_sdk._types import (
+ AppendAck,
+ AppendInput,
+ Compression,
+ Retry,
+ metered_bytes,
+)
+from s2_sdk._validators import validate_append_input
+
+
+class _UnackedBatch(NamedTuple):
+ ack_fut: asyncio.Future[AppendAck]
+ metered_bytes: int
+
+
+class AppendSession:
+ """Session for high-throughput appending with backpressure control.
+
+ Supports pipelining multiple :class:`AppendInput`\ s while preserving
+ submission order.
+
+ Caution:
+ Returned by :meth:`S2Stream.append_session`. Do not instantiate directly.
+ """
+
+ __slots__ = (
+ "_closed",
+ "_client",
+ "_compression",
+ "_error",
+ "_permits",
+ "_queue",
+ "_retry",
+ "_stream_name",
+ "_task",
+ "_unacked",
+ )
+
+ def __init__(
+ self,
+ client: HttpClient,
+ stream_name: str,
+ retry: Retry,
+ compression: Compression,
+ max_unacked_bytes: int,
+ max_unacked_batches: int | None,
+ ) -> None:
+ self._client = client
+ self._stream_name = stream_name
+ self._retry = retry
+ self._compression = compression
+ self._permits = _AppendPermits(max_unacked_bytes, max_unacked_batches)
+
+ self._queue: asyncio.Queue[AppendInput | None] = asyncio.Queue()
+ self._unacked: deque[_UnackedBatch] = deque()
+ self._closed = False
+ self._error: BaseException | None = None
+
+ self._task = asyncio.get_running_loop().create_task(self._run())
+
+ async def submit(self, inp: AppendInput) -> BatchSubmitTicket:
+ """Submit a batch of records for appending.
+
+ Waits when backpressure limits are reached.
+ """
+ self._check_ready()
+ batch_bytes = metered_bytes(inp.records)
+ validate_append_input(len(inp.records), batch_bytes)
+
+ await self._permits.acquire(batch_bytes)
+ # Re-check after potentially waiting on backpressure.
+ try:
+ self._check_ready()
+ except BaseException:
+ self._permits.release(batch_bytes)
+ raise
+
+ ack_fut: asyncio.Future[AppendAck] = asyncio.get_running_loop().create_future()
+ self._unacked.append(_UnackedBatch(ack_fut, batch_bytes))
+ await self._queue.put(inp)
+ return BatchSubmitTicket(ack_fut)
+
+ def _check_ready(self) -> None:
+ if self._closed:
+ raise S2ClientError("AppendSession is closed")
+ if self._error is not None:
+ raise self._error
+
+ async def close(self) -> None:
+ """Close the session and wait for all submitted batches to be appended."""
+ if self._closed:
+ return
+ self._closed = True
+ await self._queue.put(None)
+ await self._task
+ if self._error is not None:
+ raise self._error
+
+ async def __aenter__(self) -> Self:
+ return self
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
+ await self.close()
+ return False
+
+ async def _run(self) -> None:
+ try:
+ async for ack in run_append_session(
+ self._client,
+ self._stream_name,
+ self._input_iter(),
+ retry=self._retry,
+ compression=self._compression,
+ ack_timeout=self._client._request_timeout,
+ ):
+ self._resolve_next(ack)
+ except BaseException as e:
+ # Unwrap single-exception ExceptionGroups so callers see
+ # the original exception type (e.g. S2ServerError, SeqNumMismatchError).
+ exc = e
+ while isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1:
+ exc = exc.exceptions[0]
+ self._fail_all(exc)
+
+ async def _input_iter(self) -> AsyncIterable[AppendInput]:
+ while True:
+ item = await self._queue.get()
+ if item is None:
+ return
+ yield item
+
+ def _resolve_next(self, ack: AppendAck) -> None:
+ unacked = self._unacked.popleft()
+ self._permits.release(unacked.metered_bytes)
+ unacked.ack_fut.set_result(ack)
+
+ def _fail_all(self, error: BaseException) -> None:
+ self._error = error
+ for unacked in self._unacked:
+ self._permits.release(unacked.metered_bytes)
+ if not unacked.ack_fut.done():
+ unacked.ack_fut.set_exception(error)
+ self._unacked.clear()
+ # Drain queue
+ while not self._queue.empty():
+ try:
+ self._queue.get_nowait()
+ except asyncio.QueueEmpty:
+ break
+
+
+class BatchSubmitTicket:
+ """Awaitable that resolves to an :class:`AppendAck` once the batch is appended."""
+
+ __slots__ = ("_ack_fut",)
+
+ def __init__(self, ack_fut: asyncio.Future[AppendAck]) -> None:
+ self._ack_fut = ack_fut
+
+ def __await__(self):
+ return self._ack_fut.__await__()
+
+
+class _Semaphore:
+ __slots__ = ("_event", "_lock", "_value")
+
+ def __init__(self, value: int) -> None:
+ self._value = value
+ self._event = asyncio.Event()
+ self._event.set()
+ self._lock = asyncio.Lock()
+
+ async def acquire(self, n: int) -> None:
+ while True:
+ async with self._lock:
+ if self._value >= n:
+ self._value -= n
+ return
+ self._event.clear()
+ if self._value >= n:
+ continue
+ await self._event.wait()
+
+ def release(self, n: int) -> None:
+ self._value += n
+ self._event.set()
+
+
+class _AppendPermits:
+ __slots__ = ("_bytes", "_count")
+
+ def __init__(self, max_bytes: int, max_count: int | None = None) -> None:
+ self._bytes = _Semaphore(max_bytes)
+ self._count = _Semaphore(max_count) if max_count is not None else None
+
+ async def acquire(self, n_bytes: int) -> None:
+ if self._count is not None:
+ await self._count.acquire(1)
+ try:
+ await self._bytes.acquire(n_bytes)
+ except BaseException:
+ if self._count is not None:
+ self._count.release(1)
+ raise
+
+ def release(self, n_bytes: int) -> None:
+ self._bytes.release(n_bytes)
+ if self._count is not None:
+ self._count.release(1)
diff --git a/src/s2_sdk/_batching.py b/src/s2_sdk/_batching.py
new file mode 100644
index 0000000..dccad1a
--- /dev/null
+++ b/src/s2_sdk/_batching.py
@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import asyncio
+from typing import AsyncIterable
+
+from s2_sdk._types import AppendInput, Batching, Record, metered_bytes
+from s2_sdk._validators import validate_batching
+
+
+class BatchAccumulator:
+ __slots__ = ("_batching", "_bytes", "_records")
+
+ def __init__(self, batching: Batching) -> None:
+ self._batching = batching
+ self._records: list[Record] = []
+ self._bytes = 0
+
+ def add(self, record: Record) -> None:
+ self._records.append(record)
+ self._bytes += metered_bytes((record,))
+
+ def take(self) -> list[Record]:
+ records = list(self._records)
+ self._records.clear()
+ self._bytes = 0
+ return records
+
+ def is_full(self) -> bool:
+ return (
+ len(self._records) >= self._batching.max_records
+ or self._bytes >= self._batching.max_bytes
+ )
+
+ def is_empty(self) -> bool:
+ return len(self._records) == 0
+
+ @property
+ def linger(self) -> float:
+ return self._batching.linger.total_seconds()
+
+
+async def append_record_batches(
+ records: AsyncIterable[Record],
+ *,
+ batching: Batching | None = None,
+) -> AsyncIterable[list[Record]]:
+ """Group records into batches based on count, bytes, and linger time."""
+ if batching is None:
+ batching = Batching()
+ validate_batching(batching.max_records, batching.max_bytes)
+ acc = BatchAccumulator(batching)
+ linger_secs = batching.linger.total_seconds()
+ aiter = records.__aiter__()
+
+ while True:
+ try:
+ record = await anext(aiter)
+ except StopAsyncIteration:
+ break
+
+ acc.add(record)
+ if acc.is_full():
+ yield acc.take()
+ continue
+
+ try:
+ while not acc.is_full():
+ if linger_secs > 0:
+ record = await asyncio.wait_for(anext(aiter), timeout=linger_secs)
+ else:
+ record = await anext(aiter)
+ acc.add(record)
+ except StopAsyncIteration:
+ pass
+ except TimeoutError:
+ pass
+
+ if not acc.is_empty():
+ yield acc.take()
+
+
+async def append_inputs(
+ records: AsyncIterable[Record],
+ *,
+ match_seq_num: int | None = None,
+ fencing_token: str | None = None,
+ batching: Batching | None = None,
+) -> AsyncIterable[AppendInput]:
+ """Group records into :class:`AppendInput` batches based on count, bytes, and linger time.
+
+ If ``match_seq_num`` is set, it applies to the first input and is auto-incremented for subsequent ones.
+ """
+ if batching is None:
+ batching = Batching()
+ async for batch in append_record_batches(records, batching=batching):
+ if not batch:
+ continue
+ append_input = AppendInput(
+ records=batch,
+ match_seq_num=match_seq_num,
+ fencing_token=fencing_token,
+ )
+ if match_seq_num is not None:
+ match_seq_num += len(batch)
+ yield append_input
diff --git a/src/s2_sdk/_client.py b/src/s2_sdk/_client.py
new file mode 100644
index 0000000..e139a05
--- /dev/null
+++ b/src/s2_sdk/_client.py
@@ -0,0 +1,1070 @@
+from __future__ import annotations
+
+import asyncio
+import json as json_lib
+import ssl
+import time
+from collections.abc import AsyncGenerator, AsyncIterator, Callable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from importlib.metadata import version
+from typing import Any, NamedTuple
+from urllib.parse import urlencode, urlsplit
+
+import h2.config
+import h2.connection
+import h2.events
+
+from s2_sdk._compression import compress, decompress
+from s2_sdk._exceptions import (
+ UNKNOWN_CODE,
+ ConnectError,
+ ConnectionClosedError,
+ ProtocolError,
+ ReadTimeoutError,
+ S2ClientError,
+ S2ServerError,
+ TransportError,
+ raise_for_412,
+ raise_for_416,
+)
+from s2_sdk._types import Compression
+
+_VERSION = version("s2-sdk")
+_USER_AGENT = f"s2-sdk-python/{_VERSION}"
+
+DEFAULT_MAX_STREAMS_PER_CONN = 100
+IDLE_TIMEOUT = 90.0
+REAPER_INTERVAL = 30.0
+_DRAIN_BUFFER_THRESHOLD = 65536 # 64 KiB — drain when write buffer exceeds this
+
+
+_COMPRESSION_ENCODING = {
+ Compression.ZSTD: "zstd",
+ Compression.GZIP: "gzip",
+}
+
+_ENCODING_COMPRESSION = {v: k for k, v in _COMPRESSION_ENCODING.items()}
+
+
+class HttpClient:
+ __slots__ = (
+ "_base_url",
+ "_pool",
+ "_request_timeout",
+ "_scheme",
+ "_authority",
+ "_headers",
+ "_compression",
+ )
+
+ def __init__(
+ self,
+ pool: ConnectionPool,
+ base_url: str,
+ request_timeout: float,
+ headers: dict[str, str] | None = None,
+ compression: Compression = Compression.NONE,
+ ) -> None:
+ self._pool = pool
+ self._base_url = base_url
+ self._request_timeout = request_timeout
+ origin = _origin(base_url)
+ self._scheme = origin.scheme
+ default_port = 443 if origin.scheme == "https" else 80
+ self._authority = (
+ origin.host
+ if origin.port == default_port
+ else f"{origin.host}:{origin.port}"
+ )
+ self._headers = headers
+ self._compression = compression
+
+ async def unary_request(
+ self,
+ method: str,
+ path: str,
+ *,
+ json: Any = None,
+ params: dict[str, Any] | None = None,
+ headers: dict[str, str] | None = None,
+ content: bytes | None = None,
+ ) -> Response:
+ pc, state = await self._pool.checkout(self._base_url)
+ conn = pc._conn
+ stream_id: int | None = None
+ try:
+ url_path = _build_path(path, params)
+ h2_headers = self._build_headers(method, url_path, headers)
+
+ # Build body
+ body: bytes | None = None
+ if json is not None:
+ body = json_lib.dumps(json).encode("utf-8")
+ h2_headers.append(("content-type", "application/json"))
+ elif content is not None:
+ body = content
+
+ # Compress request body
+ if body is not None and self._compression != Compression.NONE:
+ body = compress(body, self._compression)
+ h2_headers.append(
+ ("content-encoding", _COMPRESSION_ENCODING[self._compression])
+ )
+
+ if body is not None:
+ h2_headers.append(("content-length", str(len(body))))
+
+ end_stream = body is None
+ stream_id = await conn.send_headers(
+ state, h2_headers, end_stream=end_stream
+ )
+
+ if body is not None:
+ assert stream_id is not None
+ await conn.send_data(stream_id, body, end_stream=True)
+
+ # Wait for response headers
+ resp_headers = await asyncio.wait_for(
+ state.response_headers,
+ timeout=self._request_timeout,
+ )
+ status_code = _status_from_headers(resp_headers)
+
+ # Read full body
+ chunks: list[bytes] = []
+ while True:
+ item = await asyncio.wait_for(
+ state.data_queue.get(),
+ timeout=self._request_timeout,
+ )
+ if item is None:
+ break
+ chunk, flow_bytes = _queue_item_parts(item)
+ chunks.append(chunk)
+ if flow_bytes > 0:
+ state.unacked_flow_bytes -= flow_bytes
+ await conn.ack_data(stream_id, flow_bytes)
+
+ if state.error is not None and not state.end_stream_received:
+ raise state.error
+
+ resp_body = b"".join(chunks)
+
+ # Decompress response body
+ content_encoding = _header_value(resp_headers, "content-encoding")
+ if (
+ content_encoding is not None
+ and content_encoding in _ENCODING_COMPRESSION
+ ):
+ resp_body = decompress(
+ resp_body, _ENCODING_COMPRESSION[content_encoding]
+ )
+
+ response = Response(status_code, resp_body)
+ except TransportError:
+ raise
+ except asyncio.TimeoutError:
+ raise ReadTimeoutError("Request timed out")
+ finally:
+ if stream_id is not None:
+ nbytes = _take_all_unacked_flow_bytes(state)
+ if nbytes > 0:
+ try:
+ await conn.ack_data(stream_id, nbytes)
+ except Exception:
+ pass
+ if not state.ended.is_set():
+ await conn.reset_stream(stream_id)
+ conn.release_stream(stream_id, state)
+ pc.touch_idle()
+
+ retry_after_ms = _header_value(resp_headers, "retry-after-ms")
+ _raise_for_status(response, retry_after_ms=retry_after_ms)
+ return response
+
+ @asynccontextmanager
+ async def streaming_request(
+ self,
+ method: str,
+ path: str,
+ *,
+ params: dict[str, Any] | None = None,
+ headers: dict[str, str] | None = None,
+ content: Any = None,
+ frame_signal: Any = None,
+ ) -> AsyncIterator[StreamingResponse]:
+ pc, state = await self._pool.checkout(self._base_url)
+ conn = pc._conn
+ stream_id: int | None = None
+ send_task: asyncio.Task[None] | None = None
+ try:
+ url_path = _build_path(path, params)
+ h2_headers = self._build_headers(method, url_path, headers)
+
+ has_body = content is not None
+ stream_id = await conn.send_headers(
+ state, h2_headers, end_stream=not has_body
+ )
+
+ if has_body:
+ on_write = frame_signal.signal if frame_signal is not None else None
+ send_task = asyncio.get_running_loop().create_task(
+ # stream_id is assigned by send_headers above.
+ _drain_body(conn, stream_id, content, on_write)
+ )
+
+ # Propagate sender errors to stream state so the response
+ # reader doesn't hang forever on data_queue.get().
+ def _on_send_done(
+ task: asyncio.Task[None], _state: Any = state
+ ) -> None:
+ if task.cancelled():
+ return
+ exc = task.exception()
+ if exc is not None:
+ conn._fail_stream(_state, exc)
+
+ assert send_task is not None
+ send_task.add_done_callback(_on_send_done)
+
+ # Wait for response headers
+ resp_headers = await asyncio.wait_for(
+ state.response_headers,
+ timeout=self._request_timeout,
+ )
+ status_code = _status_from_headers(resp_headers)
+
+ async def _ack_stream_data(nbytes: int) -> None:
+ assert stream_id is not None
+ await conn.ack_data(stream_id, nbytes)
+
+ response = StreamingResponse(
+ status_code=status_code,
+ data_queue=state.data_queue,
+ ended=state.ended,
+ stream_state=state,
+ ack=_ack_stream_data,
+ )
+ yield response
+ except TransportError:
+ raise
+ except asyncio.TimeoutError:
+ raise ReadTimeoutError("Streaming request timed out")
+ finally:
+ if send_task is not None and not send_task.done():
+ send_task.cancel()
+ try:
+ await send_task
+ except (asyncio.CancelledError, Exception):
+ pass
+ # Ack remaining flow bytes to keep connection window healthy
+ if stream_id is not None:
+ nbytes = _take_all_unacked_flow_bytes(state)
+ if nbytes > 0:
+ try:
+ await conn.ack_data(stream_id, nbytes)
+ except Exception:
+ pass
+ if not state.ended.is_set():
+ await conn.reset_stream(stream_id)
+ conn.release_stream(stream_id, state)
+ pc.touch_idle()
+
+ def _build_headers(
+ self,
+ method: str,
+ url_path: str,
+ extra_headers: dict[str, str] | None = None,
+ ) -> list[tuple[str, str]]:
+ h = [
+ (":method", method),
+ (":path", url_path),
+ (":scheme", self._scheme),
+ (":authority", self._authority),
+ ("user-agent", _USER_AGENT),
+ ]
+ if self._compression != Compression.NONE:
+ h.append(("accept-encoding", _COMPRESSION_ENCODING[self._compression]))
+ if self._headers:
+ for k, v in self._headers.items():
+ h.append((k.lower(), v))
+ if extra_headers:
+ for k, v in extra_headers.items():
+ h.append((k.lower(), v))
+ return h
+
+
+class ConnectionPool:
+ __slots__ = (
+ "_closed",
+ "_connect_timeout",
+ "_hosts",
+ "_host_locks",
+ "_reaper_task",
+ "_ssl_context",
+ )
+
+ def __init__(self, connect_timeout: float) -> None:
+ self._connect_timeout = connect_timeout
+ self._hosts: dict[str, list[_PooledConnection]] = {}
+ self._host_locks: dict[str, asyncio.Lock] = {}
+ self._reaper_task: asyncio.Task[None] | None = None
+ self._closed = False
+ self._ssl_context = ssl.create_default_context()
+ self._ssl_context.set_alpn_protocols(["h2"])
+
+ async def checkout(self, base_url: str) -> _Checkout:
+ if self._closed:
+ raise S2ClientError("Pool is closed")
+
+ self._ensure_reaper()
+
+ result = self._try_checkout(base_url)
+ if result is not None:
+ return result
+
+ lock = self._host_locks.get(base_url)
+ if lock is None:
+ lock = asyncio.Lock()
+ self._host_locks[base_url] = lock
+
+ async with lock:
+ # Re-check after acquiring lock — another caller may have
+ # created a connection while we waited.
+ result = self._try_checkout(base_url)
+ if result is not None:
+ return result
+
+ scheme, host, port = _origin(base_url)
+ use_ssl = self._ssl_context if scheme == "https" else None
+
+ conn = Connection(
+ host=host,
+ port=port,
+ ssl_context=use_ssl,
+ connect_timeout=self._connect_timeout,
+ )
+ await conn.connect()
+
+ # Wait briefly for the server's initial SETTINGS frame so
+ # max_concurrent_streams reflects the real limit.
+ try:
+ await asyncio.wait_for(
+ conn._settings_received.wait(),
+ timeout=self._connect_timeout,
+ )
+ except asyncio.TimeoutError:
+ pass # Proceed with h2 defaults
+ if conn._recv_dead:
+ await conn.close()
+ raise ConnectError(
+ f"Connection to {host}:{port} closed before HTTP/2 SETTINGS"
+ )
+
+ pc = _PooledConnection(conn)
+ conns = self._hosts.get(base_url)
+ if conns is None:
+ conns = [pc]
+ self._hosts[base_url] = conns
+ else:
+ conns.append(pc)
+ if conn._settings_received.is_set() and conn.max_concurrent_streams <= 0:
+ await pc.close()
+ conns.remove(pc)
+ raise ProtocolError("Connection has no available stream capacity")
+ state = pc._conn.reserve_stream()
+ return _Checkout(pc, state)
+
+ def _try_checkout(self, base_url: str) -> _Checkout | None:
+ conns = self._hosts.get(base_url)
+ if conns is not None:
+ for pc in conns:
+ if pc.has_capacity:
+ state = pc._conn.reserve_stream()
+ return _Checkout(pc, state)
+ return None
+
+ def _ensure_reaper(self) -> None:
+ if self._reaper_task is None or self._reaper_task.done():
+ self._reaper_task = asyncio.get_running_loop().create_task(
+ self._reap_idle()
+ )
+
+ async def _reap_idle(self) -> None:
+ while not self._closed:
+ await asyncio.sleep(REAPER_INTERVAL)
+ empty_hosts: list[str] = []
+ for base_url, conns in tuple(self._hosts.items()):
+ to_close: list[_PooledConnection] = []
+ for pc in conns:
+ if not pc._conn.is_available:
+ to_close.append(pc)
+ elif (
+ pc.is_idle
+ and pc.idle_for() > IDLE_TIMEOUT
+ and len(conns) - len(to_close) > 1
+ ):
+ to_close.append(pc)
+ for pc in to_close:
+ conns.remove(pc)
+ for pc in to_close:
+ await pc.close()
+ if not conns:
+ empty_hosts.append(base_url)
+ for base_url in empty_hosts:
+ self._hosts.pop(base_url, None)
+ lock = self._host_locks.get(base_url)
+ if lock is not None and not lock.locked():
+ self._host_locks.pop(base_url, None)
+
+ async def close(self) -> None:
+ self._closed = True
+ if self._reaper_task is not None:
+ self._reaper_task.cancel()
+ try:
+ await self._reaper_task
+ except asyncio.CancelledError:
+ pass
+ for conns in self._hosts.values():
+ for pc in conns:
+ await pc.close()
+ self._hosts.clear()
+ self._host_locks.clear()
+
+
+class Response:
+ __slots__ = ("status_code", "_content")
+
+ def __init__(self, status_code: int, content: bytes) -> None:
+ self.status_code = status_code
+ self._content = content
+
+ @property
+ def content(self) -> bytes:
+ return self._content
+
+ @property
+ def text(self) -> str:
+ return self._content.decode("utf-8", errors="replace")
+
+ def json(self) -> Any:
+ return json_lib.loads(self._content)
+
+
+class StreamingResponse:
+ __slots__ = (
+ "status_code",
+ "_data_queue",
+ "_ended",
+ "_stream_state",
+ "_buf",
+ "_ack",
+ )
+
+ def __init__(
+ self,
+ status_code: int,
+ data_queue: asyncio.Queue[tuple[bytes, int] | bytes | None],
+ ended: asyncio.Event,
+ stream_state: Any,
+ ack: Callable[[int], Any] | None = None,
+ ) -> None:
+ self.status_code = status_code
+ self._data_queue = data_queue
+ self._ended = ended
+ self._stream_state = stream_state
+ self._buf = bytearray()
+ self._ack = ack
+
+ async def aread(self) -> bytes:
+ chunks: list[bytes] = []
+ if self._buf:
+ chunks.append(bytes(self._buf))
+ self._buf.clear()
+ while True:
+ item = await self._data_queue.get()
+ if item is None:
+ break
+ chunk, flow_bytes = _queue_item_parts(item)
+ chunks.append(chunk)
+ if self._ack is not None and flow_bytes > 0:
+ self._stream_state.unacked_flow_bytes -= flow_bytes
+ await self._ack(flow_bytes)
+ if (
+ self._stream_state.error is not None
+ and not self._stream_state.end_stream_received
+ ):
+ raise self._stream_state.error
+ return b"".join(chunks)
+
+ async def aiter_bytes(self) -> AsyncGenerator[bytes, None]:
+ if self._buf:
+ yield bytes(self._buf)
+ self._buf.clear()
+ while True:
+ item = await self._data_queue.get()
+ if item is None:
+ if (
+ self._stream_state.error is not None
+ and not self._stream_state.end_stream_received
+ ):
+ raise self._stream_state.error
+ return
+ chunk, flow_bytes = _queue_item_parts(item)
+ try:
+ yield chunk
+ finally:
+ if self._ack is not None and flow_bytes > 0:
+ self._stream_state.unacked_flow_bytes -= flow_bytes
+ await self._ack(flow_bytes)
+
+
+class _PooledConnection:
+ __slots__ = ("_conn", "_idle_since")
+
+ def __init__(self, conn: Connection) -> None:
+ self._conn = conn
+ self._idle_since: float | None = None
+
+ @property
+ def has_capacity(self) -> bool:
+ return (
+ self._conn.is_available
+ and self._conn._settings_received.is_set()
+ and self._conn.open_stream_count < self._conn.max_concurrent_streams
+ )
+
+ @property
+ def is_idle(self) -> bool:
+ return self._conn.open_stream_count == 0
+
+ def touch_idle(self) -> None:
+ if self._conn.open_stream_count == 0:
+ self._idle_since = time.monotonic()
+
+ def idle_for(self) -> float:
+ if self._idle_since is None:
+ return 0.0
+ return time.monotonic() - self._idle_since
+
+ async def close(self) -> None:
+ await self._conn.close()
+
+
+@dataclass
+class _StreamState:
+ response_headers: asyncio.Future[list[tuple[str, str]]] = field(
+ default_factory=lambda: asyncio.get_running_loop().create_future()
+ )
+ data_queue: asyncio.Queue[tuple[bytes, int] | bytes | None] = field(
+ default_factory=asyncio.Queue
+ )
+ ended: asyncio.Event = field(default_factory=asyncio.Event)
+ window_updated: asyncio.Event = field(default_factory=asyncio.Event)
+ error: BaseException | None = None
+ end_stream_received: bool = False
+ unacked_flow_bytes: int = 0
+
+
+class Connection:
+ __slots__ = (
+ "_host",
+ "_port",
+ "_ssl_context",
+ "_connect_timeout",
+ "_reader",
+ "_writer",
+ "_h2",
+ "_write_lock",
+ "_streams",
+ "_pending_streams",
+ "_recv_task",
+ "_closed",
+ "_goaway_received",
+ "_recv_dead",
+ "_settings_received",
+ )
+
+ def __init__(
+ self,
+ host: str,
+ port: int,
+ ssl_context: ssl.SSLContext | None,
+ connect_timeout: float,
+ ) -> None:
+ self._host = host
+ self._port = port
+ self._ssl_context = ssl_context
+ self._connect_timeout = connect_timeout
+ self._reader: asyncio.StreamReader | None = None
+ self._writer: asyncio.StreamWriter | None = None
+ self._h2: h2.connection.H2Connection | None = None
+ self._write_lock = asyncio.Lock()
+ self._streams: dict[int, _StreamState] = {}
+ self._pending_streams: dict[int, _StreamState] = {}
+ self._recv_task: asyncio.Task[None] | None = None
+ self._closed = False
+ self._goaway_received = False
+ self._recv_dead = False
+ self._settings_received = asyncio.Event()
+
+ async def connect(self) -> None:
+ try:
+ self._reader, self._writer = await asyncio.wait_for(
+ asyncio.open_connection(
+ self._host,
+ self._port,
+ ssl=self._ssl_context,
+ ),
+ timeout=self._connect_timeout,
+ )
+ except asyncio.TimeoutError:
+ raise ConnectError(f"Connection to {self._host}:{self._port} timed out")
+ except OSError as e:
+ raise ConnectError(str(e)) from e
+
+ if self._ssl_context is not None:
+ assert self._writer is not None
+ ssl_object = self._writer.get_extra_info("ssl_object")
+ if ssl_object is not None:
+ alpn = ssl_object.selected_alpn_protocol()
+ if alpn != "h2":
+ self._writer.close()
+ raise ConnectError(
+ f"ALPN negotiation failed: expected 'h2', got {alpn!r}"
+ )
+
+ config = h2.config.H2Configuration(client_side=True, header_encoding="utf-8")
+ self._h2 = h2.connection.H2Connection(config=config)
+ self._h2.initiate_connection()
+ self._flush_h2_data_sync()
+ assert self._writer is not None
+ await self._writer.drain()
+ self._recv_task = asyncio.get_running_loop().create_task(self._recv_loop())
+
+ def reserve_stream(self) -> _StreamState:
+ """Reserve per-connection capacity for a future outbound stream."""
+ state = _StreamState()
+ self._pending_streams[id(state)] = state
+ return state
+
+ async def send_headers(
+ self,
+ state: _StreamState,
+ headers: list[tuple[str, str]],
+ end_stream: bool = False,
+ ) -> int:
+ """Atomically allocate a stream ID and send request headers."""
+ assert self._h2 is not None
+ async with self._write_lock:
+ if state.error is not None:
+ raise state.error
+ pending_state = self._pending_streams.pop(id(state), None)
+ if pending_state is None:
+ raise ProtocolError("Stream reservation missing")
+ stream_id = self._h2.get_next_available_stream_id()
+ self._streams[stream_id] = state
+ try:
+ self._h2.send_headers(stream_id, headers, end_stream=end_stream)
+ await self._flush_h2_data_and_drain()
+ except Exception:
+ self._streams.pop(stream_id, None)
+ raise
+ return stream_id
+
+ async def send_data(
+ self,
+ stream_id: int,
+ data: bytes,
+ end_stream: bool = False,
+ on_write: Callable[[], None] | None = None,
+ ) -> None:
+ """Send data on a stream, respecting flow control windows.
+
+ Chunks by flow control window + max frame size. If ``on_write`` is
+ provided, it is called after each chunk is flushed to the socket.
+ """
+ assert self._h2 is not None
+ offset = 0
+ while offset < len(data):
+ state = self._streams.get(stream_id)
+ if state and state.error:
+ raise state.error
+
+ # Window check + send in a single lock acquisition to avoid TOCTOU.
+ sent = False
+ async with self._write_lock:
+ window = self._h2.local_flow_control_window(stream_id)
+ if window > 0:
+ max_frame = self._h2.max_outbound_frame_size
+ chunk_size = min(len(data) - offset, window, max_frame)
+ chunk = data[offset : offset + chunk_size]
+ is_last_chunk = offset + chunk_size >= len(data)
+
+ self._h2.send_data(
+ stream_id,
+ chunk,
+ end_stream=end_stream and is_last_chunk,
+ )
+ await self._flush_h2_data()
+ offset += chunk_size
+ sent = True
+
+ if sent:
+ if on_write:
+ on_write()
+ continue
+
+ # Window exhausted — wait for update (lock released).
+ # Caller-level timeouts (unary or streaming) will cancel this
+ # if it takes too long.
+ state = self._streams.get(stream_id)
+ if state:
+ state.window_updated.clear()
+ # Re-check under lock to avoid missing an update.
+ async with self._write_lock:
+ window = self._h2.local_flow_control_window(stream_id)
+ if window <= 0:
+ await state.window_updated.wait()
+
+ # Handle empty data with end_stream
+ if not data and end_stream:
+ async with self._write_lock:
+ self._h2.send_data(stream_id, b"", end_stream=True)
+ await self._flush_h2_data_and_drain()
+ if on_write:
+ on_write()
+
+ async def end_stream(self, stream_id: int) -> None:
+ """Send END_STREAM on a stream."""
+ assert self._h2 is not None
+ async with self._write_lock:
+ self._h2.send_data(stream_id, b"", end_stream=True)
+ await self._flush_h2_data_and_drain()
+
+ async def ack_data(self, stream_id: int, nbytes: int) -> None:
+ """Acknowledge received data to update the flow control window."""
+ assert self._h2 is not None
+ async with self._write_lock:
+ self._h2.acknowledge_received_data(nbytes, stream_id)
+ await self._flush_h2_data()
+
+ async def reset_stream(self, stream_id: int) -> None:
+ """Send RST_STREAM to tell the peer to stop sending."""
+ assert self._h2 is not None
+ try:
+ async with self._write_lock:
+ self._h2.reset_stream(stream_id)
+ await self._flush_h2_data_and_drain()
+ except Exception:
+ pass # Best effort
+
+ def release_stream(
+ self, stream_id: int | None, state: _StreamState | None = None
+ ) -> None:
+ """Clean up active or reserved stream state."""
+ if stream_id is not None:
+ self._streams.pop(stream_id, None)
+ if state is not None:
+ self._pending_streams.pop(id(state), None)
+
+ @property
+ def is_available(self) -> bool:
+ """Connection is usable for new streams."""
+ return not self._closed and not self._goaway_received and not self._recv_dead
+
+ @property
+ def max_concurrent_streams(self) -> int:
+ assert self._h2 is not None
+ advertised = self._h2.remote_settings.max_concurrent_streams
+ if advertised is None:
+ return DEFAULT_MAX_STREAMS_PER_CONN
+ return int(advertised)
+
+ @property
+ def open_stream_count(self) -> int:
+ return len(self._streams) + len(self._pending_streams)
+
+ async def close(self) -> None:
+ """Send GOAWAY, cancel recv_loop, close socket."""
+ if self._closed:
+ return
+ self._closed = True
+ self._fail_all_streams(ConnectionClosedError("Connection closed"))
+
+ if self._h2 is not None and self._writer is not None:
+ try:
+ async with self._write_lock:
+ self._h2.close_connection()
+ await self._flush_h2_data_and_drain()
+ except Exception:
+ pass
+
+ if self._recv_task is not None:
+ self._recv_task.cancel()
+ try:
+ await self._recv_task
+ except (asyncio.CancelledError, Exception):
+ pass
+
+ if self._writer is not None:
+ try:
+ self._writer.close()
+ await self._writer.wait_closed()
+ except Exception:
+ pass
+
+ def _flush_h2_data_sync(self) -> None:
+ """Write pending h2 bytes to socket. Must be called under _write_lock or during init."""
+ assert self._h2 is not None
+ assert self._writer is not None
+ data = self._h2.data_to_send()
+ if data:
+ self._writer.write(data)
+
+ async def _flush_h2_data(self) -> None:
+ """Write pending h2 bytes, draining only when the buffer is large.
+
+ This allows small writes to coalesce in the kernel buffer, reducing
+ the number of syscalls and await-points on high-rate paths.
+ """
+ assert self._writer is not None
+ self._flush_h2_data_sync()
+ transport = self._writer.transport
+ if (
+ transport is not None
+ and transport.get_write_buffer_size() >= _DRAIN_BUFFER_THRESHOLD
+ ):
+ await self._writer.drain()
+
+ async def _flush_h2_data_and_drain(self) -> None:
+ """Write pending h2 bytes and unconditionally drain the socket."""
+ assert self._writer is not None
+ self._flush_h2_data_sync()
+ await self._writer.drain()
+
+ async def _recv_loop(self) -> None:
+ """Background task: read from socket, feed to h2, dispatch events."""
+ assert self._reader is not None
+ assert self._h2 is not None
+ try:
+ while not self._closed:
+ data = await self._reader.read(65535)
+ if not data:
+ self._fail_all_streams(
+ ConnectionClosedError("Connection closed by remote")
+ )
+ return
+
+ # All h2 state mutations must be under _write_lock to avoid
+ # concurrent access with send_data/send_headers.
+ async with self._write_lock:
+ events = self._h2.receive_data(data)
+ for event in events:
+ self._handle_event(event)
+ # Flush h2 data generated by event handling (e.g. window update ACKs)
+ await self._flush_h2_data()
+ except Exception as e:
+ if not self._closed:
+ self._fail_all_streams(ConnectionClosedError(f"recv_loop error: {e}"))
+ finally:
+ self._recv_dead = True
+
+ def _handle_event(self, event: h2.events.Event) -> None:
+ if isinstance(event, h2.events.ResponseReceived):
+ state = self._streams.get(event.stream_id)
+ if state and not state.response_headers.done():
+ # h2 returns str tuples when header_encoding is set,
+ # but type stubs declare bytes.
+ headers = [(str(n), str(v)) for n, v in event.headers]
+ state.response_headers.set_result(headers)
+
+ elif isinstance(event, h2.events.DataReceived):
+ state = self._streams.get(event.stream_id)
+ if state:
+ state.data_queue.put_nowait((event.data, event.flow_controlled_length))
+ state.unacked_flow_bytes += event.flow_controlled_length
+
+ elif isinstance(event, h2.events.StreamEnded):
+ state = self._streams.get(event.stream_id)
+ if state:
+ state.end_stream_received = True
+ state.data_queue.put_nowait(None)
+ state.ended.set()
+
+ elif isinstance(event, h2.events.WindowUpdated):
+ if event.stream_id == 0:
+ # Connection-level window update — wake all streams.
+ for state in self._streams.values():
+ state.window_updated.set()
+ else:
+ state = self._streams.get(event.stream_id)
+ if state:
+ state.window_updated.set()
+
+ elif isinstance(event, h2.events.StreamReset):
+ state = self._streams.get(event.stream_id)
+ if state:
+ err = ProtocolError(
+ f"Stream reset with error code {event.error_code}",
+ error_code=event.error_code,
+ )
+ self._fail_stream(state, err)
+
+ elif isinstance(event, h2.events.ConnectionTerminated):
+ self._goaway_received = True
+ err = ProtocolError(
+ f"GOAWAY received: error_code={event.error_code}, "
+ f"last_stream_id={event.last_stream_id}",
+ error_code=event.error_code,
+ )
+ # Only fail streams the server never processed.
+ if event.last_stream_id is not None:
+ for stream_id, state in self._streams.items():
+ if stream_id > event.last_stream_id:
+ self._fail_stream(state, err)
+ else:
+ self._fail_all_streams(err)
+
+ elif isinstance(event, h2.events.RemoteSettingsChanged):
+ self._settings_received.set()
+ for state in self._streams.values():
+ state.window_updated.set()
+
+ def _fail_stream(self, state: _StreamState, error: BaseException) -> None:
+ if state.error is None:
+ state.error = error
+ if not state.response_headers.done():
+ state.response_headers.set_exception(state.error)
+ if not state.ended.is_set():
+ state.data_queue.put_nowait(None)
+ state.ended.set()
+ state.window_updated.set()
+
+ def _fail_all_streams(self, error: BaseException) -> None:
+ for state in self._streams.values():
+ self._fail_stream(state, error)
+ for state in self._pending_streams.values():
+ self._fail_stream(state, error)
+
+
+async def _drain_body(
+ conn: Connection,
+ stream_id: int,
+ content: Any,
+ on_write: Any | None,
+) -> None:
+ """Drain an async generator body into h2 send_data calls."""
+ try:
+ async for chunk in content:
+ await conn.send_data(stream_id, chunk, on_write=on_write)
+ await conn.end_stream(stream_id)
+ except (asyncio.CancelledError, Exception):
+ await conn.reset_stream(stream_id)
+ raise
+
+
+def _queue_item_parts(item: tuple[bytes, int] | bytes) -> tuple[bytes, int]:
+ if isinstance(item, tuple):
+ return item
+ return item, len(item)
+
+
+def _take_all_unacked_flow_bytes(state: _StreamState) -> int:
+ nbytes = state.unacked_flow_bytes
+ state.unacked_flow_bytes = 0
+ return nbytes
+
+
+def _parse_retry_after_ms(raw: str | None) -> float | None:
+ if raw is None:
+ return None
+ try:
+ return int(raw) / 1000.0
+ except (ValueError, TypeError):
+ return None
+
+
+def _raise_for_status(response: Response, *, retry_after_ms: str | None = None) -> None:
+ status = response.status_code
+ if 200 <= status < 300:
+ return
+
+ retry_after = _parse_retry_after_ms(retry_after_ms)
+ body: Any | None
+ try:
+ body = response.json()
+ except Exception:
+ body = None
+
+ if status == 412 and isinstance(body, dict):
+ code = body.get("code", UNKNOWN_CODE)
+ raise_for_412(body, code)
+
+ if status == 416 and isinstance(body, dict):
+ code = body.get("code", UNKNOWN_CODE)
+ raise_for_416(body, code)
+
+ if isinstance(body, dict):
+ message = body.get("message", response.text)
+ code = body.get("code", UNKNOWN_CODE)
+ else:
+ message = response.text
+ code = UNKNOWN_CODE
+
+ err = S2ServerError(code, message, status)
+ err._retry_after = retry_after
+ raise err
+
+
+class _Checkout(NamedTuple):
+ connection: _PooledConnection
+ state: _StreamState
+
+
+class _Origin(NamedTuple):
+ scheme: str
+ host: str
+ port: int
+
+
+def _origin(base_url: str) -> _Origin:
+ url = base_url if "://" in base_url else f"https://{base_url}"
+ parsed = urlsplit(url)
+ scheme = parsed.scheme.lower()
+ if scheme not in {"http", "https"}:
+ raise S2ClientError(f"Unsupported URL scheme: {parsed.scheme!r}")
+ if parsed.path not in {"", "/"} or parsed.query or parsed.fragment:
+ raise S2ClientError("Endpoint URL must be an origin without path or query")
+ if parsed.username is not None or parsed.password is not None:
+ raise S2ClientError("Endpoint URL must not include userinfo")
+ if parsed.hostname is None:
+ raise S2ClientError("Endpoint URL must include a host")
+ default_port = 443 if scheme == "https" else 80
+ try:
+ port = parsed.port if parsed.port is not None else default_port
+ except ValueError as e:
+ raise S2ClientError(f"Invalid endpoint URL: {e}") from e
+ return _Origin(scheme, parsed.hostname, port)
+
+
+def _build_path(path: str, params: dict[str, Any] | None) -> str:
+ if not params:
+ return path
+ qs = urlencode({k: v for k, v in params.items() if v is not None})
+ return f"{path}?{qs}" if qs else path
+
+
+def _status_from_headers(headers: list[tuple[str, str]]) -> int:
+ for name, value in headers:
+ if name == ":status":
+ return int(value)
+ raise TransportError("No :status header in response")
+
+
+def _header_value(headers: list[tuple[str, str]], name: str) -> str | None:
+ for k, v in headers:
+ if k == name:
+ return v
+ return None
diff --git a/src/s2_sdk/_compression.py b/src/s2_sdk/_compression.py
new file mode 100644
index 0000000..c083e2a
--- /dev/null
+++ b/src/s2_sdk/_compression.py
@@ -0,0 +1,32 @@
+"""S2S message-level compression (zstd and gzip)."""
+
+import gzip
+
+import zstandard
+
+from s2_sdk._types import Compression
+
+_zstd_compressor = zstandard.ZstdCompressor()
+_zstd_dctx = zstandard.ZstdDecompressor()
+
+
+def compress(data: bytes, compression: Compression) -> bytes:
+ match compression:
+ case Compression.ZSTD:
+ return _zstd_compressor.compress(data)
+ case Compression.GZIP:
+ return gzip.compress(data)
+ case _:
+ return data
+
+
+def decompress(data: bytes, compression: Compression) -> bytes:
+ match compression:
+ case Compression.ZSTD:
+ # stream_reader handles frames without content size in header.
+ with _zstd_dctx.stream_reader(data) as reader:
+ return reader.read()
+ case Compression.GZIP:
+ return gzip.decompress(data)
+ case _:
+ return data
diff --git a/src/s2_sdk/_exceptions.py b/src/s2_sdk/_exceptions.py
new file mode 100644
index 0000000..ca7df3a
--- /dev/null
+++ b/src/s2_sdk/_exceptions.py
@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from functools import wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from s2_sdk._types import StreamPosition
+
+
+class S2Error(Exception):
+ """Base class for all S2 related exceptions."""
+
+
+class S2ClientError(S2Error):
+ """Error originating from the client."""
+
+
+UNKNOWN_CODE = "unknown"
+
+
+class S2ServerError(S2Error):
+ """Error originating from the server.
+
+ Attributes:
+ code: Error code from the server. See `error codes 
