"""Shared protocol helpers: H2 configuration, tuning, transport base class, pump loop, and peer identity.
This is the core module. Nearly every other module in the package imports
from here.
"""
from __future__ import annotations
import asyncio
import grp
import inspect
import os
import pwd
import re
import socket
import struct
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from grpclib.config import Configuration
from grpclib.encoding.proto import ProtoCodec
from grpclib.events import _DispatchServerEvents # pyright: ignore[reportPrivateUsage] -- grpclib internal API
from grpclib.protocol import H2Protocol
from grpclib.server import Handler as ServerHandler
from h2 import __version__ as h2_version
from h2.config import H2Configuration
from h2.connection import H2Connection
from h2.errors import ErrorCodes
from h2.exceptions import StreamClosedError, StreamIDTooLowError
from h2.settings import SettingCodes
from hyperframe.frame import RstStreamFrame
from grpclib_transports.limits import limit_services_concurrency
if TYPE_CHECKING:
from collections.abc import Sequence
from grpclib._typing import IServable
from grpclib.const import Handler
from grpclib.encoding.base import StatusDetailsCodecBase
# **`SO_PEERCRED` is Linux only, and reading it through the module is what
# breaks.** Python defines the constant on Linux alone. `_unix_socket_identity`
# already asks `hasattr` before it reads the name, so the runtime is correct,
# but pyright does not narrow a module attribute through `hasattr` and reports
# three errors on macOS: an unknown member, an unknown argument, and the access
# itself.
#
# Reading it once here answers all three, and it is the shape `stdio.py` uses
# for `F_SETPIPE_SZ`.
#
# **macOS has a partial alternative, and this does not take it.** Python
# exposes `socket.LOCAL_PEERCRED` there, which `getsockopt` answers with a
# `struct xucred`: an effective uid and a group list, and no pid. The pid needs
# `LOCAL_PEERPID`, which Python does not expose on macOS at all. So a macOS
# branch would fill two of the four fields of `PeerIdentity` and leave `pid`
# empty, and it would have to unpack a second, differently shaped structure to
# do it.
#
# Nothing asks for that yet. `PeerIdentity` has no reader outside this library
# -- `pipes.py` and this module are the only two files that name it -- and
# `AGENTS.md` says to change this library when nanopynix needs a different
# behaviour, not before. **Write the macOS branch when a caller reads the
# identity of a unix peer, and record then whether a uid with no pid is
# enough.** Until then this returns None off Linux, which
# `peer_identity_from_transport` already turns into `transport="unknown"`.
_SO_PEERCRED: int | None = getattr(socket, "SO_PEERCRED", None)
_SIZE_UNITS = {
"": 1,
"b": 1,
"k": 1000,
"kb": 1000,
"m": 1000**2,
"mb": 1000**2,
"g": 1000**3,
"gb": 1000**3,
"ki": 1024,
"kib": 1024,
"mi": 1024**2,
"mib": 1024**2,
"gi": 1024**3,
"gib": 1024**3,
}
def _parse_size(value: str) -> int:
match = re.fullmatch(r"(\d+)\s*([a-zA-Z]*)", value.strip())
if match is None:
raise ValueError("expected an integer byte count or a size like 8MiB")
amount = int(match.group(1))
unit = match.group(2).lower()
multiplier = _SIZE_UNITS.get(unit)
if multiplier is None:
raise ValueError(f"unknown size unit {unit!r}")
return amount * multiplier
def _env_size(name: str, default: int) -> int:
raw = os.environ.get(name)
if raw is None:
return default
try:
value = _parse_size(raw)
except ValueError as e:
raise ValueError(f"{name} must be an integer byte count or size literal") from e
if value <= 0:
raise ValueError(f"{name} must be positive")
return value
_h2_fast_receive_patch_installed = False
#: The private methods of `H2Connection` that `_receive_frame_without_trace_repr`
#: calls. `install_h2_fast_receive_patch` refuses to patch without them.
_H2_PRIVATE_MEMBERS = (
"_prepare_for_sending",
"_stream_is_closed_by_end",
"_stream_is_closed_by_reset",
)
#: Set on a connection whose `H2Connection.__init__` returned. Issue #299 finds
#: a connection carrying all 17 attributes that `__init__` assigns before
#: `_frame_dispatch_table` and not that one, which is its last statement. Two
#: readings fit: the object left `__init__` early, or it lost that one
#: attribute afterwards. This mark tells them apart, and nothing else can:
#: `_frame_dispatch_table` is the only evidence either way and it is the thing
#: that is missing.
_H2_INIT_COMPLETE = "_grpclib_transports_init_complete"
def _mark_init_complete(original: Any) -> Any:
def marked_init(self: H2Connection, *args: Any, **kwargs: Any) -> None:
original(self, *args, **kwargs)
setattr(self, _H2_INIT_COMPLETE, True)
return marked_init
# A copy of `h2.connection.H2Connection._receive_frame` with the trace-logging
# frame reprs removed -- those reprs are formatted on every frame, whether or
# not trace logging is enabled, and they dominated the profile of a saturated
# stdio transport.
#
# It departs from h2 in one place, and the dispatch lookup below says why.
#
# Being a copy of a private method, every line that touches `self` touches an
# h2 private member, and h2 exposes no public equivalent for any of them. Each
# `pyright: ignore` below therefore says the same thing, which is why the
# reasons are short: the argument is here. `install_h2_fast_receive_patch`
# checks h2's own signature before installing this and raises if upstream has
# changed it, which is the guard that actually protects the copy.
def _receive_frame_without_trace_repr( # pyright: ignore[reportPrivateUsage] -- a copy of an h2 private method, so the name is private too
self: H2Connection,
frame: Any,
) -> list[Any]:
# **The lookup is outside the `try`, and it reports what it found.**
# `H2Connection` builds this table in `__init__`, so the class cannot
# answer for it and `install_h2_fast_receive_patch` cannot check it. One CI
# job reads it off an instance that does not have it, twice out of two
# runs, and it reproduces on no machine here -- see nanopynix issue #299.
# The bare `AttributeError` reaches the caller as `WorkerDiedError:
# Protocol error` and says nothing about the object, so the next
# occurrence says what the object really is instead.
#
# `getattr` and not `except AttributeError` around the call: a handler in
# the table raises `AttributeError` of its own, and catching here would
# rewrite that one as this failure.
dispatch: Any = getattr(self, "_frame_dispatch_table", None)
if dispatch is None:
raise RuntimeError(
f"h2 {h2_version}: {type(self).__name__} has no _frame_dispatch_table. "
f"init_complete={getattr(self, _H2_INIT_COMPLETE, False)} "
f"mro={[cls.__name__ for cls in type(self).__mro__]} attributes={sorted(vars(self))}"
)
try:
frames, events = dispatch[frame.__class__](frame) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportUnknownArgumentType] -- h2 ships no annotations for this private member
except StreamClosedError as e:
if self._stream_is_closed_by_reset(e.stream_id): # pyright: ignore[reportPrivateUsage] -- h2 private member, no public equivalent
f = RstStreamFrame(e.stream_id)
f.error_code = e.error_code
self._prepare_for_sending([f]) # pyright: ignore[reportPrivateUsage] -- h2 private member, no public equivalent
events = e._events # pyright: ignore[reportPrivateUsage,reportUnknownMemberType] -- h2 ships no annotations for this private member
else:
raise
except StreamIDTooLowError as e:
if self._stream_is_closed_by_reset(e.stream_id): # pyright: ignore[reportPrivateUsage] -- h2 private member, no public equivalent
f = RstStreamFrame(e.stream_id)
f.error_code = ErrorCodes.STREAM_CLOSED
self._prepare_for_sending([f]) # pyright: ignore[reportPrivateUsage] -- h2 private member, no public equivalent
events = []
elif self._stream_is_closed_by_end(e.stream_id): # pyright: ignore[reportPrivateUsage] -- h2 private member, no public equivalent
raise StreamClosedError(e.stream_id) from e
else:
raise
else:
self._prepare_for_sending(frames) # pyright: ignore[reportPrivateUsage,reportUnknownArgumentType] -- h2 ships no annotations for this private member
return events # pyright: ignore[reportUnknownVariableType] -- h2 ships no annotations for this private member
[docs]
def install_h2_fast_receive_patch() -> None:
"""Patch ``h2`` to skip expensive trace-logging frame reprs.
Called automatically at import time. Idempotent — subsequent calls are
no-ops. Raises :exc:`RuntimeError` if the ``h2`` internals have changed
in a way the patch cannot handle.
"""
global _h2_fast_receive_patch_installed
if _h2_fast_receive_patch_installed:
return
params = tuple(inspect.signature(H2Connection._receive_frame).parameters) # pyright: ignore[reportPrivateUsage] -- h2 private member, no public equivalent
if params != ("self", "frame"):
raise RuntimeError(f"Unsupported h2 H2Connection._receive_frame signature: {params!r}")
# The copy above reads three more private methods, and the signature says
# nothing about any of them. A release that renames one gives an
# `AttributeError` from inside the frame pump of a live connection, which
# arrives as `WorkerDiedError: Protocol error` and names neither h2 nor
# this patch. Refuse at import instead, and name the member.
#
# `_frame_dispatch_table` is absent here on purpose: `H2Connection` sets it
# in `__init__`, so the class does not carry it and only an instance can
# answer for it. Issue #299 has one CI job where an instance did not.
missing = [name for name in _H2_PRIVATE_MEMBERS if not hasattr(H2Connection, name)]
if missing:
raise RuntimeError(f"h2 {h2_version} H2Connection is missing the private members this patch reads: {missing}")
H2Connection._receive_frame = _receive_frame_without_trace_repr # pyright: ignore[reportPrivateUsage] -- h2 private member, no public equivalent
H2Connection.__init__ = _mark_init_complete(H2Connection.__init__)
_h2_fast_receive_patch_installed = True
install_h2_fast_receive_patch()
MAX_FRAME_SIZE = 2**24 - 1
DEFAULT_HTTP2_MAX_FRAME_SIZE = 1024 * 1024
[docs]
@dataclass(frozen=True)
class TransportTuning:
"""Buffer, window, and chunk size knobs for transport performance.
All values are in bytes. Create a tuned instance with :meth:`from_env`
or read the pre-computed :data:`DEFAULT_TUNING`.
Fields:
buffer_size: OS pipe buffer size used when bumping subprocess pipes.
read_chunk_size: Size of each ``read()`` call in the pump loop.
write_high_water: High-water mark for write buffer limits.
write_low_water: Low-water mark for write buffer limits.
http2_stream_window_size: Per-stream HTTP/2 flow-control window.
http2_connection_window_size: Connection-level HTTP/2 flow-control window.
http2_max_frame_size: Maximum HTTP/2 frame size sent to the peer.
transfer_chunk_size: Default chunk size for file/data chunk iterators.
"""
buffer_size: int
read_chunk_size: int
write_high_water: int
write_low_water: int
http2_stream_window_size: int
http2_connection_window_size: int
http2_max_frame_size: int
transfer_chunk_size: int
@classmethod
def from_env(cls) -> TransportTuning:
buffer_size = _env_size("GRPCLAB_BUFFER_SIZE", 8 * 1024 * 1024)
transfer_chunk_size = _env_size(
"GRPCLAB_TRANSFER_CHUNK_SIZE",
256 * 1024,
)
stream_window_size = _env_size(
"GRPCLAB_HTTP2_STREAM_WINDOW_SIZE",
max(16 * 1024 * 1024, buffer_size * 2),
)
connection_window_size = _env_size(
"GRPCLAB_HTTP2_CONNECTION_WINDOW_SIZE",
max(64 * 1024 * 1024, stream_window_size * 4),
)
max_frame_size = _env_size(
"GRPCLAB_HTTP2_MAX_FRAME_SIZE",
min(DEFAULT_HTTP2_MAX_FRAME_SIZE, buffer_size, MAX_FRAME_SIZE),
)
if max_frame_size > MAX_FRAME_SIZE:
raise ValueError("GRPCLAB_HTTP2_MAX_FRAME_SIZE exceeds HTTP/2 maximum frame size")
return cls(
buffer_size=buffer_size,
read_chunk_size=buffer_size,
write_high_water=buffer_size,
write_low_water=buffer_size // 2,
http2_stream_window_size=stream_window_size,
http2_connection_window_size=connection_window_size,
http2_max_frame_size=max_frame_size,
transfer_chunk_size=transfer_chunk_size,
)
DEFAULT_TUNING = TransportTuning.from_env()
"""Pre-computed :class:`TransportTuning` from environment variables.
Overridable via ``GRPCLAB_BUFFER_SIZE``, ``GRPCLAB_TRANSFER_CHUNK_SIZE``,
``GRPCLAB_HTTP2_STREAM_WINDOW_SIZE``, ``GRPCLAB_HTTP2_CONNECTION_WINDOW_SIZE``,
and ``GRPCLAB_HTTP2_MAX_FRAME_SIZE`` (all support suffixes like ``8MiB``).
"""
[docs]
@dataclass(frozen=True)
class PeerIdentity:
"""Identity information about the remote peer of a transport.
Fields:
transport: Transport type (``"stdio"``, ``"ssh"``, ``"unix"``, ``"unknown"``).
username: OS username of the peer, if available.
uid: Unix user ID of the peer.
gid: Unix group ID of the peer.
pid: Process ID of the peer.
group_ids: Supplementary group IDs.
group_names: Supplementary group names (resolved from *group_ids*).
"""
transport: str
username: str | None
uid: int | None = None
gid: int | None = None
pid: int | None = None
group_ids: tuple[int, ...] | None = None
group_names: tuple[str, ...] | None = None
def _username_for_uid(uid: int) -> str | None:
try:
return pwd.getpwuid(uid).pw_name
except KeyError:
return None
def _group_names(group_ids: Sequence[int]) -> tuple[str, ...]:
names: list[str] = []
for group_id in group_ids:
try:
names.append(grp.getgrgid(group_id).gr_name)
except KeyError:
names.append(str(group_id))
return tuple(names)
def _groups_for_user(username: str, gid: int) -> tuple[int, ...] | None:
try:
return tuple(os.getgrouplist(username, gid))
except OSError:
return None
[docs]
def local_process_identity(*, transport: str) -> PeerIdentity:
"""Build a :class:`PeerIdentity` for the current process."""
uid = os.geteuid()
gid = os.getegid()
username = _username_for_uid(uid)
group_ids = tuple(os.getgroups())
return PeerIdentity(
transport=transport,
username=username,
uid=uid,
gid=gid,
pid=os.getpid(),
group_ids=group_ids,
group_names=_group_names(group_ids),
)
def _unix_socket_identity(sock: socket.socket) -> PeerIdentity | None:
if _SO_PEERCRED is None:
return None
try:
raw = sock.getsockopt(
socket.SOL_SOCKET,
_SO_PEERCRED,
struct.calcsize("3i"),
)
except OSError:
return None
pid, uid, gid = struct.unpack("3i", raw)
username = _username_for_uid(uid)
group_ids = _groups_for_user(username, gid) if username is not None else None
return PeerIdentity(
transport="unix",
username=username,
uid=uid,
gid=gid,
pid=pid,
group_ids=group_ids,
group_names=_group_names(group_ids) if group_ids is not None else None,
)
[docs]
def peer_identity_from_transport(transport: Any) -> PeerIdentity:
"""Extract :class:`PeerIdentity` from an asyncio transport.
Checks for a stored ``peer_identity`` extra, an SSH username, or a
Unix-domain socket with ``SO_PEERCRED``.
"""
identity = transport.get_extra_info("peer_identity")
if isinstance(identity, PeerIdentity):
return identity
ssh_username = transport.get_extra_info("username")
if isinstance(ssh_username, str):
return PeerIdentity(transport="ssh", username=ssh_username)
sock = transport.get_extra_info("socket")
if isinstance(sock, socket.socket) and sock.family == socket.AF_UNIX:
identity = _unix_socket_identity(sock)
if identity is not None:
return identity
return PeerIdentity(transport="unknown", username=None)
[docs]
def peer_identity_from_stream(stream: Any) -> PeerIdentity:
"""Extract :class:`PeerIdentity` from a gRPC stream's underlying transport."""
transport = stream.peer._transport
return peer_identity_from_transport(transport)
[docs]
class BaseCustomTransport(asyncio.Transport):
"""Base class for custom asyncio transports that wrap a reader/writer pair.
Concrete transports (StdioTransport, SshTransport) override:
- ``write()`` — forward data to the underlying writer
- ``get_extra_info()`` — delegate to the appropriate underlying object
- ``get_write_buffer_size()`` — introspect the write buffer (for flow control)
- ``abort()`` — hard reset (transport-specific)
- ``can_write_eof()`` / ``write_eof()`` — EOF support (transport-specific)
Flow-control bridging (pause_writing/resume_writing forwarding to the
H2Protocol) is handled by each concrete transport's own mechanism.
"""
def __init__(self) -> None:
super().__init__()
self._protocol: asyncio.BaseProtocol | None = None
self._closing = False
[docs]
def is_closing(self) -> bool:
return self._closing
[docs]
def get_protocol(self) -> asyncio.BaseProtocol:
protocol = self._protocol
if protocol is None:
raise RuntimeError("protocol has not been set yet")
return protocol
[docs]
def set_protocol(self, protocol: asyncio.BaseProtocol) -> None:
self._protocol = protocol
[docs]
def pause_reading(self) -> None:
pass
[docs]
def resume_reading(self) -> None:
pass
[docs]
def pause_h2_protocol(protocol: asyncio.BaseProtocol | None) -> None:
"""Pause writing on an H2 protocol, if set."""
if protocol is not None:
protocol.pause_writing()
[docs]
def resume_h2_protocol(protocol: asyncio.BaseProtocol | None) -> None:
"""Resume writing on an H2 protocol, if set."""
if protocol is not None:
protocol.resume_writing()
[docs]
async def pump(
protocol: H2Protocol,
reader: Any,
*,
tuning: TransportTuning = DEFAULT_TUNING,
) -> None:
"""Read from a byte stream and feed data into an H2 protocol.
Blocks in a loop calling ``reader.read()``. Stops on EOF, on error, or
once the connection closes. Calls ``protocol.connection_lost()`` on the
way out.
**It stops at a closed connection, and that is not an optimisation.**
`grpclib.protocol.Connection.close` deletes
`self._connection._frame_dispatch_table` to break the one reference cycle
an `H2Connection` holds. Under asyncio a transport delivers nothing after
close, so upstream may do that safely. This loop reads a pipe and calls
`data_received` by hand, so bytes already in flight arrive after the close
and reach `H2Connection.receive_data` on a connection that has no dispatch
table left. nanopynix issue #299: deterministic on a CI runner, where the
ordering always puts those bytes after the close, and absent on a fast
workstation over 24 runs.
"""
exc: BaseException | None = None
try:
while True:
data = await reader.read(tuning.read_chunk_size)
if not data:
break
connection = getattr(protocol, "connection", None)
if connection is not None and connection.is_closing():
break
protocol.data_received(data)
except (ConnectionError, OSError) as e:
exc = e
finally:
protocol.connection_lost(exc)
[docs]
async def serve_h2(
handlers: Sequence[IServable],
reader: Any,
transport: Any,
*,
tuning: TransportTuning = DEFAULT_TUNING,
max_concurrency: int | None = None,
status_details_codec: StatusDetailsCodecBase | None = None,
) -> None:
"""Build a server protocol, wire it to a transport, and pump frames."""
mapping = build_mapping(
limit_services_concurrency(
handlers,
max_concurrency=max_concurrency,
)
)
protocol = make_server_protocol(
mapping,
tuning=tuning,
status_details_codec=status_details_codec,
)
init_h2_transport(protocol, transport, tuning=tuning)
await pump(protocol, reader, tuning=tuning)
[docs]
def make_h2_config(*, client_side: bool) -> H2Configuration:
"""Build an ``h2`` configuration with strict validation off.
Disables inbound/outbound header validation and normalization so
protobuf-based headers pass through unchanged.
"""
return H2Configuration(
client_side=client_side,
header_encoding="ascii",
validate_inbound_headers=False,
validate_outbound_headers=False,
normalize_inbound_headers=False,
normalize_outbound_headers=False,
)
[docs]
def make_config(tuning: TransportTuning = DEFAULT_TUNING) -> Configuration:
"""Build a grpclib :class:`Configuration` with tuned window sizes."""
return Configuration(
http2_connection_window_size=tuning.http2_connection_window_size,
http2_stream_window_size=tuning.http2_stream_window_size,
)
[docs]
def make_server_protocol(
mapping: dict[str, Handler],
*,
tuning: TransportTuning = DEFAULT_TUNING,
status_details_codec: StatusDetailsCodecBase | None = None,
) -> H2Protocol:
"""Build a server-side H2 protocol with the given handler mapping.
``status_details_codec`` is what lets a handler's ``GRPCError.details``
reach the client: grpclib encodes it into the ``grpc-status-details-bin``
trailer only when a codec is configured, and drops it silently otherwise.
The client's channel must be given the same codec to decode it.
"""
config = make_config(tuning).__for_server__()
h2_config = make_h2_config(client_side=False)
handler = ServerHandler(mapping, ProtoCodec(), status_details_codec, _DispatchServerEvents()) # pyright: ignore[reportPrivateUsage] -- grpclib internal API
return H2Protocol(handler, config, h2_config)
[docs]
def init_h2_transport(
protocol: H2Protocol,
transport: Any,
*,
tuning: TransportTuning = DEFAULT_TUNING,
) -> None:
"""Wire an H2 protocol to a transport and advertise a tuned max frame size."""
transport.set_protocol(protocol)
protocol.connection_made(transport)
protocol.connection._connection.update_settings( # pyright: ignore[reportPrivateUsage] -- h2 internal attribute
{
SettingCodes.MAX_FRAME_SIZE: tuning.http2_max_frame_size,
}
)
protocol.connection.flush()
[docs]
def build_mapping(handlers: Sequence[IServable]) -> dict[str, Handler]:
"""Merge ``__mapping__()`` from a sequence of servable handlers into one dict."""
mapping: dict[str, Handler] = {}
for h in handlers:
mapping.update(h.__mapping__())
return mapping
[docs]
def signal_stop(stop: asyncio.Future[None]) -> None:
"""Signal a stop future, guarding against duplicate completion."""
if not stop.done():
stop.set_result(None)