299 lines
12 KiB
Python
299 lines
12 KiB
Python
"""Native activation transport over direct gRPC or the existing relay RPC.
|
|
|
|
This is deliberately a *seam adapter*, not a new relay protocol. Direct
|
|
peers use one generated ``ShardRuntime.Session`` bidi stream for the lifetime
|
|
of a Route Session. A relayed peer uses the relay's existing HTTP-shaped
|
|
binary-body contract: each body is exactly a serialized ``SessionRequest`` or
|
|
``SessionResponse``. The relay only routes those bytes and restores its own
|
|
request id; it does not deserialize a native frame.
|
|
|
|
The correlation headers are duplicated outside the opaque frame solely for
|
|
the existing tracker/relay observability and billing path. The authoritative
|
|
work, route, epoch, deadline, and cancellation information remains in the
|
|
versioned protobuf frame and is validated before it is sent.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Iterator
|
|
from dataclasses import dataclass
|
|
from queue import Empty, Full, Queue
|
|
import threading
|
|
import time
|
|
from typing import Protocol
|
|
|
|
from .native_protocol import pb
|
|
|
|
NATIVE_RELAY_PATH = "/native/session"
|
|
NATIVE_FRAME_CONTENT_TYPE = "application/x-protobuf"
|
|
|
|
|
|
class NativeActivationSeamError(RuntimeError):
|
|
"""The activation seam cannot safely continue this Route Session."""
|
|
|
|
|
|
class NativeActivationBufferFull(NativeActivationSeamError):
|
|
"""The caller exceeded the negotiated local hand-off buffer."""
|
|
|
|
|
|
class NativeActivationDisconnected(NativeActivationSeamError):
|
|
"""A direct or relay transport disconnected with an uncertain outcome."""
|
|
|
|
|
|
class RelayRequest(Protocol):
|
|
"""The existing ``_RelayHopClient.request`` shape, kept dependency-free."""
|
|
|
|
def __call__(
|
|
self, path: str, body: bytes, headers: dict[str, str]
|
|
) -> tuple[int, dict[str, str], bytes]: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class NativeFrameContext:
|
|
"""Correlation owned by Meshnet around one opaque native frame."""
|
|
|
|
request_id: str
|
|
node_id: str
|
|
route_session_id: str
|
|
route_epoch: int
|
|
work_id: str = ""
|
|
deadline_unix_nanos: int = 0
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.request_id or not self.node_id or not self.route_session_id:
|
|
raise ValueError("request, node, and Route Session identities are required")
|
|
if self.route_epoch < 0 or self.deadline_unix_nanos < 0:
|
|
raise ValueError("route epoch and deadline must be non-negative")
|
|
|
|
def headers(self) -> dict[str, str]:
|
|
"""Headers retained by the existing relay/Tracker accounting path."""
|
|
return {
|
|
"Content-Type": NATIVE_FRAME_CONTENT_TYPE,
|
|
"X-Meshnet-Native-Frame": "shard-runtime/v1",
|
|
"X-Meshnet-Request-Id": self.request_id,
|
|
"X-Meshnet-Node-Id": self.node_id,
|
|
"X-Meshnet-Session": self.route_session_id,
|
|
"X-Meshnet-Route-Epoch": str(self.route_epoch),
|
|
"X-Meshnet-Work-Id": self.work_id,
|
|
"X-Meshnet-Deadline-Unix-Nanos": str(self.deadline_unix_nanos),
|
|
# The relay request id is restored on reply and is intentionally
|
|
# distinct from the caller/billing request id above.
|
|
"X-Meshnet-Activation-Id": self.request_id,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class NativeSeamTelemetry:
|
|
transport: str
|
|
request_id: str
|
|
node_id: str
|
|
work_id: str
|
|
request_bytes: int
|
|
response_bytes: int
|
|
elapsed_seconds: float
|
|
|
|
|
|
TelemetrySink = Callable[[NativeSeamTelemetry], None]
|
|
|
|
|
|
def _request_identity(request: pb.SessionRequest) -> tuple[str, int, str, int]:
|
|
kind = request.WhichOneof("kind")
|
|
if kind == "open":
|
|
return request.open.route_session_id, request.open.route_epoch, "", 0
|
|
if kind == "chunk":
|
|
item = request.chunk.envelope
|
|
return item.route_session_id, item.route_epoch, item.work_id, item.deadline_unix_nanos
|
|
if kind == "decode":
|
|
# DecodeStep relies on the already opened Route Session, while work
|
|
# identity/deadline are carried on every decode frame.
|
|
return "", 0, request.decode.work_id, request.decode.deadline_unix_nanos
|
|
if kind in {"cancel", "release"}:
|
|
item = getattr(request, kind)
|
|
return item.route_session_id, item.route_epoch, item.work_id, 0
|
|
if kind == "flow_control":
|
|
return "", 0, "", 0
|
|
raise NativeActivationSeamError("native SessionRequest has no frame kind")
|
|
|
|
|
|
def _validate_request(request: pb.SessionRequest, context: NativeFrameContext) -> None:
|
|
if request.ByteSize() == 0:
|
|
raise NativeActivationSeamError("empty native SessionRequest is not a versioned frame")
|
|
route_session, epoch, work_id, deadline = _request_identity(request)
|
|
if route_session and route_session != context.route_session_id:
|
|
raise NativeActivationSeamError("native frame Route Session differs from seam context")
|
|
if route_session and epoch != context.route_epoch:
|
|
raise NativeActivationSeamError("native frame route epoch differs from seam context")
|
|
if context.work_id and work_id and work_id != context.work_id:
|
|
raise NativeActivationSeamError("native frame work identity differs from seam context")
|
|
if context.deadline_unix_nanos and deadline and deadline != context.deadline_unix_nanos:
|
|
raise NativeActivationSeamError("native frame deadline differs from seam context")
|
|
|
|
|
|
def _response_work_id(response: pb.SessionResponse) -> str:
|
|
kind = response.WhichOneof("kind")
|
|
if kind == "chunk":
|
|
return response.chunk.envelope.work_id
|
|
if kind == "ack":
|
|
return response.ack.work_id
|
|
if kind == "status":
|
|
return response.status.work_id
|
|
return ""
|
|
|
|
|
|
class NativeActivationSeam:
|
|
"""One Route-Session-to-worker seam with bounded direct buffering.
|
|
|
|
``direct_stub`` is the generated ``ShardRuntimeStub`` and is selected when
|
|
it is available. ``relay_request`` has the exact signature of the
|
|
existing persistent relay client; no relay server or bridge API changes
|
|
are needed. Relay calls are intentionally not retried: a failed send may
|
|
already have mutated downstream Hot KV state.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
context: NativeFrameContext,
|
|
*,
|
|
direct_stub=None,
|
|
relay_request: RelayRequest | None = None,
|
|
max_buffered_frames: int = 8,
|
|
telemetry: TelemetrySink | None = None,
|
|
) -> None:
|
|
if (direct_stub is None) == (relay_request is None):
|
|
raise ValueError("provide exactly one of direct_stub or relay_request")
|
|
if max_buffered_frames < 1:
|
|
raise ValueError("max_buffered_frames must be positive")
|
|
self.context = context
|
|
self._direct_stub = direct_stub
|
|
self._relay_request = relay_request
|
|
self._telemetry = telemetry
|
|
self._closed = False
|
|
self._failure: BaseException | None = None
|
|
self._responses: Queue[pb.SessionResponse | BaseException] = Queue(maxsize=max_buffered_frames)
|
|
self._requests: Queue[pb.SessionRequest | object] | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self._stop = object()
|
|
if direct_stub is not None:
|
|
self._requests = Queue(maxsize=max_buffered_frames)
|
|
self._thread = threading.Thread(target=self._run_direct, daemon=True, name="native-activation-grpc")
|
|
self._thread.start()
|
|
|
|
@property
|
|
def transport(self) -> str:
|
|
return "direct-grpc" if self._direct_stub is not None else "relay"
|
|
|
|
def _direct_requests(self) -> Iterator[pb.SessionRequest]:
|
|
assert self._requests is not None
|
|
while True:
|
|
item = self._requests.get()
|
|
if item is self._stop:
|
|
return
|
|
assert isinstance(item, pb.SessionRequest)
|
|
yield item
|
|
|
|
def _run_direct(self) -> None:
|
|
try:
|
|
assert self._direct_stub is not None
|
|
for response in self._direct_stub.Session(self._direct_requests()):
|
|
self._put_response(response)
|
|
except BaseException as exc:
|
|
self._failure = exc
|
|
self._put_response(exc)
|
|
|
|
def _put_response(self, value: pb.SessionResponse | BaseException) -> None:
|
|
# A worker may finish while a caller is abandoning the session. Do not
|
|
# let an unconsumed response turn into an unbounded producer queue.
|
|
try:
|
|
self._responses.put(value, timeout=0.1)
|
|
except Full:
|
|
self._failure = NativeActivationBufferFull("native response buffer is full")
|
|
|
|
def send(self, request: pb.SessionRequest) -> pb.SessionResponse | None:
|
|
"""Send one already-versioned protobuf frame without rewriting it."""
|
|
if self._closed:
|
|
raise NativeActivationDisconnected("native activation seam is closed")
|
|
if self._failure is not None:
|
|
raise NativeActivationDisconnected("native activation stream failed") from self._failure
|
|
_validate_request(request, self.context)
|
|
frame = request.SerializeToString()
|
|
if self._direct_stub is not None:
|
|
assert self._requests is not None
|
|
try:
|
|
self._requests.put_nowait(request)
|
|
except Full as exc:
|
|
raise NativeActivationBufferFull("native direct request buffer is full") from exc
|
|
return None
|
|
|
|
assert self._relay_request is not None
|
|
started = time.monotonic()
|
|
try:
|
|
status, _, response_frame = self._relay_request(NATIVE_RELAY_PATH, frame, self.context.headers())
|
|
except Exception as exc:
|
|
self._closed = True
|
|
raise NativeActivationDisconnected("relay outcome is uncertain; refusing replay") from exc
|
|
if status != 200:
|
|
self._closed = True
|
|
raise NativeActivationDisconnected(f"relay native frame returned HTTP {status}")
|
|
response = pb.SessionResponse()
|
|
try:
|
|
response.ParseFromString(response_frame)
|
|
except Exception as exc:
|
|
self._closed = True
|
|
raise NativeActivationSeamError("relay returned a malformed native response frame") from exc
|
|
self._validate_response(response)
|
|
self._record(len(frame), len(response_frame), started)
|
|
return response
|
|
|
|
def receive(self, timeout: float | None = None) -> pb.SessionResponse:
|
|
"""Receive the next response from the one long-lived direct stream."""
|
|
if self._direct_stub is None:
|
|
raise NativeActivationSeamError("relay sends return their response synchronously")
|
|
try:
|
|
value = self._responses.get(timeout=timeout)
|
|
except Empty as exc:
|
|
raise TimeoutError("timed out waiting for native direct response") from exc
|
|
if isinstance(value, BaseException):
|
|
raise NativeActivationDisconnected("native direct stream disconnected") from value
|
|
self._validate_response(value)
|
|
# gRPC owns its framing, but this records the actual protobuf payload
|
|
# size at the seam for the same telemetry shape as relay.
|
|
self._record(0, len(value.SerializeToString()), time.monotonic())
|
|
return value
|
|
|
|
def cancel(self, reason: str = "cancelled") -> pb.SessionResponse | None:
|
|
"""Propagate cancellation through the same path and correlation fields."""
|
|
return self.send(pb.SessionRequest(cancel=pb.CancelSignal(
|
|
route_session_id=self.context.route_session_id,
|
|
route_epoch=self.context.route_epoch,
|
|
work_id=self.context.work_id,
|
|
reason=reason,
|
|
)))
|
|
|
|
def _validate_response(self, response: pb.SessionResponse) -> None:
|
|
work_id = _response_work_id(response)
|
|
if self.context.work_id and work_id and work_id != self.context.work_id:
|
|
raise NativeActivationSeamError("native response work identity differs from seam context")
|
|
|
|
def _record(self, request_bytes: int, response_bytes: int, started: float) -> None:
|
|
if self._telemetry is not None:
|
|
self._telemetry(NativeSeamTelemetry(
|
|
transport=self.transport, request_id=self.context.request_id,
|
|
node_id=self.context.node_id, work_id=self.context.work_id,
|
|
request_bytes=request_bytes, response_bytes=response_bytes,
|
|
elapsed_seconds=max(0.0, time.monotonic() - started),
|
|
))
|
|
|
|
def close(self) -> None:
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
if self._requests is not None:
|
|
try:
|
|
self._requests.put_nowait(self._stop)
|
|
except Full:
|
|
# The bounded queue is intentionally never expanded during
|
|
# shutdown; the worker will observe process/session teardown.
|
|
pass
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=1.0)
|