422 lines
17 KiB
Python
422 lines
17 KiB
Python
"""TLS gRPC candidate carrying the existing raw-first binary wire unchanged.
|
|
|
|
One stream per resident profile. The trusted local controller issues one-shot,
|
|
short-lived capabilities AFTER choosing the input epoch. Network code cannot
|
|
activate models, renew leases, reset temporal stores or authorize a reconnect.
|
|
No insecure listener/channel, route upload, background source queue or retry.
|
|
|
|
The gRPC method uses identity byte serialization (not generated protobufs):
|
|
each request is <=64 KiB of the existing MCI2 byte stream, each response one
|
|
bounded MCR3 envelope. Source/result semantics remain in the existing adapters.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hmac
|
|
import secrets
|
|
import socket
|
|
import threading
|
|
import time
|
|
from collections.abc import Callable, Iterable
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass, field
|
|
from hashlib import sha256
|
|
from typing import Any
|
|
|
|
import grpc # type: ignore[import-untyped] # Upstream wheel does not ship type stubs.
|
|
|
|
from k1link.compute.live_perception import LiveIngressEvent
|
|
|
|
from . import streaming_wire as wire
|
|
from .realtime_contract import StreamStart
|
|
from .streaming_ingress import StreamingIngress
|
|
from .streaming_lifecycle import StreamingLifecycle
|
|
from .streaming_network import (
|
|
BRIDGE_RESERVATION,
|
|
INPUT_CHUNK,
|
|
MAX_REPLY,
|
|
LatestReplies,
|
|
decode_reply,
|
|
encode_reply,
|
|
)
|
|
|
|
SERVICE = "missioncore.perception.v1.BinaryStream"
|
|
METHOD = f"/{SERVICE}/Exchange"
|
|
OPTIONS = (
|
|
("grpc.enable_retries", 0),
|
|
("grpc.http2.bdp_probe", 0),
|
|
# A fixed 1 MiB HTTP/2 receive window, not the default 64 KiB which
|
|
# serializes large observations into many WAN/SSH round trips. Keep BDP
|
|
# auto-growth disabled; native buffering remains part of the RSS gate.
|
|
("grpc.http2.lookahead_bytes", wire.MAX_FRAGMENT),
|
|
("grpc.http2.write_buffer_size", INPUT_CHUNK),
|
|
("grpc.max_metadata_size", 4096),
|
|
("grpc.max_concurrent_streams", 2),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StreamAccess:
|
|
epoch: StreamStart
|
|
session_id: str
|
|
session_generation: int
|
|
token: str = field(repr=False)
|
|
|
|
def metadata(self) -> tuple[tuple[str, str], ...]:
|
|
return (
|
|
("authorization", "Bearer " + self.token),
|
|
("x-stream-binding", wire.binding(self.epoch)),
|
|
)
|
|
|
|
|
|
class GrpcStreamEndpoint:
|
|
def __init__(
|
|
self,
|
|
runtime: StreamingLifecycle,
|
|
consume: Callable[[LiveIngressEvent], None],
|
|
notice: Callable[[dict[str, Any]], None],
|
|
*,
|
|
idle_timeout: float = 2.0,
|
|
io_timeout: float = 0.25,
|
|
) -> None:
|
|
if runtime.continuity is None:
|
|
raise ValueError("external stream requires a resumable resident runtime")
|
|
if not 0 < idle_timeout <= 30 or not 0 < io_timeout <= 2:
|
|
raise ValueError("invalid bounded network timeouts")
|
|
self.runtime, self.consume, self.notice = runtime, consume, notice
|
|
self.idle_timeout, self.io_timeout = idle_timeout, io_timeout
|
|
self.lock = threading.Lock()
|
|
self.grant: tuple[StreamAccess, bytes, float] | None = None
|
|
self.active: tuple[StreamStart, LatestReplies] | None = None
|
|
self.last: dict[str, Any] | None = None
|
|
self._quarantine_thread: threading.Thread | None = None
|
|
|
|
def issue(self, epoch: StreamStart, session_id: str, session_generation: int) -> StreamAccess:
|
|
"""Local control plane only. Delivery of this secret is outside data-plane."""
|
|
wire.identifier(session_id)
|
|
wire.decimal(session_generation)
|
|
if session_generation < 1:
|
|
raise ValueError("invalid acquisition generation")
|
|
self.runtime.check_input(epoch, synchronizing=True)
|
|
with self.lock:
|
|
if self.active is not None:
|
|
raise wire.StreamWireError("previous network stream has not drained")
|
|
access = StreamAccess(epoch, session_id, session_generation, secrets.token_hex(32))
|
|
# Do not retain the plaintext secret on the endpoint.
|
|
self.grant = (
|
|
StreamAccess(epoch, session_id, session_generation, ""),
|
|
sha256(access.token.encode()).digest(),
|
|
time.monotonic() + 30,
|
|
)
|
|
return access
|
|
|
|
def _claim(self, metadata: Iterable[tuple[str, str | bytes]]) -> StreamAccess:
|
|
values: dict[str, str | bytes] = {}
|
|
for key, value in metadata:
|
|
if key in ("authorization", "x-stream-binding"):
|
|
if key in values:
|
|
raise wire.StreamWireError("stream authorization rejected")
|
|
values[key] = value
|
|
token = values.get("authorization")
|
|
if not isinstance(token, str) or len(token) != 71 or not token.startswith("Bearer "):
|
|
raise wire.StreamWireError("stream authorization rejected")
|
|
with self.lock:
|
|
grant = self.grant
|
|
if (
|
|
grant is None
|
|
or self.active is not None
|
|
or time.monotonic() >= grant[2]
|
|
or values.get("x-stream-binding") != wire.binding(grant[0].epoch)
|
|
or not hmac.compare_digest(sha256(token[7:].encode()).digest(), grant[1])
|
|
):
|
|
raise wire.StreamWireError("stream authorization rejected")
|
|
self.runtime.check_input(grant[0].epoch, synchronizing=True)
|
|
self.grant = None
|
|
self.active = (grant[0].epoch, LatestReplies())
|
|
return grant[0]
|
|
|
|
def publish(self, epoch: StreamStart, sequence: int, payload: bytes) -> bool:
|
|
self.runtime.check_input(epoch)
|
|
with self.lock:
|
|
if self.active is None or self.active[0] != epoch:
|
|
return False
|
|
return self.active[1].put(sequence, payload)
|
|
|
|
def finish_results(self, epoch: StreamStart) -> None:
|
|
"""Graph controller signals output EOF AFTER all admitted work is accounted."""
|
|
with self.lock:
|
|
if self.active is not None and self.active[0] == epoch:
|
|
self.active[1].finish()
|
|
|
|
async def serve(
|
|
self,
|
|
address: str,
|
|
*,
|
|
certificate: bytes,
|
|
private_key: bytes,
|
|
control_handlers: tuple[Any, ...] = (),
|
|
) -> Any:
|
|
if not certificate or not private_key:
|
|
raise ValueError("TLS identity is required")
|
|
server = grpc.aio.server(
|
|
options=OPTIONS
|
|
+ (
|
|
("grpc.max_receive_message_length", INPUT_CHUNK),
|
|
("grpc.max_send_message_length", MAX_REPLY),
|
|
),
|
|
# aio reserves its next accept before awaiting a request. Leave
|
|
# one bounded accept slot beside Exchange + a completing Poll;
|
|
# otherwise the next Poll can inherit a stale "limit exceeded".
|
|
# Data ownership still admits exactly one Exchange, never 2 GPUs.
|
|
maximum_concurrent_rpcs=3 if control_handlers else 2,
|
|
compression=grpc.Compression.NoCompression,
|
|
)
|
|
server.add_generic_rpc_handlers(
|
|
(
|
|
grpc.method_handlers_generic_handler(
|
|
SERVICE, {"Exchange": grpc.stream_stream_rpc_method_handler(self.exchange)}
|
|
),
|
|
*control_handlers,
|
|
)
|
|
)
|
|
port = server.add_secure_port(
|
|
address, grpc.ssl_server_credentials(((private_key, certificate),))
|
|
)
|
|
if port == 0:
|
|
raise RuntimeError("TLS stream listener failed to bind")
|
|
await server.start()
|
|
return server, port
|
|
|
|
async def exchange(self, unused_iterator: Any, context: Any) -> None:
|
|
if b"ssl" not in context.auth_context().get("transport_security_type", ()):
|
|
await context.abort(grpc.StatusCode.UNAUTHENTICATED, "TLS required")
|
|
try:
|
|
access = self._claim(context.invocation_metadata())
|
|
except (ValueError, RuntimeError):
|
|
await context.abort(grpc.StatusCode.UNAUTHENTICATED, "stream authorization rejected")
|
|
return
|
|
assert self.active is not None
|
|
outbox = self.active[1]
|
|
left = right = receiver = reservation = None
|
|
tasks: list[asyncio.Task[None]] = []
|
|
try:
|
|
reservation = self.runtime.mailbox.reserve_ingress(BRIDGE_RESERVATION)
|
|
left, right = socket.socketpair()
|
|
for sock in (left, right):
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 32 * 1024)
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 32 * 1024)
|
|
sock.setblocking(False)
|
|
receiver = StreamingIngress(
|
|
right,
|
|
self.runtime,
|
|
access.session_id,
|
|
access.session_generation,
|
|
self.consume,
|
|
self.notice,
|
|
input_epoch=access.epoch,
|
|
fragment_timeout=self.io_timeout,
|
|
idle_timeout=self.idle_timeout,
|
|
)
|
|
receiver.start()
|
|
await context.send_initial_metadata((("x-stream-binding", wire.binding(access.epoch)),))
|
|
|
|
async def read_input() -> None:
|
|
while True:
|
|
async with asyncio.timeout(self.idle_timeout):
|
|
raw = await context.read()
|
|
if raw is grpc.aio.EOF:
|
|
# Explicit wire End can already have closed the peer.
|
|
with suppress(OSError):
|
|
left.shutdown(socket.SHUT_WR)
|
|
return
|
|
if not isinstance(raw, bytes) or not 0 < len(raw) <= INPUT_CHUNK:
|
|
raise wire.StreamWireError("invalid network chunk")
|
|
async with asyncio.timeout(self.io_timeout):
|
|
await asyncio.get_running_loop().sock_sendall(left, raw)
|
|
|
|
async def write_output() -> None:
|
|
while True:
|
|
if not receiver.thread.is_alive() and receiver.terminal != "end":
|
|
raise ConnectionError("input interrupted or rejected")
|
|
self.runtime.check_input(access.epoch, synchronizing=True)
|
|
item = outbox.take()
|
|
if item is not None:
|
|
self.runtime.check_input(access.epoch)
|
|
raw = encode_reply(access.epoch, *item)
|
|
async with asyncio.timeout(self.io_timeout):
|
|
await context.write(raw)
|
|
del raw
|
|
item = None
|
|
elif receiver.terminal == "end" and outbox.drained():
|
|
return
|
|
else:
|
|
await asyncio.sleep(0.01)
|
|
|
|
tasks = [asyncio.create_task(read_input()), asyncio.create_task(write_output())]
|
|
await asyncio.gather(*tasks)
|
|
except (TimeoutError, OSError, ValueError, RuntimeError):
|
|
try:
|
|
# Even error-status delivery can block behind an unread result.
|
|
# It must never delay closing ingress/releasing local buffers.
|
|
async with asyncio.timeout(self.io_timeout):
|
|
await context.abort(
|
|
grpc.StatusCode.UNAVAILABLE,
|
|
"stream interrupted; fresh controller grant required",
|
|
)
|
|
except TimeoutError:
|
|
context.set_code(grpc.StatusCode.UNAVAILABLE)
|
|
except Exception as exc:
|
|
# A blocked peer can cancel the HTTP/2 call while the server is
|
|
# sending its final status. Cleanup still owns the real outcome.
|
|
if not isinstance(exc, grpc.RpcError):
|
|
raise
|
|
finally:
|
|
for task in tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
if left is not None:
|
|
left.close()
|
|
if receiver is not None and receiver.thread.ident is not None:
|
|
deadline = time.monotonic() + 1
|
|
while receiver.thread.is_alive() and time.monotonic() < deadline:
|
|
await asyncio.sleep(0.01)
|
|
elif right is not None:
|
|
right.close()
|
|
outbox.close()
|
|
drained = receiver is None or not receiver.thread.is_alive()
|
|
if drained:
|
|
self._complete(receiver, reservation, outbox)
|
|
else:
|
|
self._quarantine_thread = threading.Thread(
|
|
target=self._finish_quarantine,
|
|
args=(receiver, reservation, outbox),
|
|
daemon=True,
|
|
name="perception-network-quarantine",
|
|
)
|
|
self._quarantine_thread.start()
|
|
# A hung trusted consumer is quarantined: no replacement grant or
|
|
# premature reservation release while it still owns raw bytes.
|
|
|
|
def _finish_quarantine(
|
|
self, receiver: StreamingIngress, reservation: Any, outbox: LatestReplies
|
|
) -> None:
|
|
receiver.thread.join() # Bounded to this one already-quarantined stream.
|
|
self._complete(receiver, reservation, outbox)
|
|
|
|
def _complete(
|
|
self, receiver: StreamingIngress | None, reservation: Any, outbox: LatestReplies
|
|
) -> None:
|
|
if reservation is not None:
|
|
reservation.release()
|
|
with self.lock:
|
|
self.last = {
|
|
"ingress": receiver.snapshot() if receiver else None,
|
|
"reply_drops": outbox.dropped,
|
|
"network_threads_drained": True,
|
|
"transport": "grpc-tls-binary/v1",
|
|
"actuation_allowed": False,
|
|
}
|
|
self.active = None
|
|
|
|
|
|
class GrpcStreamClient:
|
|
"""Explicit async source/result API. Exactly one writer and one reader."""
|
|
|
|
def __init__(
|
|
self, target: str, roots: bytes, access: StreamAccess, *, timeout: float = 0.25
|
|
) -> None:
|
|
if not roots or not 0 < timeout <= 2:
|
|
raise ValueError("trusted TLS roots and bounded timeout required")
|
|
self.access, self.timeout = access, timeout
|
|
self.channel = grpc.aio.secure_channel(
|
|
target,
|
|
grpc.ssl_channel_credentials(root_certificates=roots),
|
|
options=OPTIONS
|
|
+ (
|
|
("grpc.max_send_message_length", INPUT_CHUNK),
|
|
("grpc.max_receive_message_length", MAX_REPLY),
|
|
),
|
|
compression=grpc.Compression.NoCompression,
|
|
)
|
|
self.call = self.channel.stream_stream(METHOD)(
|
|
metadata=access.metadata(), wait_for_ready=False
|
|
)
|
|
self.last_reply = -1
|
|
self._sending = self._reading = False
|
|
|
|
async def open(self) -> None:
|
|
try:
|
|
async with asyncio.timeout(2):
|
|
metadata = await self.call.initial_metadata()
|
|
if dict(metadata).get("x-stream-binding") != wire.binding(self.access.epoch):
|
|
raise wire.StreamWireError("server did not admit current input")
|
|
await self._send(
|
|
(
|
|
wire.open_packet(
|
|
self.access.epoch, self.access.session_id, self.access.session_generation
|
|
),
|
|
)
|
|
)
|
|
except BaseException:
|
|
self.call.cancel()
|
|
raise
|
|
|
|
async def _send(self, pieces: Iterable[bytes | memoryview]) -> None:
|
|
if self._sending:
|
|
self.call.cancel()
|
|
raise wire.StreamWireError("one source writer per stream")
|
|
self._sending = True
|
|
try:
|
|
async with asyncio.timeout(self.timeout):
|
|
for piece in pieces:
|
|
for offset in range(0, len(piece), INPUT_CHUNK):
|
|
await self.call.write(bytes(piece[offset : offset + INPUT_CHUNK]))
|
|
except BaseException:
|
|
self.call.cancel()
|
|
raise
|
|
finally:
|
|
self._sending = False
|
|
|
|
async def send(self, event: LiveIngressEvent) -> None:
|
|
if (event.session_id, event.session_generation) != (
|
|
self.access.session_id,
|
|
self.access.session_generation,
|
|
):
|
|
self.call.cancel()
|
|
raise wire.StreamWireError("source acquisition binding changed")
|
|
await self._send(
|
|
piece for pair in wire.event_packets(self.access.epoch, event) for piece in pair
|
|
)
|
|
|
|
async def end(self) -> None:
|
|
await self._send((wire.terminal_packet(self.access.epoch),))
|
|
await self.call.done_writing()
|
|
|
|
async def receive(self, *, timeout: float = 2) -> tuple[int, bytes] | None:
|
|
if not 0 < timeout <= 30:
|
|
raise ValueError("invalid bounded result deadline")
|
|
if self._reading:
|
|
self.call.cancel()
|
|
raise wire.StreamWireError("one result reader per stream")
|
|
self._reading = True
|
|
try:
|
|
async with asyncio.timeout(timeout):
|
|
raw = await self.call.read()
|
|
if raw is grpc.aio.EOF:
|
|
return None
|
|
sequence, payload = decode_reply(self.access.epoch, raw)
|
|
if sequence <= self.last_reply:
|
|
raise wire.StreamWireError("network result sequence regressed")
|
|
self.last_reply = sequence
|
|
return sequence, payload
|
|
except BaseException:
|
|
self.call.cancel()
|
|
raise
|
|
finally:
|
|
self._reading = False
|
|
|
|
async def close(self) -> None:
|
|
self.call.cancel()
|
|
await self.channel.close()
|