ringo-task-queue 0.1.0.dev6__py3-none-win_amd64.whl → 0.1.0.dev7__py3-none-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,1071 +1,1198 @@
1
- """Connection management and the user-facing client/queue APIs."""
2
-
3
- from __future__ import annotations
4
-
5
- import asyncio
6
- import os
7
- import secrets
8
- from datetime import timedelta
9
- from typing import TYPE_CHECKING, Any, Awaitable, Callable, Coroutine, Iterable, Literal
10
-
11
- import grpc
12
- from grpc_health.v1 import health_pb2, health_pb2_grpc
13
-
14
- from . import _pb
15
- from ._proto.ringo.v1 import queue_pb2, queue_pb2_grpc
16
- from ._version import PROTOCOL_MAJOR, PROTOCOL_MINOR, SDK_VERSION
17
- from .concurrency import AdaptiveConcurrency, ResourceConcurrency
18
- from .errors import (
19
- IncompatibleVersionError,
20
- InvalidArgumentError,
21
- RingoError,
22
- UnavailableError,
23
- map_rpc_error,
24
- )
25
- from .models import (
26
- BatchResult,
27
- EnqueueResult,
28
- PostgresStorage,
29
- SQLiteStorage,
30
- Task,
31
- TaskSpec,
32
- )
33
-
34
- if TYPE_CHECKING: # pragma: no cover
35
- from .daemon import LocalDaemon
36
- from .shared import SharedLocalDaemon
37
- from .worker import Worker
38
-
39
- SDK_NAME = "ringo-python"
40
-
41
- DEFAULT_RPC_TIMEOUT = 30.0
42
- _MAX_ENQUEUE_BATCH = 1000
43
- _RECONNECT_ATTEMPTS = 3
44
- _BACKOFF_SECONDS = (0.2, 0.6, 1.5)
45
-
46
- #: Supported values for the remote ``compression`` option: ``"none"``
47
- #: (the default) or ``"gzip"``. ``None`` means "not specified": the channel
48
- #: default applies and a legacy raw ``channel_options`` compression entry
49
- #: keeps its meaning. The runtime validation in
50
- #: :func:`_normalize_compression` still rejects any other value for dynamic
51
- #: callers.
52
- CompressionOption = Literal["none", "gzip"]
53
- _COMPRESSION_CHOICES: tuple[str, ...] = ("none", "gzip")
54
- #: gRPC channel option that controls the same thing as ``compression``. An
55
- #: explicit ``compression`` (even ``"none"``) combined with this raw option is
56
- #: an ambiguity error, mirroring the Node SDK.
57
- _COMPRESSION_CHANNEL_OPTION = "grpc.default_compression_algorithm"
58
-
59
-
60
- def _normalize_compression(value: CompressionOption | None) -> grpc.Compression | None:
61
- """Map the public ``compression`` option onto ``grpc.Compression``.
62
-
63
- ``None`` (unspecified) maps to ``None`` so the channel default applies.
64
- Anything else outside ``_COMPRESSION_CHOICES`` is a clear
65
- ``InvalidArgumentError`` raised before any network activity.
66
- """
67
- if value is None:
68
- return None
69
- if not isinstance(value, str) or value not in _COMPRESSION_CHOICES:
70
- raise InvalidArgumentError(
71
- f"unsupported compression {value!r}; expected one of"
72
- f" {', '.join(_COMPRESSION_CHOICES)}"
73
- )
74
- if value == "gzip":
75
- return grpc.Compression.Gzip
76
- return grpc.Compression.NoCompression
77
-
78
-
79
- def _check_compression_channel_option_conflict(
80
- channel_options: Iterable[tuple[str, Any]] | None,
81
- compression: grpc.Compression | None,
82
- ) -> None:
83
- """Reject an explicit ``compression`` alongside the raw gRPC option.
84
-
85
- When ``compression`` is unspecified (``None``), raw ``channel_options``
86
- keep their legacy meaning for advanced users.
87
- """
88
- if compression is None:
89
- return
90
- for key, _ in channel_options or ():
91
- if key == _COMPRESSION_CHANNEL_OPTION:
92
- raise InvalidArgumentError(
93
- f"compression and channel_options[{_COMPRESSION_CHANNEL_OPTION!r}]"
94
- " must not both be set"
95
- )
96
-
97
-
98
- def _consume_task_exception(task: asyncio.Task[Any]) -> None:
99
- if not task.cancelled():
100
- task.exception()
101
-
102
-
103
- def _parse_endpoint(endpoint: str) -> tuple[str, bool]:
104
- """Return ``(authority, use_tls)`` for an endpoint string."""
105
- if not endpoint:
106
- raise InvalidArgumentError("endpoint is required")
107
- use_tls = False
108
- rest = endpoint
109
- for scheme, tls in (
110
- ("grpc://", False),
111
- ("http://", False),
112
- ("grpcs://", True),
113
- ("https://", True),
114
- ):
115
- if rest.startswith(scheme):
116
- rest = rest[len(scheme) :]
117
- use_tls = tls
118
- break
119
- rest = rest.rstrip("/")
120
- if ":" not in rest:
121
- rest = f"{rest}:7233"
122
- return rest, use_tls
123
-
124
-
125
- #: Bounded automatic restarts of a crashed local daemon (design §31).
126
- _MAX_DAEMON_RESTARTS = 3
127
-
128
-
129
- class _Connection:
130
- """Owns the grpc.aio channel, negotiation, and retried unary calls."""
131
-
132
- def __init__(
133
- self,
134
- endpoint: str,
135
- token: str | None = None,
136
- *,
137
- local_token: str | None = None,
138
- local_generation: str | None = None,
139
- channel_options: Iterable[tuple[str, Any]] | None = None,
140
- compression: grpc.Compression | None = None,
141
- ) -> None:
142
- self._endpoint, self._tls = _parse_endpoint(endpoint)
143
- self._token = token
144
- self._local_token = local_token
145
- self._local_generation = local_generation
146
- self._attach_call: Any = None
147
- self._attach_reader_task: asyncio.Task[Any] | None = None
148
- self._shared_disconnect: Callable[[], Awaitable[None]] | None = None
149
- self._channel_options = list(channel_options or ())
150
- if local_token is not None:
151
- configured = {key for key, _ in self._channel_options}
152
- for key, value in (
153
- ("grpc.keepalive_time_ms", 15000),
154
- ("grpc.keepalive_timeout_ms", 10000),
155
- ("grpc.keepalive_permit_without_calls", 1),
156
- ):
157
- if key not in configured:
158
- self._channel_options.append((key, value))
159
- _check_compression_channel_option_conflict(
160
- self._channel_options, compression
161
- )
162
- # Channel-level default compression; survives every reconnect because
163
- # _open_channel() re-applies it to each new aio channel. ``None`` means
164
- # unspecified, so legacy raw channel options decide the encoding.
165
- self._compression = compression
166
- self._channel: grpc.aio.Channel | None = None
167
- self._stub: queue_pb2_grpc.QueueServiceStub | None = None
168
- self._session_minor = PROTOCOL_MINOR
169
- self._session_token = ""
170
- self._lock = asyncio.Lock()
171
- self._closed = False
172
- #: Optional client-provided reconnect hook (daemon restart + channel
173
- #: rebuild). Invoked under ``_lock`` instead of plain ``negotiate``.
174
- self.on_reconnect: Callable[[], Awaitable[None]] | None = None
175
-
176
- # -- lifecycle ---------------------------------------------------------
177
-
178
- def _open_channel(self) -> None:
179
- if self._tls:
180
- credentials = grpc.ssl_channel_credentials()
181
- self._channel = grpc.aio.secure_channel(
182
- self._endpoint,
183
- credentials,
184
- options=self._channel_options,
185
- compression=self._compression,
186
- )
187
- else:
188
- self._channel = grpc.aio.insecure_channel(
189
- self._endpoint,
190
- options=self._channel_options,
191
- compression=self._compression,
192
- )
193
- self._stub = queue_pb2_grpc.QueueServiceStub(self._channel)
194
-
195
- async def connect(self, *, check_health: bool = False) -> None:
196
- self._open_channel()
197
- if check_health:
198
- await self._await_serving()
199
- if self._local_token is not None:
200
- await self.attach_local()
201
- await self.negotiate()
202
-
203
- async def attach_local(self) -> None:
204
- """Register this SDK process as a shared-daemon member."""
205
- if self._local_token is None or self._local_generation is None:
206
- return
207
- stub = self._require_stub()
208
- try:
209
- request = queue_pb2.AttachLocalRequest(generation=self._local_generation)
210
- call = stub.AttachLocal(request, metadata=self._metadata())
211
- response = await call.read()
212
- if response is grpc.aio.EOF or response is None:
213
- raise UnavailableError("shared daemon closed AttachLocal during handshake")
214
- if response.generation != self._local_generation:
215
- call.cancel()
216
- raise UnavailableError("shared daemon generation changed during AttachLocal")
217
- except grpc.aio.AioRpcError as exc:
218
- raise map_rpc_error(exc) from exc
219
- self._attach_call = call
220
- self._attach_reader_task = asyncio.create_task(
221
- self._consume_attach(call), name="ringo-shared-attach"
222
- )
223
-
224
- async def _consume_attach(self, call: Any) -> None:
225
- try:
226
- while True:
227
- response = await call.read()
228
- if response is grpc.aio.EOF:
229
- break
230
- if response is None:
231
- break
232
- if response.generation != self._local_generation:
233
- break
234
- except asyncio.CancelledError:
235
- return
236
- except grpc.aio.AioRpcError:
237
- pass
238
- finally:
239
- if call is self._attach_call:
240
- self._attach_call = None
241
- if not self._closed and self._shared_disconnect is not None:
242
- asyncio.ensure_future(self._shared_disconnect())
243
-
244
- def set_shared_session(
245
- self, token: str, generation: str, on_disconnect: Callable[[], Awaitable[None]]
246
- ) -> None:
247
- self._local_token = token
248
- self._local_generation = generation
249
- self._shared_disconnect = on_disconnect
250
-
251
- async def _close_attach(self) -> None:
252
- task, self._attach_reader_task = self._attach_reader_task, None
253
- call, self._attach_call = self._attach_call, None
254
- if call is not None:
255
- try:
256
- call.cancel()
257
- except Exception:
258
- pass
259
- if task is not None and task is not asyncio.current_task():
260
- task.cancel()
261
- await asyncio.gather(task, return_exceptions=True)
262
-
263
- async def _await_serving(self, timeout: float = 5.0) -> None:
264
- """Wait until the health service reports SERVING (local daemon
265
- readiness is announced on stdout before the gRPC health state may be
266
- observable, so poll briefly instead of assuming)."""
267
- assert self._channel is not None
268
- health_stub = health_pb2_grpc.HealthStub(self._channel)
269
- deadline = asyncio.get_running_loop().time() + timeout
270
- last: Exception | None = None
271
- while True:
272
- try:
273
- response = await health_stub.Check(
274
- health_pb2.HealthCheckRequest(),
275
- timeout=2.0,
276
- metadata=self._metadata(),
277
- )
278
- if response.status == health_pb2.HealthCheckResponse.SERVING:
279
- return
280
- last = RingoError("daemon health status is not SERVING")
281
- except grpc.aio.AioRpcError as exc:
282
- last = exc
283
- if asyncio.get_running_loop().time() >= deadline:
284
- raise RingoError(f"daemon did not report SERVING health: {last}")
285
- await asyncio.sleep(0.05)
286
-
287
- async def reset(self, endpoint: str) -> None:
288
- """Point the connection at a new endpoint (local daemon restart)."""
289
- await self._close_attach()
290
- self._endpoint, self._tls = _parse_endpoint(endpoint)
291
- channel, self._channel = self._channel, None
292
- self._stub = None
293
- if channel is not None:
294
- await channel.close()
295
- self._open_channel()
296
-
297
- async def negotiate(self) -> None:
298
- """Perform the major=1/minor=0 version handshake (ADR 0002)."""
299
- stub = self._require_stub()
300
- request = queue_pb2.NegotiateRequest(
301
- sdk_name=SDK_NAME,
302
- sdk_version=SDK_VERSION,
303
- protocol_major=PROTOCOL_MAJOR,
304
- protocol_minor=PROTOCOL_MINOR,
305
- )
306
- try:
307
- response = await stub.Negotiate(
308
- request, timeout=DEFAULT_RPC_TIMEOUT, metadata=self._metadata()
309
- )
310
- except grpc.aio.AioRpcError as exc:
311
- raise map_rpc_error(exc) from exc
312
- if response.protocol_major != PROTOCOL_MAJOR:
313
- raise IncompatibleVersionError(
314
- f"server speaks protocol major {response.protocol_major},"
315
- f" SDK requires major {PROTOCOL_MAJOR}; upgrade one side"
316
- )
317
- self._session_minor = min(response.protocol_minor, PROTOCOL_MINOR)
318
- self._session_token = response.session_token
319
-
320
- async def close(self) -> None:
321
- self._closed = True
322
- await self._close_attach()
323
- channel, self._channel = self._channel, None
324
- self._stub = None
325
- if channel is not None:
326
- await channel.close()
327
-
328
- # -- helpers -----------------------------------------------------------
329
-
330
- def _require_stub(self) -> queue_pb2_grpc.QueueServiceStub:
331
- if self._stub is None:
332
- raise RingoError("connection is closed")
333
- return self._stub
334
-
335
- def _metadata(self) -> list[tuple[str, str]]:
336
- metadata: list[tuple[str, str]] = []
337
- if self._token:
338
- metadata.append(("authorization", f"Bearer {self._token}"))
339
- if self._local_token:
340
- metadata.append(("x-ringo-local-token", self._local_token))
341
- return metadata
342
-
343
- def request_context(self) -> queue_pb2.RequestContext:
344
- return queue_pb2.RequestContext(
345
- protocol_major=PROTOCOL_MAJOR,
346
- protocol_minor=self._session_minor,
347
- session_token=self._session_token,
348
- )
349
-
350
- @property
351
- def endpoint(self) -> str:
352
- return self._endpoint
353
-
354
- async def call(
355
- self,
356
- name: str,
357
- request: Any,
358
- *,
359
- timeout: float = DEFAULT_RPC_TIMEOUT,
360
- retry_safe: bool = False,
361
- context_attr: str | None = "context",
362
- ) -> Any:
363
- """Issue a single unary RPC.
364
-
365
- Submission is never silently replayed: an ``UNAVAILABLE`` may mean the
366
- server committed the operation (a lost ACK response, a claimed lease)
367
- so the original error is raised after best-effort recovery (daemon
368
- restart + fresh negotiation) prepares the connection for *future*
369
- calls. Only genuinely read-only RPCs (``GetTask``) pass
370
- ``retry_safe=True`` and may be re-sent, each time with a freshly
371
- negotiated ``RequestContext``.
372
- """
373
- rpc, sent_token = await self._begin_unary(
374
- name, request, timeout=timeout, context_attr=context_attr
375
- )
376
- cause: BaseException | None = None
377
- err: RingoError | None = None
378
- try:
379
- return await rpc
380
- except grpc.aio.AioRpcError as exc:
381
- cause = exc
382
- err = map_rpc_error(exc)
383
- except asyncio.CancelledError as exc:
384
- # grpc.aio raises CancelledError when a reconnect closes its old
385
- # channel. Preserve genuine caller cancellation, but surface this
386
- # transport race as UNAVAILABLE.
387
- task = asyncio.current_task()
388
- if task is not None and task.cancelling():
389
- raise
390
- cause = exc
391
- err = UnavailableError("RPC interrupted by connection recovery")
392
-
393
- assert err is not None and cause is not None
394
- if self._closed:
395
- raise err from cause
396
- if isinstance(err, IncompatibleVersionError):
397
- # Session rejection happens in authorization before the operation
398
- # can execute. It is therefore safe to negotiate a fresh short-lived
399
- # token and retry once, including ordinary token expiry with no
400
- # concurrent daemon restart.
401
- await self._recover_best_effort()
402
- async with self._lock:
403
- rotated = self._session_token != sent_token
404
- if rotated:
405
- rpc, _ = await self._begin_unary(
406
- name, request, timeout=timeout, context_attr=context_attr
407
- )
408
- try:
409
- return await rpc
410
- except grpc.aio.AioRpcError as retry_exc:
411
- raise map_rpc_error(retry_exc) from retry_exc
412
- raise err from cause
413
- if not isinstance(err, UnavailableError):
414
- raise err from cause
415
-
416
- # Recovery prepares the connection for future calls. Mutations are not
417
- # replayed because an unavailable response may hide a committed write.
418
- await self._recover_best_effort()
419
- if not retry_safe:
420
- raise err from cause
421
- rpc, _ = await self._begin_unary(
422
- name, request, timeout=timeout, context_attr=context_attr
423
- )
424
- try:
425
- return await rpc
426
- except grpc.aio.AioRpcError as retry_exc:
427
- raise map_rpc_error(retry_exc) from retry_exc
428
-
429
- async def _begin_unary(
430
- self,
431
- name: str,
432
- request: Any,
433
- *,
434
- timeout: float,
435
- context_attr: str | None,
436
- ) -> tuple[Any, str]:
437
- """Snapshot a stub and fresh session context outside reconnect races."""
438
- async with self._lock:
439
- if context_attr is not None:
440
- getattr(request, context_attr).CopyFrom(self.request_context())
441
- stub = self._require_stub()
442
- token = self._session_token
443
- rpc = getattr(stub, name)(
444
- request, timeout=timeout, metadata=self._metadata()
445
- )
446
- return rpc, token
447
-
448
- async def _recover_best_effort(self) -> None:
449
- """Bounded reconnect so later calls can succeed; failures here are
450
- swallowed because the original call error is what the caller sees."""
451
- for attempt in range(_RECONNECT_ATTEMPTS):
452
- try:
453
- await asyncio.sleep(
454
- _BACKOFF_SECONDS[min(attempt, len(_BACKOFF_SECONDS) - 1)]
455
- )
456
- async with self._lock:
457
- if self._closed:
458
- return
459
- if self.on_reconnect is not None:
460
- await self.on_reconnect()
461
- else:
462
- await self.negotiate()
463
- return
464
- except Exception:
465
- continue
466
-
467
-
468
- def new_worker_id() -> str:
469
- return f"py-{os.getpid()}-{secrets.token_hex(4)}"
470
-
471
-
472
- # ---------------------------------------------------------------------------
473
- # Public client
474
-
475
-
476
- class RingoClient:
477
- """A connected (or connectable) Ringo client.
478
-
479
- Use as an async context manager, or call :meth:`start` / :meth:`close`
480
- explicitly::
481
-
482
- async with Ringo.connect("127.0.0.1:7233") as client:
483
- queue = client.queue("crawl")
484
- """
485
-
486
- def __init__(
487
- self,
488
- *,
489
- endpoint: str | None = None,
490
- token: str | None = None,
491
- daemon: "LocalDaemon | None" = None,
492
- shared_daemon: "SharedLocalDaemon | None" = None,
493
- channel_options: Iterable[tuple[str, Any]] | None = None,
494
- compression: CompressionOption | None = None,
495
- ) -> None:
496
- self._endpoint = endpoint
497
- self._token = token
498
- self._channel_options = channel_options
499
- # The value is always validated, but ``compression`` only applies to
500
- # remote connections; local mode stays on the (uncompressed) default.
501
- normalized = _normalize_compression(compression)
502
- if daemon is None and shared_daemon is None:
503
- _check_compression_channel_option_conflict(
504
- self._channel_options, normalized
505
- )
506
- self._compression = normalized
507
- self._compression_label: CompressionOption = (
508
- compression if compression is not None else "none"
509
- )
510
- else:
511
- # Local mode never enables compression; an explicit value is
512
- # validated above and then ignored (source compatibility).
513
- self._compression = None
514
- self._compression_label = "none"
515
- self._daemon = daemon
516
- self._shared_daemon = shared_daemon
517
- if daemon is not None and shared_daemon is not None:
518
- raise InvalidArgumentError("daemon and shared_daemon are mutually exclusive")
519
- # ``_conn`` is built in start(): a local client learns the daemon's
520
- # random loopback port only after the daemon announces readiness.
521
- self._conn: _Connection | None = None
522
- self._queues: dict[str, TaskQueue] = {}
523
- self._workers: set[Worker] = set()
524
- self._started = False
525
- self._closed = False
526
- self._start_lock = asyncio.Lock()
527
- self._close_lock = asyncio.Lock()
528
- self._close_task: asyncio.Task[None] | None = None
529
- self._daemon_restart_count = 0
530
-
531
- # -- lifecycle ---------------------------------------------------------
532
-
533
- async def start(self) -> "RingoClient":
534
- async with self._start_lock:
535
- if self._closed:
536
- raise RingoError("client is closed")
537
- if not self._started:
538
- try:
539
- for attempt in range(2 if self._shared_daemon is not None else 1):
540
- if self._daemon is not None:
541
- self._endpoint = await self._daemon.start()
542
- elif self._shared_daemon is not None:
543
- self._endpoint = await self._shared_daemon.start()
544
- conn = _Connection(
545
- self._endpoint or "",
546
- self._token,
547
- local_token=(self._shared_daemon.token if self._shared_daemon else None),
548
- local_generation=(self._shared_daemon.generation if self._shared_daemon else None),
549
- channel_options=self._channel_options,
550
- compression=self._compression,
551
- )
552
- self._conn = conn
553
- conn.on_reconnect = self._reconnect
554
- if self._shared_daemon is not None:
555
- conn.set_shared_session(
556
- self._shared_daemon.token or "",
557
- self._shared_daemon.generation or "",
558
- self._shared_stream_lost,
559
- )
560
- try:
561
- await conn.connect(
562
- check_health=self._daemon is not None or self._shared_daemon is not None
563
- )
564
- break
565
- except Exception:
566
- await conn.close()
567
- self._conn = None
568
- if self._shared_daemon is None or attempt > 0:
569
- raise
570
- # Stale PID/port, wrong token or generation is never
571
- # accepted as identity. Re-elect once, boundedly.
572
- self._shared_daemon.invalidate_current()
573
- except BaseException:
574
- # Fast cleanup: never leak a half-started channel or local daemon.
575
- failed_conn = self._conn
576
- self._conn = None
577
- if failed_conn is not None:
578
- await failed_conn.close()
579
- await self._stop_daemon()
580
- raise
581
- self._started = True
582
- return self
583
-
584
- def _require_conn(self) -> _Connection:
585
- conn = self._conn
586
- if conn is None:
587
- raise RingoError(
588
- "client is not started; use 'async with' or await client.start()"
589
- )
590
- return conn
591
-
592
- async def _reconnect(self) -> None:
593
- """Reconnect after a failed transport and negotiate a fresh session."""
594
- conn = self._require_conn()
595
- daemon = self._daemon
596
- shared = self._shared_daemon
597
- if daemon is not None and daemon.is_dead:
598
- if self._daemon_restart_count >= _MAX_DAEMON_RESTARTS:
599
- raise UnavailableError(
600
- f"local daemon restart budget ({_MAX_DAEMON_RESTARTS}) exhausted"
601
- )
602
- self._daemon_restart_count += 1
603
- delay = 0.3 * (2 ** (self._daemon_restart_count - 1))
604
- await asyncio.sleep(delay)
605
- endpoint = await daemon.start()
606
- await conn.reset(endpoint)
607
- elif shared is not None:
608
- try:
609
- endpoint = await shared.refresh()
610
- await conn.reset(endpoint)
611
- conn.set_shared_session(
612
- shared.token or "", shared.generation or "", self._shared_stream_lost
613
- )
614
- await conn.attach_local()
615
- await conn.negotiate()
616
- except Exception:
617
- shared.invalidate_current()
618
- raise
619
- return
620
- await conn.negotiate()
621
-
622
- async def _shared_stream_lost(self) -> None:
623
- """Proactively reconnect an idle shared client after AttachLocal EOF."""
624
- if self._closed or not self._started or self._shared_daemon is None:
625
- return
626
- conn = self._conn
627
- if conn is None or conn._closed:
628
- return
629
- # A stream can report EOF more than once while a replacement is being
630
- # opened. Serialize recovery through the connection's normal lock.
631
- for attempt in range(_MAX_DAEMON_RESTARTS + 1):
632
- if self._closed:
633
- return
634
- try:
635
- await asyncio.sleep(_BACKOFF_SECONDS[min(attempt, len(_BACKOFF_SECONDS) - 1)])
636
- async with conn._lock:
637
- await self._reconnect()
638
- return
639
- except Exception:
640
- continue
641
-
642
- async def __aenter__(self) -> "RingoClient":
643
- return await self.start()
644
-
645
- async def close(self) -> None:
646
- """Drain and stop once; caller cancellation cannot cancel cleanup."""
647
- async with self._close_lock:
648
- if self._close_task is None:
649
- self._closed = True
650
- self._close_task = asyncio.create_task(
651
- self._close_impl(), name="ringo-client-close"
652
- )
653
- self._close_task.add_done_callback(_consume_task_exception)
654
- close_task = self._close_task
655
- await asyncio.shield(close_task)
656
-
657
- async def _close_impl(self) -> None:
658
- """Attempt every cleanup step and aggregate any failures."""
659
- errors: list[Exception] = []
660
- # 1. Drain workers (stop taking tasks, wait grace, cancel remainder).
661
- for worker in list(self._workers):
662
- try:
663
- await worker.close()
664
- except Exception as exc: # keep cleaning up the rest
665
- errors.append(exc)
666
- self._workers.clear()
667
- # 2. Close the gRPC channel.
668
- conn, self._conn = self._conn, None
669
- if conn is not None:
670
- try:
671
- await conn.close()
672
- except Exception as exc:
673
- errors.append(exc)
674
- # 3. Shut the local daemon down via the lifecycle pipe, then wait.
675
- try:
676
- await self._stop_daemon()
677
- except Exception as exc:
678
- errors.append(exc)
679
- if errors:
680
- raise ExceptionGroup("errors while closing Ringo client", errors)
681
-
682
- async def _stop_daemon(self) -> None:
683
- daemon, self._daemon = self._daemon, None
684
- if daemon is not None:
685
- await daemon.stop()
686
-
687
- async def __aexit__(self, exc_type, exc, tb) -> None:
688
- try:
689
- await self.close()
690
- except ExceptionGroup as cleanup_errors:
691
- # A business exception takes precedence; cleanup failures must
692
- # not mask it, but should not vanish either.
693
- if exc is None:
694
- raise
695
- for error in cleanup_errors.exceptions:
696
- exc.__cause__ = exc.__cause__ or error
697
-
698
- @property
699
- def is_local(self) -> bool:
700
- return self._daemon is not None or self._shared_daemon is not None
701
-
702
- @property
703
- def compression(self) -> CompressionOption:
704
- """The effective channel compression: ``"none"`` (default) or ``"gzip"``."""
705
- return self._compression_label
706
-
707
- # -- queues ------------------------------------------------------------
708
-
709
- def queue(self, name: str) -> "TaskQueue":
710
- if not name:
711
- raise InvalidArgumentError("queue name is required")
712
- if self._closed:
713
- raise RingoError("client is closed")
714
- existing = self._queues.get(name)
715
- if existing is None:
716
- existing = TaskQueue(self, name)
717
- self._queues[name] = existing
718
- return existing
719
-
720
- def _register_worker(self, worker: "Worker") -> None:
721
- self._workers.add(worker)
722
-
723
- def _unregister_worker(self, worker: "Worker") -> None:
724
- self._workers.discard(worker)
725
-
726
-
727
- class Ringo:
728
- """Factory for local (embedded daemon) and remote clients.
729
-
730
- Both factories are synchronous and return a client whose
731
- ``async with``/``await start()`` performs the actual connect and
732
- negotiation, so ``async with Ringo.local(...)`` works directly.
733
- """
734
-
735
- @staticmethod
736
- def local(
737
- data_dir: str | os.PathLike[str],
738
- storage: str | SQLiteStorage | PostgresStorage = "sqlite",
739
- daemon_path: str | os.PathLike[str] | None = None,
740
- *,
741
- shared: bool = False,
742
- startup_timeout: float = 15.0,
743
- stop_timeout: float = 10.0,
744
- ) -> RingoClient:
745
- """Start a private daemon, or opt into a shared SQLite daemon."""
746
- if not isinstance(shared, bool):
747
- raise InvalidArgumentError("shared must be a boolean")
748
- if storage == "sqlite" or isinstance(storage, SQLiteStorage):
749
- storage_name = "sqlite"
750
- postgres_dsn = None
751
- elif isinstance(storage, PostgresStorage):
752
- storage_name = "postgres"
753
- postgres_dsn = storage.dsn
754
- else:
755
- raise InvalidArgumentError(f"unsupported local storage {storage!r}")
756
- if shared and storage_name != "sqlite":
757
- raise InvalidArgumentError("shared local mode supports SQLite only")
758
- # Imported here so remote-only users never pay for subprocess imports.
759
- if shared:
760
- from .shared import SharedLocalDaemon
761
-
762
- shared_daemon = SharedLocalDaemon(
763
- data_dir, daemon_path=daemon_path, startup_timeout=startup_timeout
764
- )
765
- return RingoClient(shared_daemon=shared_daemon)
766
- from .daemon import LocalDaemon
767
-
768
- daemon = LocalDaemon(
769
- data_dir,
770
- storage=storage_name,
771
- postgres_dsn=postgres_dsn,
772
- daemon_path=daemon_path,
773
- startup_timeout=startup_timeout,
774
- stop_timeout=stop_timeout,
775
- )
776
- return RingoClient(daemon=daemon)
777
-
778
- @staticmethod
779
- def connect(
780
- endpoint: str,
781
- token: str | None = None,
782
- *,
783
- channel_options: Iterable[tuple[str, Any]] | None = None,
784
- compression: CompressionOption | None = None,
785
- ) -> RingoClient:
786
- """Create a client for a remote (or already-running) endpoint.
787
-
788
- ``compression`` is an optional whole-channel setting: ``"none"``
789
- (the default) or ``"gzip"``. ``None`` means "not specified", leaving
790
- the gRPC channel default (and any raw ``channel_options`` compression
791
- entry) in charge. An explicit value combined with the raw channel
792
- option ``grpc.default_compression_algorithm`` raises
793
- :class:`InvalidArgumentError`; any other value does so as well,
794
- immediately and before any network activity.
795
- """
796
- return RingoClient(
797
- endpoint=endpoint,
798
- token=token,
799
- channel_options=channel_options,
800
- compression=compression,
801
- )
802
-
803
-
804
- # ---------------------------------------------------------------------------
805
- # TaskQueue
806
-
807
-
808
- Handler = Callable[[Task, Any], Coroutine[Any, Any, None]]
809
-
810
-
811
- class TaskQueue:
812
- """A client bound to one logical queue (``queue + uniqueKey`` dedupe)."""
813
-
814
- def __init__(self, client: RingoClient, name: str) -> None:
815
- self._client = client
816
- self.name = name
817
-
818
- # -- enqueue ------------------------------------------------------------
819
-
820
- async def enqueue(
821
- self,
822
- spec: TaskSpec | None = None,
823
- /,
824
- **kwargs: Any,
825
- ) -> EnqueueResult:
826
- """Enqueue one task; duplicates return ``created=False`` (no raise)."""
827
- if spec is None:
828
- spec = TaskSpec(**kwargs)
829
- request = queue_pb2.EnqueueRequest(queue=self.name, task=spec.to_proto())
830
- request.context.CopyFrom(self._client._require_conn().request_context())
831
- response = await self._client._require_conn().call("Enqueue", request)
832
- return EnqueueResult.from_proto(response.result)
833
-
834
- async def enqueue_many(self, specs: Iterable[TaskSpec]) -> BatchResult:
835
- """Enqueue any number of tasks, auto-batching under the server limit.
836
-
837
- The server enforces a hard limit of ``_MAX_ENQUEUE_BATCH`` (1000)
838
- tasks per gRPC call. This method transparently splits larger
839
- iterables into sequential sub-batches, preserving creation order,
840
- and aggregates the per-item results into a single :class:`BatchResult`.
841
-
842
- Fail-fast: if any sub-batch fails, the exception propagates and no
843
- partial result is returned.
844
- """
845
- specs = list(specs)
846
- all_results: list[EnqueueResult] = []
847
- total_created = 0
848
- total_duplicates = 0
849
-
850
- for start in range(0, len(specs), _MAX_ENQUEUE_BATCH):
851
- chunk = specs[start : start + _MAX_ENQUEUE_BATCH]
852
- request = queue_pb2.EnqueueManyRequest(
853
- queue=self.name, tasks=[s.to_proto() for s in chunk]
854
- )
855
- request.context.CopyFrom(self._client._require_conn().request_context())
856
- response = await self._client._require_conn().call("EnqueueMany", request)
857
- all_results.extend(EnqueueResult.from_proto(r) for r in response.results)
858
- total_created += response.created_count
859
- total_duplicates += response.duplicate_count
860
-
861
- return BatchResult(
862
- results=all_results,
863
- created_count=total_created,
864
- duplicate_count=total_duplicates,
865
- )
866
-
867
- # -- queries ------------------------------------------------------------
868
-
869
- async def get_by_id(self, task_id: str) -> Task:
870
- if not task_id:
871
- raise InvalidArgumentError("task_id is required")
872
- request = queue_pb2.GetTaskRequest(queue=self.name, task_id=task_id)
873
- request.context.CopyFrom(self._client._require_conn().request_context())
874
- response = await self._client._require_conn().call(
875
- "GetTask", request, retry_safe=True
876
- )
877
- return Task.from_proto(response.task)
878
-
879
- async def get_by_unique_key(self, unique_key: str) -> Task:
880
- if not unique_key:
881
- raise InvalidArgumentError("unique_key is required")
882
- request = queue_pb2.GetTaskRequest(queue=self.name, unique_key=unique_key)
883
- request.context.CopyFrom(self._client._require_conn().request_context())
884
- response = await self._client._require_conn().call(
885
- "GetTask", request, retry_safe=True
886
- )
887
- return Task.from_proto(response.task)
888
-
889
- async def retry(self, unique_key: str) -> Task:
890
- """Reactivate a SUCCEEDED/DEAD task by its business unique key.
891
-
892
- This always resolves through the proto ``unique_key`` field; use
893
- :meth:`retry_by_id` for the task-ID form. No string is ever guessed.
894
- """
895
- if not unique_key:
896
- raise InvalidArgumentError("unique_key is required")
897
- request = queue_pb2.RetryTaskRequest(queue=self.name)
898
- request.unique_key = unique_key
899
- request.context.CopyFrom(self._client._require_conn().request_context())
900
- response = await self._client._require_conn().call("RetryTask", request)
901
- return Task.from_proto(response.task)
902
-
903
- async def retry_by_id(self, task_id: str) -> Task:
904
- """Reactivate a SUCCEEDED/DEAD task by its server task ID."""
905
- if not task_id:
906
- raise InvalidArgumentError("task_id is required")
907
- request = queue_pb2.RetryTaskRequest(queue=self.name)
908
- request.task_id = task_id
909
- request.context.CopyFrom(self._client._require_conn().request_context())
910
- response = await self._client._require_conn().call("RetryTask", request)
911
- return Task.from_proto(response.task)
912
-
913
- # -- manual claim ---------------------------------------------------------
914
-
915
- async def claim(
916
- self,
917
- limit: int = 1,
918
- *,
919
- task_types: Iterable[str] | None = None,
920
- wait_timeout: timedelta | None = None,
921
- lease_duration: timedelta | None = None,
922
- worker_id: str | None = None,
923
- ) -> list["LeasedTask"]:
924
- """Manually claim tasks (advanced API; not the default worker path)."""
925
- if limit < 1 or limit > _MAX_ENQUEUE_BATCH:
926
- raise InvalidArgumentError(
927
- f"limit must be between 1 and {_MAX_ENQUEUE_BATCH}"
928
- )
929
- if wait_timeout is not None and wait_timeout < timedelta(0):
930
- raise InvalidArgumentError("wait_timeout must not be negative")
931
- if lease_duration is not None and lease_duration < timedelta(0):
932
- raise InvalidArgumentError("lease_duration must not be negative")
933
- request = queue_pb2.ClaimRequest(
934
- queue=self.name,
935
- worker_id=worker_id or new_worker_id(),
936
- task_types=list(task_types or []),
937
- limit=limit,
938
- )
939
- if wait_timeout is not None:
940
- request.wait_timeout.CopyFrom(_pb.td_to_dur(wait_timeout))
941
- if lease_duration is not None:
942
- request.lease_duration.CopyFrom(_pb.td_to_dur(lease_duration))
943
- request.context.CopyFrom(self._client._require_conn().request_context())
944
- response = await self._client._require_conn().call("Claim", request)
945
- return [LeasedTask(self, t) for t in response.tasks]
946
-
947
- # -- worker ---------------------------------------------------------------
948
-
949
- def worker(
950
- self,
951
- handler: Handler,
952
- *,
953
- task_types: Iterable[str] | None = None,
954
- concurrency: int | ResourceConcurrency | AdaptiveConcurrency = 1,
955
- weight: int = 1,
956
- lease_duration: timedelta = timedelta(seconds=60),
957
- grace_period: timedelta = timedelta(seconds=30),
958
- on_idle: Callable[[], Any] | None = None,
959
- ) -> "Worker":
960
- """Create a credit-controlled worker over the frozen Work stream.
961
-
962
- ``on_idle`` is an optional zero-argument callable (sync or awaitable)
963
- for bounded batch runs: once the worker has registered and sent its
964
- initial credit, the SDK polls queue activity every 250 ms via the
965
- read-only ``GetQueueActivity`` RPC scoped to this worker's
966
- ``task_types``; after the activity count has been continuously zero
967
- for at least 1 s the callback fires exactly once, inline on the
968
- ``run()`` coroutine. A callback that returns normally leaves the
969
- worker running (and never fires again); any exception it raises —
970
- including ``SystemExit`` — propagates verbatim out of
971
- ``await worker.run()`` without becoming a task NACK. Workers without
972
- ``on_idle`` never send activity RPCs and stay compatible with
973
- protocol minor 0 servers; ``on_idle`` requires minor >= 1 and raises
974
- :class:`IncompatibleVersionError` from ``run()`` otherwise. Never
975
- call ``worker.close()`` from inside the callback (it would wait on
976
- the very coroutine executing it).
977
- """
978
- from .worker import Worker
979
-
980
- if not callable(handler):
981
- raise InvalidArgumentError("handler must be callable")
982
- if on_idle is not None and not callable(on_idle):
983
- raise InvalidArgumentError("on_idle must be callable")
984
- if isinstance(concurrency, bool) or not isinstance(
985
- concurrency, (int, ResourceConcurrency, AdaptiveConcurrency)
986
- ):
987
- raise InvalidArgumentError(
988
- "concurrency must be an integer, ResourceConcurrency, or AdaptiveConcurrency"
989
- )
990
- if isinstance(concurrency, int) and (concurrency < 1 or concurrency > 2**31 - 1):
991
- raise InvalidArgumentError(
992
- "initial fixed concurrency must be between 1 and 2**31-1"
993
- )
994
- if isinstance(weight, bool) or not isinstance(weight, int) or weight < 1:
995
- raise InvalidArgumentError("weight must be a positive integer")
996
- if lease_duration < timedelta(0):
997
- raise InvalidArgumentError("lease_duration must not be negative")
998
- if grace_period < timedelta(0):
999
- raise InvalidArgumentError("grace_period must not be negative")
1000
- worker = Worker(
1001
- self,
1002
- handler,
1003
- task_types=list(task_types or []),
1004
- concurrency=concurrency,
1005
- lease_duration=lease_duration,
1006
- grace_period=grace_period,
1007
- on_idle=on_idle,
1008
- weight=weight,
1009
- )
1010
- self._client._register_worker(worker)
1011
- return worker
1012
-
1013
-
1014
- # ---------------------------------------------------------------------------
1015
- # Manual lease
1016
-
1017
-
1018
- class LeasedTask:
1019
- """A task claimed via :meth:`TaskQueue.claim` with its current lease."""
1020
-
1021
- def __init__(self, queue: TaskQueue, msg: queue_pb2.LeasedTask) -> None:
1022
- self._queue = queue
1023
- self.task = Task.from_proto(msg.task)
1024
- self.attempt: int = msg.lease.attempt
1025
- self.worker_id: str = msg.lease.worker_id
1026
- self._token: str = msg.lease.lease_token
1027
- self._lease_until = _pb.ts_to_dt(msg.lease.until)
1028
-
1029
- @property
1030
- def lease_until(self):
1031
- return self._lease_until
1032
-
1033
- def _ref(self) -> queue_pb2.LeaseRef:
1034
- return queue_pb2.LeaseRef(
1035
- task_id=self.task.id,
1036
- attempt=self.attempt,
1037
- lease_token=self._token,
1038
- )
1039
-
1040
- async def _terminal(self, name: str, request: Any) -> None:
1041
- request.lease.CopyFrom(self._ref())
1042
- conn = self._queue._client._require_conn()
1043
- request.context.CopyFrom(conn.request_context())
1044
- await conn.call(name, request)
1045
-
1046
- async def ack(self) -> None:
1047
- """Acknowledge success for the current attempt/lease token."""
1048
- await self._terminal("Ack", queue_pb2.AckRequest())
1049
-
1050
- async def nack(self, error: str = "", *, delay: timedelta | None = None) -> None:
1051
- """Negative-acknowledge; the task is retried per policy."""
1052
- request = queue_pb2.NackRequest(error=error)
1053
- if delay is not None:
1054
- request.retry_delay_override.CopyFrom(_pb.td_to_dur(delay))
1055
- await self._terminal("Nack", request)
1056
-
1057
- async def reject(self, error: str = "") -> None:
1058
- """Permanently fail the task (DEAD) for the current attempt."""
1059
- await self._terminal("Reject", queue_pb2.RejectRequest(error=error))
1060
-
1061
- async def extend_lease(self, extension: timedelta) -> None:
1062
- """Extend the lease to ``server_now + extension`` (absolute reset)."""
1063
- if extension <= timedelta(0):
1064
- raise InvalidArgumentError("extension must be positive")
1065
- request = queue_pb2.ExtendLeaseRequest()
1066
- request.extension.CopyFrom(_pb.td_to_dur(extension))
1067
- request.lease.CopyFrom(self._ref())
1068
- conn = self._queue._client._require_conn()
1069
- request.context.CopyFrom(conn.request_context())
1070
- response = await conn.call("ExtendLease", request)
1071
- self._lease_until = _pb.ts_to_dt(response.lease.until)
1
+ """Connection management and the user-facing client/queue APIs."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import os
7
+ import secrets
8
+ from datetime import timedelta
9
+ from itertools import islice
10
+ from typing import TYPE_CHECKING, Any, Awaitable, Callable, Coroutine, Iterable, Literal
11
+
12
+ import grpc
13
+ from grpc_health.v1 import health_pb2, health_pb2_grpc
14
+
15
+ from . import _pb
16
+ from ._proto.ringo.v1 import queue_pb2, queue_pb2_grpc
17
+ from ._version import PROTOCOL_MAJOR, PROTOCOL_MINOR, SDK_VERSION
18
+ from .concurrency import AdaptiveConcurrency, ResourceConcurrency
19
+ from .errors import (
20
+ IncompatibleVersionError,
21
+ InvalidArgumentError,
22
+ RingoError,
23
+ UnavailableError,
24
+ map_rpc_error,
25
+ )
26
+ from .models import (
27
+ AttemptPage,
28
+ BatchResult,
29
+ EnqueueResult,
30
+ PostgresStorage,
31
+ QueueStats,
32
+ SQLiteStorage,
33
+ Task,
34
+ TaskSpec,
35
+ )
36
+
37
+ if TYPE_CHECKING: # pragma: no cover
38
+ from .daemon import LocalDaemon
39
+ from .shared import SharedLocalDaemon
40
+ from .worker import Worker
41
+
42
+ SDK_NAME = "ringo-python"
43
+
44
+ DEFAULT_RPC_TIMEOUT = 30.0
45
+ _DEFAULT_MAX_MESSAGE_BYTES = 16 * 1024 * 1024
46
+ _MAX_ENQUEUE_BATCH = 1000
47
+ _RECONNECT_ATTEMPTS = 3
48
+ _BACKOFF_SECONDS = (0.2, 0.6, 1.5)
49
+ #: GetQueueStats is additive protocol 1.4; older servers reject the RPC.
50
+ _STATS_MIN_PROTOCOL_MINOR = 4
51
+ #: Maximum number of task-type filters accepted by ``TaskQueue.stats``.
52
+ _STATS_MAX_TASK_TYPES = 100
53
+ _ATTEMPT_HISTORY_MIN_PROTOCOL_MINOR = 5
54
+ _ATTEMPT_PAGE_DEFAULT_LIMIT = 32
55
+ _ATTEMPT_PAGE_MAX_LIMIT = 100
56
+ _MAX_ATTEMPT_CURSOR_BYTES = 1024
57
+
58
+ #: Supported values for the remote ``compression`` option: ``"none"``
59
+ #: (the default) or ``"gzip"``. ``None`` means "not specified": the channel
60
+ #: default applies and a legacy raw ``channel_options`` compression entry
61
+ #: keeps its meaning. The runtime validation in
62
+ #: :func:`_normalize_compression` still rejects any other value for dynamic
63
+ #: callers.
64
+ CompressionOption = Literal["none", "gzip"]
65
+ _COMPRESSION_CHOICES: tuple[str, ...] = ("none", "gzip")
66
+ #: gRPC channel option that controls the same thing as ``compression``. An
67
+ #: explicit ``compression`` (even ``"none"``) combined with this raw option is
68
+ #: an ambiguity error, mirroring the Node SDK.
69
+ _COMPRESSION_CHANNEL_OPTION = "grpc.default_compression_algorithm"
70
+
71
+
72
+ def _normalize_compression(value: CompressionOption | None) -> grpc.Compression | None:
73
+ """Map the public ``compression`` option onto ``grpc.Compression``.
74
+
75
+ ``None`` (unspecified) maps to ``None`` so the channel default applies.
76
+ Anything else outside ``_COMPRESSION_CHOICES`` is a clear
77
+ ``InvalidArgumentError`` raised before any network activity.
78
+ """
79
+ if value is None:
80
+ return None
81
+ if not isinstance(value, str) or value not in _COMPRESSION_CHOICES:
82
+ raise InvalidArgumentError(
83
+ f"unsupported compression {value!r}; expected one of"
84
+ f" {', '.join(_COMPRESSION_CHOICES)}"
85
+ )
86
+ if value == "gzip":
87
+ return grpc.Compression.Gzip
88
+ return grpc.Compression.NoCompression
89
+
90
+
91
+ def _check_compression_channel_option_conflict(
92
+ channel_options: Iterable[tuple[str, Any]] | None,
93
+ compression: grpc.Compression | None,
94
+ ) -> None:
95
+ """Reject an explicit ``compression`` alongside the raw gRPC option.
96
+
97
+ When ``compression`` is unspecified (``None``), raw ``channel_options``
98
+ keep their legacy meaning for advanced users.
99
+ """
100
+ if compression is None:
101
+ return
102
+ for key, _ in channel_options or ():
103
+ if key == _COMPRESSION_CHANNEL_OPTION:
104
+ raise InvalidArgumentError(
105
+ f"compression and channel_options[{_COMPRESSION_CHANNEL_OPTION!r}]"
106
+ " must not both be set"
107
+ )
108
+
109
+
110
+ def _consume_task_exception(task: asyncio.Task[Any]) -> None:
111
+ if not task.cancelled():
112
+ task.exception()
113
+
114
+
115
+ def _parse_endpoint(endpoint: str) -> tuple[str, bool]:
116
+ """Return ``(authority, use_tls)`` for an endpoint string."""
117
+ if not endpoint:
118
+ raise InvalidArgumentError("endpoint is required")
119
+ use_tls = False
120
+ rest = endpoint
121
+ for scheme, tls in (
122
+ ("grpc://", False),
123
+ ("http://", False),
124
+ ("grpcs://", True),
125
+ ("https://", True),
126
+ ):
127
+ if rest.startswith(scheme):
128
+ rest = rest[len(scheme) :]
129
+ use_tls = tls
130
+ break
131
+ rest = rest.rstrip("/")
132
+ if ":" not in rest:
133
+ rest = f"{rest}:7233"
134
+ return rest, use_tls
135
+
136
+
137
+ #: Bounded automatic restarts of a crashed local daemon (design §31).
138
+ _MAX_DAEMON_RESTARTS = 3
139
+
140
+
141
+ class _Connection:
142
+ """Owns the grpc.aio channel, negotiation, and retried unary calls."""
143
+
144
+ def __init__(
145
+ self,
146
+ endpoint: str,
147
+ token: str | None = None,
148
+ *,
149
+ local_token: str | None = None,
150
+ local_generation: str | None = None,
151
+ channel_options: Iterable[tuple[str, Any]] | None = None,
152
+ compression: grpc.Compression | None = None,
153
+ ) -> None:
154
+ self._endpoint, self._tls = _parse_endpoint(endpoint)
155
+ self._token = token
156
+ self._local_token = local_token
157
+ self._local_generation = local_generation
158
+ self._attach_call: Any = None
159
+ self._attach_reader_task: asyncio.Task[Any] | None = None
160
+ self._shared_disconnect: Callable[[], Awaitable[None]] | None = None
161
+ self._channel_options = list(channel_options or ())
162
+ configured = {key for key, _ in self._channel_options}
163
+ for key in ("grpc.max_receive_message_length", "grpc.max_send_message_length"):
164
+ if key not in configured:
165
+ self._channel_options.append((key, _DEFAULT_MAX_MESSAGE_BYTES))
166
+ if local_token is not None:
167
+ configured = {key for key, _ in self._channel_options}
168
+ for key, value in (
169
+ ("grpc.keepalive_time_ms", 15000),
170
+ ("grpc.keepalive_timeout_ms", 10000),
171
+ ("grpc.keepalive_permit_without_calls", 1),
172
+ ):
173
+ if key not in configured:
174
+ self._channel_options.append((key, value))
175
+ _check_compression_channel_option_conflict(
176
+ self._channel_options, compression
177
+ )
178
+ # Channel-level default compression; survives every reconnect because
179
+ # _open_channel() re-applies it to each new aio channel. ``None`` means
180
+ # unspecified, so legacy raw channel options decide the encoding.
181
+ self._compression = compression
182
+ self._channel: grpc.aio.Channel | None = None
183
+ self._stub: queue_pb2_grpc.QueueServiceStub | None = None
184
+ self._session_minor = PROTOCOL_MINOR
185
+ self._session_token = ""
186
+ self._lock = asyncio.Lock()
187
+ self._closed = False
188
+ #: Optional client-provided reconnect hook (daemon restart + channel
189
+ #: rebuild). Invoked under ``_lock`` instead of plain ``negotiate``.
190
+ self.on_reconnect: Callable[[], Awaitable[None]] | None = None
191
+
192
+ # -- lifecycle ---------------------------------------------------------
193
+
194
+ def _open_channel(self) -> None:
195
+ if self._tls:
196
+ credentials = grpc.ssl_channel_credentials()
197
+ self._channel = grpc.aio.secure_channel(
198
+ self._endpoint,
199
+ credentials,
200
+ options=self._channel_options,
201
+ compression=self._compression,
202
+ )
203
+ else:
204
+ self._channel = grpc.aio.insecure_channel(
205
+ self._endpoint,
206
+ options=self._channel_options,
207
+ compression=self._compression,
208
+ )
209
+ self._stub = queue_pb2_grpc.QueueServiceStub(self._channel)
210
+
211
+ async def connect(self, *, check_health: bool = False) -> None:
212
+ self._open_channel()
213
+ if check_health:
214
+ await self._await_serving()
215
+ if self._local_token is not None:
216
+ await self.attach_local()
217
+ await self.negotiate()
218
+
219
+ async def attach_local(self) -> None:
220
+ """Register this SDK process as a shared-daemon member."""
221
+ if self._local_token is None or self._local_generation is None:
222
+ return
223
+ stub = self._require_stub()
224
+ try:
225
+ request = queue_pb2.AttachLocalRequest(generation=self._local_generation)
226
+ call = stub.AttachLocal(request, metadata=self._metadata())
227
+ response = await call.read()
228
+ if response is grpc.aio.EOF or response is None:
229
+ raise UnavailableError("shared daemon closed AttachLocal during handshake")
230
+ if response.generation != self._local_generation:
231
+ call.cancel()
232
+ raise UnavailableError("shared daemon generation changed during AttachLocal")
233
+ except grpc.aio.AioRpcError as exc:
234
+ raise map_rpc_error(exc) from exc
235
+ self._attach_call = call
236
+ self._attach_reader_task = asyncio.create_task(
237
+ self._consume_attach(call), name="ringo-shared-attach"
238
+ )
239
+
240
+ async def _consume_attach(self, call: Any) -> None:
241
+ try:
242
+ while True:
243
+ response = await call.read()
244
+ if response is grpc.aio.EOF:
245
+ break
246
+ if response is None:
247
+ break
248
+ if response.generation != self._local_generation:
249
+ break
250
+ except asyncio.CancelledError:
251
+ return
252
+ except grpc.aio.AioRpcError:
253
+ pass
254
+ finally:
255
+ if call is self._attach_call:
256
+ self._attach_call = None
257
+ if not self._closed and self._shared_disconnect is not None:
258
+ asyncio.ensure_future(self._shared_disconnect())
259
+
260
+ def set_shared_session(
261
+ self, token: str, generation: str, on_disconnect: Callable[[], Awaitable[None]]
262
+ ) -> None:
263
+ self._local_token = token
264
+ self._local_generation = generation
265
+ self._shared_disconnect = on_disconnect
266
+
267
+ async def _close_attach(self) -> None:
268
+ task, self._attach_reader_task = self._attach_reader_task, None
269
+ call, self._attach_call = self._attach_call, None
270
+ if call is not None:
271
+ try:
272
+ call.cancel()
273
+ except Exception:
274
+ pass
275
+ if task is not None and task is not asyncio.current_task():
276
+ task.cancel()
277
+ await asyncio.gather(task, return_exceptions=True)
278
+
279
+ async def _await_serving(self, timeout: float = 5.0) -> None:
280
+ """Wait until the health service reports SERVING (local daemon
281
+ readiness is announced on stdout before the gRPC health state may be
282
+ observable, so poll briefly instead of assuming)."""
283
+ assert self._channel is not None
284
+ health_stub = health_pb2_grpc.HealthStub(self._channel)
285
+ deadline = asyncio.get_running_loop().time() + timeout
286
+ last: Exception | None = None
287
+ while True:
288
+ try:
289
+ response = await health_stub.Check(
290
+ health_pb2.HealthCheckRequest(),
291
+ timeout=2.0,
292
+ metadata=self._metadata(),
293
+ )
294
+ if response.status == health_pb2.HealthCheckResponse.SERVING:
295
+ return
296
+ last = RingoError("daemon health status is not SERVING")
297
+ except grpc.aio.AioRpcError as exc:
298
+ last = exc
299
+ if asyncio.get_running_loop().time() >= deadline:
300
+ raise RingoError(f"daemon did not report SERVING health: {last}")
301
+ await asyncio.sleep(0.05)
302
+
303
+ async def reset(self, endpoint: str) -> None:
304
+ """Point the connection at a new endpoint (local daemon restart)."""
305
+ await self._close_attach()
306
+ self._endpoint, self._tls = _parse_endpoint(endpoint)
307
+ channel, self._channel = self._channel, None
308
+ self._stub = None
309
+ if channel is not None:
310
+ await channel.close()
311
+ self._open_channel()
312
+
313
+ async def negotiate(self) -> None:
314
+ """Perform the major=1/minor=0 version handshake (ADR 0002)."""
315
+ stub = self._require_stub()
316
+ request = queue_pb2.NegotiateRequest(
317
+ sdk_name=SDK_NAME,
318
+ sdk_version=SDK_VERSION,
319
+ protocol_major=PROTOCOL_MAJOR,
320
+ protocol_minor=PROTOCOL_MINOR,
321
+ )
322
+ try:
323
+ response = await stub.Negotiate(
324
+ request, timeout=DEFAULT_RPC_TIMEOUT, metadata=self._metadata()
325
+ )
326
+ except grpc.aio.AioRpcError as exc:
327
+ raise map_rpc_error(exc) from exc
328
+ if response.protocol_major != PROTOCOL_MAJOR:
329
+ raise IncompatibleVersionError(
330
+ f"server speaks protocol major {response.protocol_major},"
331
+ f" SDK requires major {PROTOCOL_MAJOR}; upgrade one side"
332
+ )
333
+ self._session_minor = min(response.protocol_minor, PROTOCOL_MINOR)
334
+ self._session_token = response.session_token
335
+
336
+ async def close(self) -> None:
337
+ self._closed = True
338
+ await self._close_attach()
339
+ channel, self._channel = self._channel, None
340
+ self._stub = None
341
+ if channel is not None:
342
+ await channel.close()
343
+
344
+ # -- helpers -----------------------------------------------------------
345
+
346
+ def _require_stub(self) -> queue_pb2_grpc.QueueServiceStub:
347
+ if self._stub is None:
348
+ raise RingoError("connection is closed")
349
+ return self._stub
350
+
351
+ def _metadata(self) -> list[tuple[str, str]]:
352
+ metadata: list[tuple[str, str]] = []
353
+ if self._token:
354
+ metadata.append(("authorization", f"Bearer {self._token}"))
355
+ if self._local_token:
356
+ metadata.append(("x-ringo-local-token", self._local_token))
357
+ return metadata
358
+
359
+ def request_context(self) -> queue_pb2.RequestContext:
360
+ return queue_pb2.RequestContext(
361
+ protocol_major=PROTOCOL_MAJOR,
362
+ protocol_minor=self._session_minor,
363
+ session_token=self._session_token,
364
+ )
365
+
366
+ @property
367
+ def endpoint(self) -> str:
368
+ return self._endpoint
369
+
370
+ @property
371
+ def protocol_minor(self) -> int:
372
+ """Protocol minor negotiated for the current session."""
373
+ return self._session_minor
374
+
375
+ async def call(
376
+ self,
377
+ name: str,
378
+ request: Any,
379
+ *,
380
+ timeout: float = DEFAULT_RPC_TIMEOUT,
381
+ retry_safe: bool = False,
382
+ context_attr: str | None = "context",
383
+ ) -> Any:
384
+ """Issue a single unary RPC.
385
+
386
+ Submission is never silently replayed: an ``UNAVAILABLE`` may mean the
387
+ server committed the operation (a lost ACK response, a claimed lease)
388
+ so the original error is raised after best-effort recovery (daemon
389
+ restart + fresh negotiation) prepares the connection for *future*
390
+ calls. Only genuinely read-only RPCs (task, history, and stats reads)
391
+ pass ``retry_safe=True`` and may be re-sent, each time with a freshly
392
+ negotiated ``RequestContext``.
393
+ """
394
+ rpc, sent_token = await self._begin_unary(
395
+ name, request, timeout=timeout, context_attr=context_attr
396
+ )
397
+ cause: BaseException | None = None
398
+ err: RingoError | None = None
399
+ try:
400
+ return await rpc
401
+ except grpc.aio.AioRpcError as exc:
402
+ cause = exc
403
+ err = map_rpc_error(exc)
404
+ except asyncio.CancelledError as exc:
405
+ # grpc.aio raises CancelledError when a reconnect closes its old
406
+ # channel. Preserve genuine caller cancellation, but surface this
407
+ # transport race as UNAVAILABLE.
408
+ task = asyncio.current_task()
409
+ if task is not None and task.cancelling():
410
+ raise
411
+ cause = exc
412
+ err = UnavailableError("RPC interrupted by connection recovery")
413
+
414
+ assert err is not None and cause is not None
415
+ if self._closed:
416
+ raise err from cause
417
+ if isinstance(err, IncompatibleVersionError):
418
+ # Session rejection happens in authorization before the operation
419
+ # can execute. It is therefore safe to negotiate a fresh short-lived
420
+ # token and retry once, including ordinary token expiry with no
421
+ # concurrent daemon restart.
422
+ await self._recover_best_effort()
423
+ async with self._lock:
424
+ rotated = self._session_token != sent_token
425
+ if rotated:
426
+ rpc, _ = await self._begin_unary(
427
+ name, request, timeout=timeout, context_attr=context_attr
428
+ )
429
+ try:
430
+ return await rpc
431
+ except grpc.aio.AioRpcError as retry_exc:
432
+ raise map_rpc_error(retry_exc) from retry_exc
433
+ raise err from cause
434
+ if not isinstance(err, UnavailableError):
435
+ raise err from cause
436
+
437
+ # Recovery prepares the connection for future calls. Mutations are not
438
+ # replayed because an unavailable response may hide a committed write.
439
+ await self._recover_best_effort()
440
+ if not retry_safe:
441
+ raise err from cause
442
+ rpc, _ = await self._begin_unary(
443
+ name, request, timeout=timeout, context_attr=context_attr
444
+ )
445
+ try:
446
+ return await rpc
447
+ except grpc.aio.AioRpcError as retry_exc:
448
+ raise map_rpc_error(retry_exc) from retry_exc
449
+
450
+ async def _begin_unary(
451
+ self,
452
+ name: str,
453
+ request: Any,
454
+ *,
455
+ timeout: float,
456
+ context_attr: str | None,
457
+ ) -> tuple[Any, str]:
458
+ """Snapshot a stub and fresh session context outside reconnect races."""
459
+ async with self._lock:
460
+ # Check every send, including a read retry after renegotiation.
461
+ if (
462
+ name == "GetTaskAttempts"
463
+ and self._session_minor < _ATTEMPT_HISTORY_MIN_PROTOCOL_MINOR
464
+ ):
465
+ raise IncompatibleVersionError(
466
+ "attempt history requires protocol minor >= 5 (GetTaskAttempts);"
467
+ f" negotiated minor is {self._session_minor}"
468
+ )
469
+ if context_attr is not None:
470
+ getattr(request, context_attr).CopyFrom(self.request_context())
471
+ stub = self._require_stub()
472
+ token = self._session_token
473
+ rpc = getattr(stub, name)(
474
+ request, timeout=timeout, metadata=self._metadata()
475
+ )
476
+ return rpc, token
477
+
478
+ async def _recover_best_effort(self) -> None:
479
+ """Bounded reconnect so later calls can succeed; failures here are
480
+ swallowed because the original call error is what the caller sees."""
481
+ for attempt in range(_RECONNECT_ATTEMPTS):
482
+ try:
483
+ await asyncio.sleep(
484
+ _BACKOFF_SECONDS[min(attempt, len(_BACKOFF_SECONDS) - 1)]
485
+ )
486
+ async with self._lock:
487
+ if self._closed:
488
+ return
489
+ if self.on_reconnect is not None:
490
+ await self.on_reconnect()
491
+ else:
492
+ await self.negotiate()
493
+ return
494
+ except Exception:
495
+ continue
496
+
497
+
498
+ def new_worker_id() -> str:
499
+ return f"py-{os.getpid()}-{secrets.token_hex(4)}"
500
+
501
+
502
+ # ---------------------------------------------------------------------------
503
+ # Public client
504
+
505
+
506
+ class RingoClient:
507
+ """A connected (or connectable) Ringo client.
508
+
509
+ Use as an async context manager, or call :meth:`start` / :meth:`close`
510
+ explicitly::
511
+
512
+ async with Ringo.connect("127.0.0.1:7233") as client:
513
+ queue = client.queue("crawl")
514
+ """
515
+
516
+ def __init__(
517
+ self,
518
+ *,
519
+ endpoint: str | None = None,
520
+ token: str | None = None,
521
+ daemon: "LocalDaemon | None" = None,
522
+ shared_daemon: "SharedLocalDaemon | None" = None,
523
+ channel_options: Iterable[tuple[str, Any]] | None = None,
524
+ compression: CompressionOption | None = None,
525
+ ) -> None:
526
+ self._endpoint = endpoint
527
+ self._token = token
528
+ self._channel_options = channel_options
529
+ # The value is always validated, but ``compression`` only applies to
530
+ # remote connections; local mode stays on the (uncompressed) default.
531
+ normalized = _normalize_compression(compression)
532
+ if daemon is None and shared_daemon is None:
533
+ _check_compression_channel_option_conflict(
534
+ self._channel_options, normalized
535
+ )
536
+ self._compression = normalized
537
+ self._compression_label: CompressionOption = (
538
+ compression if compression is not None else "none"
539
+ )
540
+ else:
541
+ # Local mode never enables compression; an explicit value is
542
+ # validated above and then ignored (source compatibility).
543
+ self._compression = None
544
+ self._compression_label = "none"
545
+ self._daemon = daemon
546
+ self._shared_daemon = shared_daemon
547
+ if daemon is not None and shared_daemon is not None:
548
+ raise InvalidArgumentError("daemon and shared_daemon are mutually exclusive")
549
+ # ``_conn`` is built in start(): a local client learns the daemon's
550
+ # random loopback port only after the daemon announces readiness.
551
+ self._conn: _Connection | None = None
552
+ self._queues: dict[str, TaskQueue] = {}
553
+ self._workers: set[Worker] = set()
554
+ self._started = False
555
+ self._closed = False
556
+ self._start_lock = asyncio.Lock()
557
+ self._close_lock = asyncio.Lock()
558
+ self._close_task: asyncio.Task[None] | None = None
559
+ self._daemon_restart_count = 0
560
+
561
+ # -- lifecycle ---------------------------------------------------------
562
+
563
+ async def start(self) -> "RingoClient":
564
+ async with self._start_lock:
565
+ if self._closed:
566
+ raise RingoError("client is closed")
567
+ if not self._started:
568
+ try:
569
+ for attempt in range(2 if self._shared_daemon is not None else 1):
570
+ if self._daemon is not None:
571
+ self._endpoint = await self._daemon.start()
572
+ elif self._shared_daemon is not None:
573
+ self._endpoint = await self._shared_daemon.start()
574
+ conn = _Connection(
575
+ self._endpoint or "",
576
+ self._token,
577
+ local_token=(self._shared_daemon.token if self._shared_daemon else None),
578
+ local_generation=(self._shared_daemon.generation if self._shared_daemon else None),
579
+ channel_options=self._channel_options,
580
+ compression=self._compression,
581
+ )
582
+ self._conn = conn
583
+ conn.on_reconnect = self._reconnect
584
+ if self._shared_daemon is not None:
585
+ conn.set_shared_session(
586
+ self._shared_daemon.token or "",
587
+ self._shared_daemon.generation or "",
588
+ self._shared_stream_lost,
589
+ )
590
+ try:
591
+ await conn.connect(
592
+ check_health=self._daemon is not None or self._shared_daemon is not None
593
+ )
594
+ break
595
+ except Exception:
596
+ await conn.close()
597
+ self._conn = None
598
+ if self._shared_daemon is None or attempt > 0:
599
+ raise
600
+ # Stale PID/port, wrong token or generation is never
601
+ # accepted as identity. Re-elect once, boundedly.
602
+ self._shared_daemon.invalidate_current()
603
+ except BaseException:
604
+ # Fast cleanup: never leak a half-started channel or local daemon.
605
+ failed_conn = self._conn
606
+ self._conn = None
607
+ if failed_conn is not None:
608
+ await failed_conn.close()
609
+ await self._stop_daemon()
610
+ raise
611
+ self._started = True
612
+ return self
613
+
614
+ def _require_conn(self) -> _Connection:
615
+ conn = self._conn
616
+ if conn is None:
617
+ raise RingoError(
618
+ "client is not started; use 'async with' or await client.start()"
619
+ )
620
+ return conn
621
+
622
+ async def _reconnect(self) -> None:
623
+ """Reconnect after a failed transport and negotiate a fresh session."""
624
+ conn = self._require_conn()
625
+ daemon = self._daemon
626
+ shared = self._shared_daemon
627
+ if daemon is not None and daemon.is_dead:
628
+ if self._daemon_restart_count >= _MAX_DAEMON_RESTARTS:
629
+ raise UnavailableError(
630
+ f"local daemon restart budget ({_MAX_DAEMON_RESTARTS}) exhausted"
631
+ )
632
+ self._daemon_restart_count += 1
633
+ delay = 0.3 * (2 ** (self._daemon_restart_count - 1))
634
+ await asyncio.sleep(delay)
635
+ endpoint = await daemon.start()
636
+ await conn.reset(endpoint)
637
+ elif shared is not None:
638
+ try:
639
+ endpoint = await shared.refresh()
640
+ await conn.reset(endpoint)
641
+ conn.set_shared_session(
642
+ shared.token or "", shared.generation or "", self._shared_stream_lost
643
+ )
644
+ await conn.attach_local()
645
+ await conn.negotiate()
646
+ except Exception:
647
+ shared.invalidate_current()
648
+ raise
649
+ return
650
+ await conn.negotiate()
651
+
652
+ async def _shared_stream_lost(self) -> None:
653
+ """Proactively reconnect an idle shared client after AttachLocal EOF."""
654
+ if self._closed or not self._started or self._shared_daemon is None:
655
+ return
656
+ conn = self._conn
657
+ if conn is None or conn._closed:
658
+ return
659
+ # A stream can report EOF more than once while a replacement is being
660
+ # opened. Serialize recovery through the connection's normal lock.
661
+ for attempt in range(_MAX_DAEMON_RESTARTS + 1):
662
+ if self._closed:
663
+ return
664
+ try:
665
+ await asyncio.sleep(_BACKOFF_SECONDS[min(attempt, len(_BACKOFF_SECONDS) - 1)])
666
+ async with conn._lock:
667
+ await self._reconnect()
668
+ return
669
+ except Exception:
670
+ continue
671
+
672
+ async def __aenter__(self) -> "RingoClient":
673
+ return await self.start()
674
+
675
+ async def close(self) -> None:
676
+ """Drain and stop once; caller cancellation cannot cancel cleanup."""
677
+ async with self._close_lock:
678
+ if self._close_task is None:
679
+ self._closed = True
680
+ self._close_task = asyncio.create_task(
681
+ self._close_impl(), name="ringo-client-close"
682
+ )
683
+ self._close_task.add_done_callback(_consume_task_exception)
684
+ close_task = self._close_task
685
+ await asyncio.shield(close_task)
686
+
687
+ async def _close_impl(self) -> None:
688
+ """Attempt every cleanup step and aggregate any failures."""
689
+ errors: list[Exception] = []
690
+ # 1. Drain workers (stop taking tasks, wait grace, cancel remainder).
691
+ for worker in list(self._workers):
692
+ try:
693
+ await worker.close()
694
+ except Exception as exc: # keep cleaning up the rest
695
+ errors.append(exc)
696
+ self._workers.clear()
697
+ # 2. Close the gRPC channel.
698
+ conn, self._conn = self._conn, None
699
+ if conn is not None:
700
+ try:
701
+ await conn.close()
702
+ except Exception as exc:
703
+ errors.append(exc)
704
+ # 3. Shut the local daemon down via the lifecycle pipe, then wait.
705
+ try:
706
+ await self._stop_daemon()
707
+ except Exception as exc:
708
+ errors.append(exc)
709
+ if errors:
710
+ raise ExceptionGroup("errors while closing Ringo client", errors)
711
+
712
+ async def _stop_daemon(self) -> None:
713
+ daemon, self._daemon = self._daemon, None
714
+ if daemon is not None:
715
+ await daemon.stop()
716
+
717
+ async def __aexit__(self, exc_type, exc, tb) -> None:
718
+ try:
719
+ await self.close()
720
+ except ExceptionGroup as cleanup_errors:
721
+ # A business exception takes precedence; cleanup failures must
722
+ # not mask it, but should not vanish either.
723
+ if exc is None:
724
+ raise
725
+ for error in cleanup_errors.exceptions:
726
+ exc.__cause__ = exc.__cause__ or error
727
+
728
+ @property
729
+ def is_local(self) -> bool:
730
+ return self._daemon is not None or self._shared_daemon is not None
731
+
732
+ @property
733
+ def compression(self) -> CompressionOption:
734
+ """The effective channel compression: ``"none"`` (default) or ``"gzip"``."""
735
+ return self._compression_label
736
+
737
+ # -- queues ------------------------------------------------------------
738
+
739
+ def queue(self, name: str) -> "TaskQueue":
740
+ if not name:
741
+ raise InvalidArgumentError("queue name is required")
742
+ if self._closed:
743
+ raise RingoError("client is closed")
744
+ existing = self._queues.get(name)
745
+ if existing is None:
746
+ existing = TaskQueue(self, name)
747
+ self._queues[name] = existing
748
+ return existing
749
+
750
+ def _register_worker(self, worker: "Worker") -> None:
751
+ self._workers.add(worker)
752
+
753
+ def _unregister_worker(self, worker: "Worker") -> None:
754
+ self._workers.discard(worker)
755
+
756
+
757
+ class Ringo:
758
+ """Factory for local (embedded daemon) and remote clients.
759
+
760
+ Both factories are synchronous and return a client whose
761
+ ``async with``/``await start()`` performs the actual connect and
762
+ negotiation, so ``async with Ringo.local(...)`` works directly.
763
+ """
764
+
765
+ @staticmethod
766
+ def local(
767
+ data_dir: str | os.PathLike[str],
768
+ storage: str | SQLiteStorage | PostgresStorage = "sqlite",
769
+ daemon_path: str | os.PathLike[str] | None = None,
770
+ *,
771
+ shared: bool = False,
772
+ startup_timeout: float = 15.0,
773
+ stop_timeout: float = 10.0,
774
+ ) -> RingoClient:
775
+ """Start a private daemon, or opt into a shared SQLite daemon."""
776
+ if not isinstance(shared, bool):
777
+ raise InvalidArgumentError("shared must be a boolean")
778
+ if storage == "sqlite" or isinstance(storage, SQLiteStorage):
779
+ storage_name = "sqlite"
780
+ postgres_dsn = None
781
+ elif isinstance(storage, PostgresStorage):
782
+ storage_name = "postgres"
783
+ postgres_dsn = storage.dsn
784
+ else:
785
+ raise InvalidArgumentError(f"unsupported local storage {storage!r}")
786
+ if shared and storage_name != "sqlite":
787
+ raise InvalidArgumentError("shared local mode supports SQLite only")
788
+ # Imported here so remote-only users never pay for subprocess imports.
789
+ if shared:
790
+ from .shared import SharedLocalDaemon
791
+
792
+ shared_daemon = SharedLocalDaemon(
793
+ data_dir, daemon_path=daemon_path, startup_timeout=startup_timeout
794
+ )
795
+ return RingoClient(shared_daemon=shared_daemon)
796
+ from .daemon import LocalDaemon
797
+
798
+ daemon = LocalDaemon(
799
+ data_dir,
800
+ storage=storage_name,
801
+ postgres_dsn=postgres_dsn,
802
+ daemon_path=daemon_path,
803
+ startup_timeout=startup_timeout,
804
+ stop_timeout=stop_timeout,
805
+ )
806
+ return RingoClient(daemon=daemon)
807
+
808
+ @staticmethod
809
+ def connect(
810
+ endpoint: str,
811
+ token: str | None = None,
812
+ *,
813
+ channel_options: Iterable[tuple[str, Any]] | None = None,
814
+ compression: CompressionOption | None = None,
815
+ ) -> RingoClient:
816
+ """Create a client for a remote (or already-running) endpoint.
817
+
818
+ ``compression`` is an optional whole-channel setting: ``"none"``
819
+ (the default) or ``"gzip"``. ``None`` means "not specified", leaving
820
+ the gRPC channel default (and any raw ``channel_options`` compression
821
+ entry) in charge. An explicit value combined with the raw channel
822
+ option ``grpc.default_compression_algorithm`` raises
823
+ :class:`InvalidArgumentError`; any other value does so as well,
824
+ immediately and before any network activity.
825
+ """
826
+ return RingoClient(
827
+ endpoint=endpoint,
828
+ token=token,
829
+ channel_options=channel_options,
830
+ compression=compression,
831
+ )
832
+
833
+
834
+ # ---------------------------------------------------------------------------
835
+ # TaskQueue
836
+
837
+
838
+ Handler = Callable[[Task, Any], Coroutine[Any, Any, None]]
839
+
840
+
841
+ class TaskQueue:
842
+ """A client bound to one logical queue (``queue + uniqueKey`` dedupe)."""
843
+
844
+ def __init__(self, client: RingoClient, name: str) -> None:
845
+ self._client = client
846
+ self.name = name
847
+
848
+ # -- enqueue ------------------------------------------------------------
849
+
850
+ async def enqueue(
851
+ self,
852
+ spec: TaskSpec | None = None,
853
+ /,
854
+ **kwargs: Any,
855
+ ) -> EnqueueResult:
856
+ """Enqueue one task; duplicates return ``created=False`` (no raise)."""
857
+ if spec is None:
858
+ spec = TaskSpec(**kwargs)
859
+ request = queue_pb2.EnqueueRequest(queue=self.name, task=spec.to_proto())
860
+ request.context.CopyFrom(self._client._require_conn().request_context())
861
+ response = await self._client._require_conn().call("Enqueue", request)
862
+ return EnqueueResult.from_proto(response.result)
863
+
864
+ async def enqueue_many(self, specs: Iterable[TaskSpec]) -> BatchResult:
865
+ """Enqueue any number of tasks, auto-batching under the server limit.
866
+
867
+ The server enforces a hard limit of ``_MAX_ENQUEUE_BATCH`` (1000)
868
+ tasks per gRPC call. This method transparently splits larger
869
+ iterables into sequential sub-batches, preserving creation order,
870
+ and aggregates the per-item results into a single :class:`BatchResult`.
871
+
872
+ Fail-fast: if any sub-batch fails, the exception propagates and no
873
+ partial result is returned.
874
+ """
875
+ specs = list(specs)
876
+ all_results: list[EnqueueResult] = []
877
+ total_created = 0
878
+ total_duplicates = 0
879
+
880
+ for start in range(0, len(specs), _MAX_ENQUEUE_BATCH):
881
+ chunk = specs[start : start + _MAX_ENQUEUE_BATCH]
882
+ request = queue_pb2.EnqueueManyRequest(
883
+ queue=self.name, tasks=[s.to_proto() for s in chunk]
884
+ )
885
+ request.context.CopyFrom(self._client._require_conn().request_context())
886
+ response = await self._client._require_conn().call("EnqueueMany", request)
887
+ all_results.extend(EnqueueResult.from_proto(r) for r in response.results)
888
+ total_created += response.created_count
889
+ total_duplicates += response.duplicate_count
890
+
891
+ return BatchResult(
892
+ results=all_results,
893
+ created_count=total_created,
894
+ duplicate_count=total_duplicates,
895
+ )
896
+
897
+ # -- queries ------------------------------------------------------------
898
+
899
+ async def get_by_id(self, task_id: str) -> Task:
900
+ if not task_id:
901
+ raise InvalidArgumentError("task_id is required")
902
+ request = queue_pb2.GetTaskRequest(queue=self.name, task_id=task_id)
903
+ request.context.CopyFrom(self._client._require_conn().request_context())
904
+ response = await self._client._require_conn().call(
905
+ "GetTask", request, retry_safe=True
906
+ )
907
+ return Task.from_proto(response.task)
908
+
909
+ async def get_by_unique_key(self, unique_key: str) -> Task:
910
+ if not unique_key:
911
+ raise InvalidArgumentError("unique_key is required")
912
+ request = queue_pb2.GetTaskRequest(queue=self.name, unique_key=unique_key)
913
+ request.context.CopyFrom(self._client._require_conn().request_context())
914
+ response = await self._client._require_conn().call(
915
+ "GetTask", request, retry_safe=True
916
+ )
917
+ return Task.from_proto(response.task)
918
+
919
+ async def get_attempts(
920
+ self,
921
+ unique_key: str,
922
+ *,
923
+ limit: int = _ATTEMPT_PAGE_DEFAULT_LIMIT,
924
+ cursor: str | None = None,
925
+ ) -> AttemptPage:
926
+ """Read oldest-first history by business key, not the bounded Task preview.
927
+
928
+ Cursors retain the first page's history high-water mark. Pages can be
929
+ shorter than ``limit`` due to the byte budget; follow ``next_cursor``
930
+ until None. Restart without a cursor to see later appended attempts.
931
+ """
932
+ return await self._get_attempts("unique_key", unique_key, limit, cursor)
933
+
934
+ async def get_attempts_by_id(
935
+ self,
936
+ task_id: str,
937
+ *,
938
+ limit: int = _ATTEMPT_PAGE_DEFAULT_LIMIT,
939
+ cursor: str | None = None,
940
+ ) -> AttemptPage:
941
+ """Read a bounded durable history page by explicit server task ID."""
942
+ return await self._get_attempts("task_id", task_id, limit, cursor)
943
+
944
+ async def _get_attempts(
945
+ self, selector: str, key: str, limit: int, cursor: str | None
946
+ ) -> AttemptPage:
947
+ if not isinstance(key, str) or not key:
948
+ raise InvalidArgumentError(f"{selector} is required")
949
+ if (
950
+ isinstance(limit, bool)
951
+ or not isinstance(limit, int)
952
+ or not 1 <= limit <= _ATTEMPT_PAGE_MAX_LIMIT
953
+ ):
954
+ raise InvalidArgumentError("limit must be an integer between 1 and 100")
955
+ if cursor is not None:
956
+ if not isinstance(cursor, str) or not cursor:
957
+ raise InvalidArgumentError("cursor must be a nonempty string or None")
958
+ try:
959
+ cursor_bytes = len(cursor.encode("utf-8"))
960
+ except UnicodeEncodeError as exc:
961
+ raise InvalidArgumentError("cursor must be valid UTF-8") from exc
962
+ if cursor_bytes > _MAX_ATTEMPT_CURSOR_BYTES:
963
+ raise InvalidArgumentError("cursor must be at most 1024 UTF-8 bytes")
964
+ conn = self._client._require_conn()
965
+ if conn.protocol_minor < _ATTEMPT_HISTORY_MIN_PROTOCOL_MINOR:
966
+ raise IncompatibleVersionError(
967
+ "attempt history requires protocol minor >= 5 (GetTaskAttempts);"
968
+ f" negotiated minor is {conn.protocol_minor}"
969
+ )
970
+ request = queue_pb2.GetTaskAttemptsRequest(
971
+ queue=self.name, limit=limit, cursor=cursor or ""
972
+ )
973
+ setattr(request, selector, key)
974
+ response = await conn.call("GetTaskAttempts", request, retry_safe=True)
975
+ return AttemptPage.from_proto(response)
976
+
977
+ async def stats(self, task_types: Iterable[str] | None = None) -> QueueStats:
978
+ """Read a point-in-time per-status count snapshot for this queue.
979
+
980
+ ``task_types`` of ``None`` or empty counts every task type; otherwise
981
+ only tasks matching any of the given types are included. All six
982
+ canonical statuses are always present (zero included); ``active`` on
983
+ the result is ``SCHEDULED + READY + LEASED``. This is a per-queue,
984
+ per-instance observation: the SDK never fabricates a cross-instance
985
+ aggregate or a time-window summary.
986
+
987
+ The read-only ``GetQueueStats`` RPC is additive protocol 1.4: when the
988
+ negotiated session minor is lower, this raises
989
+ :class:`IncompatibleVersionError` before any network activity.
990
+ ``NOT_FOUND`` (unknown queue) and ``UNAVAILABLE`` propagate unchanged.
991
+ More than 100 ``task_types`` is an :class:`InvalidArgumentError` raised
992
+ before any RPC; at most 101 elements are consumed from the iterable.
993
+ """
994
+ selected: list[str] | None = None
995
+ if task_types is not None:
996
+ # islice caps consumption so an unbounded iterable stops at 101.
997
+ selected = list(islice(task_types, _STATS_MAX_TASK_TYPES + 1))
998
+ if len(selected) > _STATS_MAX_TASK_TYPES:
999
+ raise InvalidArgumentError(
1000
+ f"task_types must contain at most {_STATS_MAX_TASK_TYPES}"
1001
+ f" values; got more than {_STATS_MAX_TASK_TYPES}"
1002
+ )
1003
+ conn = self._client._require_conn()
1004
+ if conn.protocol_minor < _STATS_MIN_PROTOCOL_MINOR:
1005
+ raise IncompatibleVersionError(
1006
+ f"queue stats requires protocol minor >= {_STATS_MIN_PROTOCOL_MINOR}"
1007
+ f" (GetQueueStats); negotiated minor is {conn.protocol_minor}"
1008
+ )
1009
+ request = queue_pb2.GetQueueStatsRequest(
1010
+ queue=self.name, task_types=selected or []
1011
+ )
1012
+ request.context.CopyFrom(conn.request_context())
1013
+ response = await conn.call("GetQueueStats", request, retry_safe=True)
1014
+ return QueueStats.from_proto(response)
1015
+
1016
+ async def retry(self, unique_key: str) -> Task:
1017
+ """Reactivate a SUCCEEDED/DEAD task by its business unique key.
1018
+
1019
+ This always resolves through the proto ``unique_key`` field; use
1020
+ :meth:`retry_by_id` for the task-ID form. No string is ever guessed.
1021
+ """
1022
+ if not unique_key:
1023
+ raise InvalidArgumentError("unique_key is required")
1024
+ request = queue_pb2.RetryTaskRequest(queue=self.name)
1025
+ request.unique_key = unique_key
1026
+ request.context.CopyFrom(self._client._require_conn().request_context())
1027
+ response = await self._client._require_conn().call("RetryTask", request)
1028
+ return Task.from_proto(response.task)
1029
+
1030
+ async def retry_by_id(self, task_id: str) -> Task:
1031
+ """Reactivate a SUCCEEDED/DEAD task by its server task ID."""
1032
+ if not task_id:
1033
+ raise InvalidArgumentError("task_id is required")
1034
+ request = queue_pb2.RetryTaskRequest(queue=self.name)
1035
+ request.task_id = task_id
1036
+ request.context.CopyFrom(self._client._require_conn().request_context())
1037
+ response = await self._client._require_conn().call("RetryTask", request)
1038
+ return Task.from_proto(response.task)
1039
+
1040
+ # -- manual claim ---------------------------------------------------------
1041
+
1042
+ async def claim(
1043
+ self,
1044
+ limit: int = 1,
1045
+ *,
1046
+ task_types: Iterable[str] | None = None,
1047
+ wait_timeout: timedelta | None = None,
1048
+ lease_duration: timedelta | None = None,
1049
+ worker_id: str | None = None,
1050
+ ) -> list["LeasedTask"]:
1051
+ """Manually claim tasks (advanced API; not the default worker path)."""
1052
+ if limit < 1 or limit > _MAX_ENQUEUE_BATCH:
1053
+ raise InvalidArgumentError(
1054
+ f"limit must be between 1 and {_MAX_ENQUEUE_BATCH}"
1055
+ )
1056
+ if wait_timeout is not None and wait_timeout < timedelta(0):
1057
+ raise InvalidArgumentError("wait_timeout must not be negative")
1058
+ if lease_duration is not None and lease_duration < timedelta(0):
1059
+ raise InvalidArgumentError("lease_duration must not be negative")
1060
+ request = queue_pb2.ClaimRequest(
1061
+ queue=self.name,
1062
+ worker_id=worker_id or new_worker_id(),
1063
+ task_types=list(task_types or []),
1064
+ limit=limit,
1065
+ )
1066
+ if wait_timeout is not None:
1067
+ request.wait_timeout.CopyFrom(_pb.td_to_dur(wait_timeout))
1068
+ if lease_duration is not None:
1069
+ request.lease_duration.CopyFrom(_pb.td_to_dur(lease_duration))
1070
+ request.context.CopyFrom(self._client._require_conn().request_context())
1071
+ response = await self._client._require_conn().call("Claim", request)
1072
+ return [LeasedTask(self, t) for t in response.tasks]
1073
+
1074
+ # -- worker ---------------------------------------------------------------
1075
+
1076
+ def worker(
1077
+ self,
1078
+ handler: Handler,
1079
+ *,
1080
+ task_types: Iterable[str] | None = None,
1081
+ concurrency: int | ResourceConcurrency | AdaptiveConcurrency = 1,
1082
+ weight: int = 1,
1083
+ lease_duration: timedelta = timedelta(seconds=60),
1084
+ grace_period: timedelta = timedelta(seconds=30),
1085
+ on_idle: Callable[[], Any] | None = None,
1086
+ ) -> "Worker":
1087
+ """Create a credit-controlled worker over the frozen Work stream.
1088
+
1089
+ ``on_idle`` is an optional zero-argument callable (sync or awaitable)
1090
+ for bounded batch runs: once the worker has registered and sent its
1091
+ initial credit, the SDK polls queue activity every 250 ms via the
1092
+ read-only ``GetQueueActivity`` RPC scoped to this worker's
1093
+ ``task_types``; after the activity count has been continuously zero
1094
+ for at least 1 s the callback fires exactly once, inline on the
1095
+ ``run()`` coroutine. A callback that returns normally leaves the
1096
+ worker running (and never fires again); any exception it raises —
1097
+ including ``SystemExit`` — propagates verbatim out of
1098
+ ``await worker.run()`` without becoming a task NACK. Workers without
1099
+ ``on_idle`` never send activity RPCs and stay compatible with
1100
+ protocol minor 0 servers; ``on_idle`` requires minor >= 1 and raises
1101
+ :class:`IncompatibleVersionError` from ``run()`` otherwise. Never
1102
+ call ``worker.close()`` from inside the callback (it would wait on
1103
+ the very coroutine executing it).
1104
+ """
1105
+ from .worker import Worker
1106
+
1107
+ if not callable(handler):
1108
+ raise InvalidArgumentError("handler must be callable")
1109
+ if on_idle is not None and not callable(on_idle):
1110
+ raise InvalidArgumentError("on_idle must be callable")
1111
+ if isinstance(concurrency, bool) or not isinstance(
1112
+ concurrency, (int, ResourceConcurrency, AdaptiveConcurrency)
1113
+ ):
1114
+ raise InvalidArgumentError(
1115
+ "concurrency must be an integer, ResourceConcurrency, or AdaptiveConcurrency"
1116
+ )
1117
+ if isinstance(concurrency, int) and (concurrency < 1 or concurrency > 2**31 - 1):
1118
+ raise InvalidArgumentError(
1119
+ "initial fixed concurrency must be between 1 and 2**31-1"
1120
+ )
1121
+ if isinstance(weight, bool) or not isinstance(weight, int) or weight < 1:
1122
+ raise InvalidArgumentError("weight must be a positive integer")
1123
+ if lease_duration < timedelta(0):
1124
+ raise InvalidArgumentError("lease_duration must not be negative")
1125
+ if grace_period < timedelta(0):
1126
+ raise InvalidArgumentError("grace_period must not be negative")
1127
+ worker = Worker(
1128
+ self,
1129
+ handler,
1130
+ task_types=list(task_types or []),
1131
+ concurrency=concurrency,
1132
+ lease_duration=lease_duration,
1133
+ grace_period=grace_period,
1134
+ on_idle=on_idle,
1135
+ weight=weight,
1136
+ )
1137
+ self._client._register_worker(worker)
1138
+ return worker
1139
+
1140
+
1141
+ # ---------------------------------------------------------------------------
1142
+ # Manual lease
1143
+
1144
+
1145
+ class LeasedTask:
1146
+ """A task claimed via :meth:`TaskQueue.claim` with its current lease."""
1147
+
1148
+ def __init__(self, queue: TaskQueue, msg: queue_pb2.LeasedTask) -> None:
1149
+ self._queue = queue
1150
+ self.task = Task.from_proto(msg.task)
1151
+ self.attempt: int = msg.lease.attempt
1152
+ self.worker_id: str = msg.lease.worker_id
1153
+ self._token: str = msg.lease.lease_token
1154
+ self._lease_until = _pb.ts_to_dt(msg.lease.until)
1155
+
1156
+ @property
1157
+ def lease_until(self):
1158
+ return self._lease_until
1159
+
1160
+ def _ref(self) -> queue_pb2.LeaseRef:
1161
+ return queue_pb2.LeaseRef(
1162
+ task_id=self.task.id,
1163
+ attempt=self.attempt,
1164
+ lease_token=self._token,
1165
+ )
1166
+
1167
+ async def _terminal(self, name: str, request: Any) -> None:
1168
+ request.lease.CopyFrom(self._ref())
1169
+ conn = self._queue._client._require_conn()
1170
+ request.context.CopyFrom(conn.request_context())
1171
+ await conn.call(name, request)
1172
+
1173
+ async def ack(self) -> None:
1174
+ """Acknowledge success for the current attempt/lease token."""
1175
+ await self._terminal("Ack", queue_pb2.AckRequest())
1176
+
1177
+ async def nack(self, error: str = "", *, delay: timedelta | None = None) -> None:
1178
+ """Negative-acknowledge; the task is retried per policy."""
1179
+ request = queue_pb2.NackRequest(error=error)
1180
+ if delay is not None:
1181
+ request.retry_delay_override.CopyFrom(_pb.td_to_dur(delay))
1182
+ await self._terminal("Nack", request)
1183
+
1184
+ async def reject(self, error: str = "") -> None:
1185
+ """Permanently fail the task (DEAD) for the current attempt."""
1186
+ await self._terminal("Reject", queue_pb2.RejectRequest(error=error))
1187
+
1188
+ async def extend_lease(self, extension: timedelta) -> None:
1189
+ """Extend the lease to ``server_now + extension`` (absolute reset)."""
1190
+ if extension <= timedelta(0):
1191
+ raise InvalidArgumentError("extension must be positive")
1192
+ request = queue_pb2.ExtendLeaseRequest()
1193
+ request.extension.CopyFrom(_pb.td_to_dur(extension))
1194
+ request.lease.CopyFrom(self._ref())
1195
+ conn = self._queue._client._require_conn()
1196
+ request.context.CopyFrom(conn.request_context())
1197
+ response = await conn.call("ExtendLease", request)
1198
+ self._lease_until = _pb.ts_to_dt(response.lease.until)