flpit 0.2.1.dev5__py3-none-any.whl → 0.2.1.dev7__py3-none-any.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.
flpit/task.py ADDED
@@ -0,0 +1,1610 @@
1
+ """Task: a small asyncio-free Task/awaitable runtime.
2
+
3
+ This is a deliberately small runtime, not an asyncio compatibility layer.
4
+
5
+ Public API
6
+ ----------
7
+ Task.configure(max_workers=8, max_pending_background=128)
8
+ Task.run(async_function, *args, token=token)
9
+ Task.run(blocking_function, *args, token=token, pass_token=True)
10
+ await Task.delay(0.5, token=token)
11
+ await IO.wait_readable(sock, token=token)
12
+ await IO.wait_writable(sock, token=token)
13
+
14
+ cts = CancellationTokenSource()
15
+ linked = CancellationTokenSource.create_linked_token_source(cts.token, other.token)
16
+ cts.cancel()
17
+
18
+ There is no application-managed Runtime object and no asyncio dependency. The
19
+ runtime is process-wide and starts lazily. Call Task.shutdown() for an explicit,
20
+ observable shutdown; it can be started again by subsequent Task.run calls.
21
+
22
+ Cancellation is cooperative for coroutines and worker functions. A worker can
23
+ receive the same token with pass_token=True and should check it between units
24
+ of work. Python cannot safely force-stop an arbitrary running thread.
25
+ """
26
+ from __future__ import annotations
27
+
28
+ import atexit
29
+ import heapq
30
+ import inspect
31
+ import logging
32
+ import os
33
+ import queue
34
+ import selectors
35
+ import socket
36
+ import threading
37
+ import time
38
+ from collections import deque
39
+ from concurrent.futures import Future, ThreadPoolExecutor
40
+ from dataclasses import dataclass
41
+ from itertools import count
42
+ from typing import Any, Callable, Coroutine, Generic, TypeVar, cast, Iterable
43
+ import functools
44
+
45
+ T = TypeVar("T")
46
+ _LOG = logging.getLogger("tinytask")
47
+ _MISSING = object()
48
+ _WAKE = object()
49
+
50
+
51
+ class OperationCanceledError(BaseException):
52
+ """Raised when cancellation is delivered to an operation or coroutine."""
53
+
54
+
55
+ class InvalidStateError(RuntimeError):
56
+ """An operation was requested before a Task reached a valid state."""
57
+
58
+
59
+ class BackgroundQueueFull(RuntimeError):
60
+ """Raised when the bounded background-work queue has no free capacity."""
61
+
62
+
63
+ class RuntimeShutdownError(RuntimeError):
64
+ """Raised when new work is submitted while the singleton runtime is stopping."""
65
+
66
+
67
+ class AggregateError(Exception):
68
+ """Raised by Task.when_all() when multiple child tasks fail.
69
+
70
+ The original exceptions are available in the ``exceptions`` tuple. A single
71
+ child failure is re-raised directly to keep the common case ergonomic.
72
+ """
73
+
74
+ def __init__(
75
+ self,
76
+ exceptions: Iterable[Exception],
77
+ message: str = "One or more errors occurred.",
78
+ ) -> None:
79
+ self.exceptions = tuple(exceptions)
80
+ super().__init__(message)
81
+
82
+
83
+ @dataclass(frozen=True)
84
+ class ShutdownReport:
85
+ """Outcome of a runtime shutdown request."""
86
+
87
+ runtime_stopped: bool
88
+ tasks_cancelled: int
89
+ workers_still_running: int
90
+ timed_out: bool
91
+
92
+
93
+ class _CancellationState:
94
+ """Thread-safe cancellation state shared by tokens and their source."""
95
+
96
+ def __init__(self) -> None:
97
+ self.lock = threading.RLock()
98
+ self.cancelled = False
99
+ self.callbacks: dict[int, Callable[[], Any]] = {}
100
+ self.next_id = count(1)
101
+
102
+ def register(self, callback: Callable[[], Any]) -> "CancellationTokenRegistration":
103
+ with self.lock:
104
+ if not self.cancelled:
105
+ key = next(self.next_id)
106
+ self.callbacks[key] = callback
107
+
108
+ return CancellationTokenRegistration(self, key)
109
+
110
+ exception = _invoke_cancellation_callback(callback)
111
+
112
+ if exception is not None:
113
+ raise exception
114
+
115
+ return CancellationTokenRegistration(None, None)
116
+
117
+ def unregister(self, key: int) -> None:
118
+ with self.lock:
119
+ self.callbacks.pop(key, None)
120
+
121
+ def cancel(self) -> bool:
122
+ with self.lock:
123
+
124
+ if self.cancelled:
125
+ return False
126
+
127
+ self.cancelled = True
128
+ callbacks = list(reversed(self.callbacks.values()))
129
+ self.callbacks.clear()
130
+
131
+ exceptions: list[Exception] = []
132
+
133
+ for callback in callbacks:
134
+ exception = _invoke_cancellation_callback(callback)
135
+ if exception is not None:
136
+ exceptions.append(exception)
137
+
138
+ if exceptions:
139
+ raise AggregateError(exceptions)
140
+
141
+ return True
142
+
143
+
144
+ def _invoke_cancellation_callback(callback: Callable[[], Any]) -> Exception | None:
145
+ try:
146
+ callback()
147
+ except Exception as exc:
148
+ return exc
149
+
150
+ return None
151
+
152
+
153
+ class CancellationTokenRegistration:
154
+ """Disposable registration returned by CancellationToken.register()."""
155
+
156
+ __slots__ = ("_state", "_key", "_disposed", "_lock")
157
+
158
+ def __init__(self, state: _CancellationState | None, key: int | None) -> None:
159
+ self._state = state
160
+ self._key = key
161
+ self._disposed = False
162
+ self._lock = threading.Lock()
163
+
164
+ def dispose(self) -> None:
165
+ with self._lock:
166
+ if self._disposed:
167
+ return
168
+ self._disposed = True
169
+ state, key = self._state, self._key
170
+ self._state = None
171
+ self._key = None
172
+ if state is not None and key is not None:
173
+ state.unregister(key)
174
+
175
+ close = dispose
176
+
177
+ def __enter__(self) -> "CancellationTokenRegistration":
178
+ return self
179
+
180
+ def __exit__(self, *_: Any) -> None:
181
+ self.dispose()
182
+
183
+
184
+ class CancellationToken:
185
+ """Read-only cancellation handle. Create cancellable tokens via a source."""
186
+
187
+ __slots__ = ("_state",)
188
+
189
+ def __init__(self, state: _CancellationState | None = None) -> None:
190
+ # A public no-argument token is intentionally non-cancellable, matching
191
+ # the role of CancellationToken.none(). Sources create tokens with state.
192
+ self._state = state
193
+
194
+ @classmethod
195
+ def none(cls) -> "CancellationToken":
196
+ return _NONE_TOKEN
197
+
198
+ @property
199
+ def can_be_cancelled(self) -> bool:
200
+ return self._state is not None
201
+
202
+ @property
203
+ def is_cancellation_requested(self) -> bool:
204
+ state = self._state
205
+ if state is None:
206
+ return False
207
+ with state.lock:
208
+ return state.cancelled
209
+
210
+ def throw_if_cancellation_requested(self) -> None:
211
+ if self.is_cancellation_requested:
212
+ raise OperationCanceledError("Cancellation was requested")
213
+
214
+ def register(self, callback: Callable[[], Any]) -> CancellationTokenRegistration:
215
+ if not callable(callback):
216
+ raise TypeError("callback must be callable")
217
+ state = self._state
218
+ if state is None:
219
+ return CancellationTokenRegistration(None, None)
220
+ return state.register(callback)
221
+
222
+ def __repr__(self) -> str:
223
+ return (
224
+ "CancellationToken(can_be_cancelled=False)"
225
+ if self._state is None
226
+ else f"CancellationToken(cancelled={self.is_cancellation_requested})"
227
+ )
228
+
229
+
230
+ _NONE_TOKEN = CancellationToken()
231
+
232
+
233
+ class CancellationTokenSource:
234
+ """Owns cancellation state and provides a token consumers can observe."""
235
+
236
+ def __init__(self) -> None:
237
+ self._state = _CancellationState()
238
+ self._token = CancellationToken(self._state)
239
+ self._lock = threading.RLock()
240
+ self._disposed = False
241
+ self._timer: threading.Timer | None = None
242
+ self._linked_registrations: list[CancellationTokenRegistration] = []
243
+
244
+ @property
245
+ def token(self) -> CancellationToken:
246
+ return self._token
247
+
248
+ @property
249
+ def is_cancellation_requested(self) -> bool:
250
+ return self._token.is_cancellation_requested
251
+
252
+ def cancel(self) -> bool:
253
+ with self._lock:
254
+ self._throw_if_disposed()
255
+ timer = self._timer
256
+ self._timer = None
257
+ links = list(self._linked_registrations)
258
+ self._linked_registrations.clear()
259
+ if timer is not None:
260
+ timer.cancel()
261
+ cancelled = self._state.cancel()
262
+ # A linked source no longer needs to hold registrations after its own
263
+ # cancellation has fired. This is safe even when called from a link callback.
264
+ for registration in links:
265
+ registration.dispose()
266
+ return cancelled
267
+
268
+ def cancel_after(self, seconds: float | None) -> None:
269
+ """Request cancellation after seconds; None or -1 disables the timer."""
270
+ with self._lock:
271
+ self._throw_if_disposed()
272
+ old_timer = self._timer
273
+ self._timer = None
274
+ if seconds is not None and seconds != -1:
275
+ if seconds < 0:
276
+ raise ValueError("seconds must be >= 0, -1, or None")
277
+ timer = threading.Timer(seconds, self._cancel_from_timer)
278
+ timer.daemon = True
279
+ self._timer = timer
280
+ else:
281
+ timer = None
282
+ if old_timer is not None:
283
+ old_timer.cancel()
284
+ if timer is not None:
285
+ timer.start()
286
+
287
+ def _cancel_from_timer(self) -> None:
288
+ try:
289
+ self.cancel()
290
+ except RuntimeError:
291
+ # The source was disposed concurrently with its timer firing.
292
+ pass
293
+
294
+ @classmethod
295
+ def create_linked_token_source(
296
+ cls, *tokens: CancellationToken
297
+ ) -> "CancellationTokenSource":
298
+ if any(not isinstance(token, CancellationToken) for token in tokens):
299
+ raise TypeError("all linked values must be CancellationToken instances")
300
+ linked = cls()
301
+ for token in tokens:
302
+ registration = token.register(linked._cancel_from_link)
303
+ linked._linked_registrations.append(registration)
304
+ # If the source token was already cancelled, the callback executes
305
+ # synchronously during register() and the linked source is now cancelled.
306
+ if linked.is_cancellation_requested:
307
+ break
308
+ if linked.is_cancellation_requested:
309
+ linked._dispose_link_registrations()
310
+ return linked
311
+
312
+ def _cancel_from_link(self) -> None:
313
+ try:
314
+ self.cancel()
315
+ except RuntimeError:
316
+ pass
317
+
318
+ def _dispose_link_registrations(self) -> None:
319
+ with self._lock:
320
+ links = list(self._linked_registrations)
321
+ self._linked_registrations.clear()
322
+ for registration in links:
323
+ registration.dispose()
324
+
325
+ def dispose(self) -> None:
326
+ with self._lock:
327
+ if self._disposed:
328
+ return
329
+ self._disposed = True
330
+ timer = self._timer
331
+ self._timer = None
332
+ links = list(self._linked_registrations)
333
+ self._linked_registrations.clear()
334
+ if timer is not None:
335
+ timer.cancel()
336
+ for registration in links:
337
+ registration.dispose()
338
+
339
+ close = dispose
340
+
341
+ def _throw_if_disposed(self) -> None:
342
+ if self._disposed:
343
+ raise RuntimeError("CancellationTokenSource has been disposed")
344
+
345
+ def __enter__(self) -> "CancellationTokenSource":
346
+ self._throw_if_disposed()
347
+ return self
348
+
349
+ def __exit__(self, *_: Any) -> None:
350
+ self.dispose()
351
+
352
+
353
+ class _Operation:
354
+ """Instruction yielded by an awaitable and consumed by the scheduler."""
355
+
356
+ def __init__(self, runtime: "_Runtime", token: CancellationToken | None = None) -> None:
357
+ self.runtime = runtime
358
+ self.token = token or CancellationToken.none()
359
+ self._token_registration: CancellationTokenRegistration | None = None
360
+
361
+ def __await__(self):
362
+ result = yield self
363
+ return result
364
+
365
+
366
+ class _Sleep(_Operation):
367
+ def __init__(self, runtime: "_Runtime", delay: float) -> None:
368
+ super().__init__(runtime)
369
+ self.delay = delay
370
+ self.deadline = time.monotonic() + delay
371
+
372
+
373
+ class _YieldNow(_Operation):
374
+ pass
375
+
376
+
377
+ class _WaitIO(_Operation):
378
+ def __init__(
379
+ self, runtime: "_Runtime", sock: socket.socket, mode: str,
380
+ token: CancellationToken | None = None,
381
+ ) -> None:
382
+ super().__init__(runtime, token)
383
+ self.sock = sock
384
+ self.mode = mode
385
+
386
+
387
+ class _WaitTask(_Operation):
388
+ def __init__(self, runtime: "_Runtime", task: "Task[Any]") -> None:
389
+ super().__init__(runtime)
390
+ self.task = task
391
+
392
+
393
+ class _WaitMany(_Operation):
394
+ """Wait for all or any of a set of tasks to complete."""
395
+
396
+ def __init__(
397
+ self, runtime: "_Runtime", tasks: tuple["Task[Any]", ...], mode: str
398
+ ) -> None:
399
+ super().__init__(runtime)
400
+ if mode not in ("all", "any"):
401
+ raise ValueError("mode must be 'all' or 'any'")
402
+ self.tasks = tasks
403
+ self.mode = mode
404
+
405
+
406
+ class _BackgroundCall(_Operation):
407
+ def __init__(
408
+ self,
409
+ runtime: "_Runtime",
410
+ function: Callable[..., Any],
411
+ args: tuple[Any, ...],
412
+ kwargs: dict[str, Any],
413
+ ) -> None:
414
+ super().__init__(runtime)
415
+ self.function = function
416
+ self.args = args
417
+ self.kwargs = kwargs
418
+ self.future: Future[Any] | None = None
419
+ self._counted_active = False
420
+
421
+
422
+ async def _await_any(awaitable: Any) -> Any:
423
+ return await awaitable
424
+
425
+
426
+ async def _invoke_async(
427
+ function: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]
428
+ ) -> Any:
429
+ return await function(*args, **kwargs)
430
+
431
+
432
+ async def _invoke_background(
433
+ runtime: "_Runtime",
434
+ function: Callable[..., Any],
435
+ args: tuple[Any, ...],
436
+ kwargs: dict[str, Any],
437
+ ) -> Any:
438
+ result = await _BackgroundCall(runtime, function, args, kwargs)
439
+ # Support sync wrappers that return an awaitable.
440
+ if inspect.isawaitable(result):
441
+ return await result
442
+ return result
443
+
444
+
445
+ async def _delay_coroutine(runtime: "_Runtime", delay: float) -> None:
446
+ await _Sleep(runtime, delay)
447
+
448
+
449
+ class Task(Generic[T]):
450
+ """Result-bearing task managed by TinyTask's process-wide scheduler."""
451
+
452
+ PENDING = "pending"
453
+ SUCCEEDED = "succeeded"
454
+ FAILED = "failed"
455
+ CANCELLED = "cancelled"
456
+
457
+ def __init__(
458
+ self,
459
+ runtime: "_Runtime",
460
+ *,
461
+ name: str | None = None,
462
+ token: CancellationToken | None = None,
463
+ ) -> None:
464
+ self._runtime = runtime
465
+ self.name = name or f"Task-{id(self):x}"
466
+ self._state = self.PENDING
467
+ self._result: Any = None
468
+ self._exception: BaseException | None = None
469
+ self._coroutine: Coroutine[Any, Any, Any] | None = None
470
+ self._awaiters: set[Task[Any]] = set()
471
+ self._waiting_on: _Operation | None = None
472
+ self._queued = False
473
+ self._running = False
474
+ self._next_value: Any = None
475
+ self._next_exception: BaseException | None = None
476
+ self._cancel_requested = False
477
+ self._completed = threading.Event()
478
+ self._token = token or CancellationToken.none()
479
+ self._token_registration: CancellationTokenRegistration | None = None
480
+ self._pass_token = False
481
+
482
+ @classmethod
483
+ def configure(
484
+ cls,
485
+ *,
486
+ max_workers: int | None = None,
487
+ max_pending_background: int | None = None,
488
+ ) -> None:
489
+ """Configure worker capacity before the runtime is first started.
490
+
491
+ Background overload policy is intentionally explicit and currently fixed
492
+ to rejection: excess work fails with BackgroundQueueFull rather than
493
+ growing an unbounded queue or blocking the scheduler thread.
494
+ """
495
+ global _runtime_config
496
+ with _runtime_lock:
497
+ if _runtime is not None and not _runtime.closed:
498
+ raise RuntimeError("Task.configure() must be called before runtime startup")
499
+ config = dict(_runtime_config)
500
+ if max_workers is not None:
501
+ if isinstance(max_workers, bool) or not isinstance(max_workers, int) or max_workers < 1:
502
+ raise ValueError("max_workers must be an integer >= 1")
503
+ config["max_workers"] = max_workers
504
+ if max_pending_background is not None:
505
+ if (
506
+ isinstance(max_pending_background, bool)
507
+ or not isinstance(max_pending_background, int)
508
+ or max_pending_background < 0
509
+ ):
510
+ raise ValueError("max_pending_background must be an integer >= 0")
511
+ config["max_pending_background"] = max_pending_background
512
+ _runtime_config = config
513
+
514
+ @classmethod
515
+ def set_exception_handler(
516
+ cls, handler: Callable[[dict[str, Any]], Any] | None
517
+ ) -> None:
518
+ """Install a handler for unexpected scheduler/backend errors.
519
+
520
+ Handler receives a context mapping containing at least ``message`` and
521
+ ``exception``. If no handler is installed, these failures are logged.
522
+ """
523
+ if handler is not None and not callable(handler):
524
+ raise TypeError("handler must be callable or None")
525
+ _get_runtime().set_exception_handler(handler)
526
+
527
+ @classmethod
528
+ def run(
529
+ cls,
530
+ work: Any,
531
+ /,
532
+ *args: Any,
533
+ name: str | None = None,
534
+ token: CancellationToken | None = None,
535
+ pass_token: bool = False,
536
+ **kwargs: Any,
537
+ ) -> "Task[Any]":
538
+ """Schedule a coroutine/awaitable or ordinary callable.
539
+
540
+ Ordinary functions run in the bounded worker pool. ``token`` controls
541
+ the Task's cancellation. With ``pass_token=True``, that same token is
542
+ also supplied to the target callable as the keyword argument ``token``
543
+ (useful for cooperative worker functions).
544
+ """
545
+ if inspect.isawaitable(work):
546
+ if args or kwargs or pass_token:
547
+ if inspect.iscoroutine(work):
548
+ work.close()
549
+ raise TypeError("args/kwargs/pass_token cannot be supplied with an awaitable instance")
550
+ elif not callable(work):
551
+ raise TypeError("Task.run() expects an awaitable or a callable")
552
+ if token is not None and not isinstance(token, CancellationToken):
553
+ if inspect.iscoroutine(work):
554
+ work.close()
555
+ raise TypeError("token must be a CancellationToken")
556
+ if pass_token and token is None:
557
+ if inspect.iscoroutine(work):
558
+ work.close()
559
+ raise ValueError("pass_token=True requires an explicit token")
560
+ runtime = _get_runtime()
561
+ task = cls(runtime, name=name, token=token)
562
+ task._pass_token = pass_token
563
+ runtime._submit(task, work, args, kwargs)
564
+ return task
565
+
566
+ @classmethod
567
+ def when_all(
568
+ cls,
569
+ *tasks: Any,
570
+ name: str | None = None,
571
+ token: CancellationToken | None = None,
572
+ ) -> "Task[list[Any]]":
573
+ """Complete when every supplied task completes, preserving input order.
574
+
575
+ Accepts Task instances and raw awaitables (including coroutine objects)
576
+ as positional arguments or in one iterable. Raw awaitables are scheduled
577
+ automatically. Child tasks are not cancelled if this
578
+ aggregate task is cancelled. A single child failure is re-raised; if
579
+ several children fail, TaskAggregateError exposes all their exceptions.
580
+ If there are no tasks, the result is an empty list.
581
+ """
582
+ if token is not None and not isinstance(token, CancellationToken):
583
+ raise TypeError("token must be a CancellationToken")
584
+ normalized = _normalize_tasks(tasks, allow_empty=True)
585
+ return cast(
586
+ "Task[list[Any]]",
587
+ cls.run(_when_all_coroutine, normalized, name=name, token=token),
588
+ )
589
+
590
+ @classmethod
591
+ def when_any(
592
+ cls,
593
+ *tasks: Any,
594
+ name: str | None = None,
595
+ token: CancellationToken | None = None,
596
+ ) -> "Task[Task[Any]]":
597
+ """Complete when one supplied task completes and return that task.
598
+
599
+ Accepts Task instances and raw awaitables (including coroutine objects)
600
+ positionally or in one iterable. Raw awaitables are scheduled
601
+ automatically. The winning task is returned even if it failed or was
602
+ canceled; inspect/await that task to observe
603
+ its outcome. Losing tasks continue running. An empty input is invalid.
604
+ If multiple tasks are already complete at observation time, input order
605
+ breaks the tie deterministically.
606
+ """
607
+ if token is not None and not isinstance(token, CancellationToken):
608
+ raise TypeError("token must be a CancellationToken")
609
+ normalized = _normalize_tasks(tasks, allow_empty=False)
610
+ return cast(
611
+ "Task[Task[Any]]",
612
+ cls.run(_when_any_coroutine, normalized, name=name, token=token),
613
+ )
614
+
615
+ @staticmethod
616
+ async def collect(*tasks: Task[Any] | list[Task[Any]]):
617
+ """
618
+ Yields tasks as they complete. Accepts either a list or loose positional tasks.
619
+ """
620
+ # FIX: If it's a list or tuple, extract and cast the FIRST element to a mutable list
621
+ if len(tasks) == 1 and isinstance(tasks[0], (list, tuple)):
622
+ pending = list(tasks[0]) # Converts the inner container into a fresh list
623
+ else:
624
+ pending = list(tasks) # Converts the *tasks tuple into a fresh list
625
+
626
+ while pending:
627
+ finished_task = await Task.when_any(pending)
628
+ pending.remove(finished_task) # Safe to call, because pending is ALWAYS a list!
629
+ yield finished_task
630
+
631
+ @classmethod
632
+ def delay(
633
+ cls,
634
+ seconds: float,
635
+ *,
636
+ token: CancellationToken | None = None,
637
+ name: str | None = None,
638
+ ) -> "Task[None]":
639
+ """Return a task completed after seconds without blocking a worker."""
640
+ if seconds < 0:
641
+ raise ValueError("delay must be >= 0")
642
+ if token is not None and not isinstance(token, CancellationToken):
643
+ raise TypeError("token must be a CancellationToken")
644
+ runtime = _get_runtime()
645
+ task: Task[None] = cls(runtime, name=name or f"Delay({seconds:g}s)", token=token)
646
+ # During graceful cancellation, allow an already-running coroutine to
647
+ # create a delay for async cleanup. External callers are still rejected.
648
+ runtime._submit_coroutine(
649
+ task,
650
+ _delay_coroutine(runtime, seconds),
651
+ allow_during_shutdown=runtime.is_runtime_thread(),
652
+ )
653
+ return task
654
+
655
+ @classmethod
656
+ def yield_now(cls) -> _YieldNow:
657
+ """Awaitable that lets other ready tasks run."""
658
+ return _YieldNow(_get_runtime())
659
+
660
+ @classmethod
661
+ def shutdown(
662
+ cls,
663
+ *,
664
+ grace_timeout: float = 1.0,
665
+ worker_timeout: float = 2.0,
666
+ ) -> ShutdownReport:
667
+ """Stop the singleton runtime and report work that could not be stopped.
668
+
669
+ Coroutines receive cancellation and get up to ``grace_timeout`` to unwind.
670
+ Running synchronous worker functions are allowed up to ``worker_timeout``
671
+ to finish; Python cannot forcibly stop them. Pending worker jobs are cancelled.
672
+ The singleton can be lazily recreated by later Task.run calls.
673
+ """
674
+ if grace_timeout < 0 or worker_timeout < 0:
675
+ raise ValueError("shutdown timeouts must be >= 0")
676
+ global _runtime
677
+ with _runtime_lock:
678
+ runtime = _runtime
679
+ if runtime is None:
680
+ return ShutdownReport(True, 0, 0, False)
681
+ report = runtime.shutdown(grace_timeout=grace_timeout, worker_timeout=worker_timeout)
682
+ with _runtime_lock:
683
+ if _runtime is runtime and runtime.closed:
684
+ _runtime = None
685
+ return report
686
+
687
+ @property
688
+ def cancellation_token(self) -> CancellationToken:
689
+ return self._token
690
+
691
+ def done(self) -> bool:
692
+ return self._completed.is_set()
693
+
694
+ def cancelled(self) -> bool:
695
+ return self._completed.is_set() and self._state == self.CANCELLED
696
+
697
+ def result(self, timeout: float | None = None) -> T:
698
+ """Block an external thread until completion and return/raise the result."""
699
+ if not self.done() and self._runtime.is_runtime_thread():
700
+ raise RuntimeError("Task.result() would block the runtime thread; use 'await task'")
701
+ if not self._completed.wait(timeout):
702
+ raise TimeoutError(f"{self.name} did not complete before the timeout")
703
+ return self._raise_or_return()
704
+
705
+ def exception(self, timeout: float | None = None) -> BaseException | None:
706
+ if not self.done() and self._runtime.is_runtime_thread():
707
+ raise RuntimeError("Task.exception() would block the runtime thread")
708
+ if not self._completed.wait(timeout):
709
+ raise TimeoutError(f"{self.name} did not complete before the timeout")
710
+ if self._state == self.CANCELLED:
711
+ raise OperationCanceledError(f"{self.name} was cancelled")
712
+ return self._exception
713
+
714
+ def cancel(self) -> bool:
715
+ """Request cancellation of this Task, independently of other token users."""
716
+ if self.done():
717
+ return False
718
+ try:
719
+ self._runtime._post(lambda: self._cancel_on_runtime(), internal=True)
720
+ except RuntimeShutdownError:
721
+ return False
722
+ return True
723
+
724
+ def __await__(self):
725
+ if self.done():
726
+ return self._raise_or_return()
727
+ result = yield _WaitTask(self._runtime, self)
728
+ return result
729
+
730
+ def _raise_or_return(self) -> T:
731
+ if self._state == self.CANCELLED:
732
+ raise OperationCanceledError(f"{self.name} was cancelled")
733
+ if self._state == self.FAILED:
734
+ assert self._exception is not None
735
+ raise self._exception
736
+ if self._state != self.SUCCEEDED:
737
+ raise InvalidStateError(f"{self.name} has not completed")
738
+ return cast(T, self._result)
739
+
740
+ def _bind_token(self) -> None:
741
+ if not self._token.can_be_cancelled or self._token_registration is not None:
742
+ return
743
+ self._token_registration = self._token.register(self.cancel)
744
+
745
+ def _cancel_on_runtime(self) -> None:
746
+ if self.done():
747
+ return
748
+ self._cancel_requested = True
749
+ self._runtime._detach_wait(self)
750
+ if not self._running:
751
+ self._runtime._schedule(self, exc=OperationCanceledError(f"{self.name} cancelled"))
752
+
753
+ def _finish(
754
+ self,
755
+ state: str,
756
+ result: Any = None,
757
+ exception: BaseException | None = None,
758
+ ) -> None:
759
+ if self.done():
760
+ return
761
+ self._state = state
762
+ self._result = result
763
+ self._exception = exception
764
+ registration = self._token_registration
765
+ self._token_registration = None
766
+ if registration is not None:
767
+ registration.dispose()
768
+ self._completed.set()
769
+ self._runtime._task_completed(self)
770
+
771
+ def _step(self) -> None:
772
+ if self.done():
773
+ return
774
+ coroutine = self._coroutine
775
+ if coroutine is None:
776
+ self._finish(self.FAILED, exception=RuntimeError("task has no coroutine"))
777
+ return
778
+ value = self._next_value
779
+ exc = self._next_exception
780
+ self._next_value = None
781
+ self._next_exception = None
782
+ self._running = True
783
+ previous_task = self._runtime._current_task
784
+ self._runtime._current_task = self
785
+ if isinstance(exc, OperationCanceledError):
786
+ self._cancel_requested = False
787
+ try:
788
+ try:
789
+ if exc is not None:
790
+ yielded = coroutine.throw(exc)
791
+ else:
792
+ yielded = coroutine.send(value)
793
+ except StopIteration as stop:
794
+ self._finish(self.SUCCEEDED, result=stop.value)
795
+ return
796
+ except OperationCanceledError as cancelled:
797
+ self._finish(self.CANCELLED, exception=cancelled)
798
+ return
799
+ except BaseException as error:
800
+ self._finish(self.FAILED, exception=error)
801
+ return
802
+ finally:
803
+ self._running = False
804
+ self._runtime._current_task = previous_task
805
+ if self._cancel_requested:
806
+ self._runtime._schedule(self, exc=OperationCanceledError(f"{self.name} cancelled"))
807
+ return
808
+ try:
809
+ self._runtime._handle_yielded(self, yielded)
810
+ except BaseException as error:
811
+ self._runtime._detach_wait(self)
812
+ self._runtime._schedule(self, exc=error)
813
+
814
+ def __repr__(self) -> str:
815
+ return f"<{self.name} state={self._state}>"
816
+
817
+
818
+ def _normalize_tasks(items: tuple[Any, ...], *, allow_empty: bool) -> tuple[Task[Any], ...]:
819
+ """Normalize task-combinator inputs, scheduling raw awaitables as Tasks.
820
+
821
+ Accepts Task instances, coroutine objects, and other awaitables supported by
822
+ this runtime. A single iterable is also accepted. If validation fails,
823
+ coroutine objects supplied to this call are closed so Python does not emit
824
+ "coroutine was never awaited" warnings for rejected inputs.
825
+ """
826
+ # A single awaitable is one operation, not an iterable of operations.
827
+ if (
828
+ len(items) == 1
829
+ and not isinstance(items[0], Task)
830
+ and not inspect.isawaitable(items[0])
831
+ and not isinstance(items[0], (str, bytes, bytearray))
832
+ ):
833
+ candidate = items[0]
834
+ try:
835
+ items = tuple(candidate)
836
+ except TypeError:
837
+ pass
838
+
839
+ if not items and not allow_empty:
840
+ raise ValueError("Task.when_any() requires at least one task or awaitable")
841
+
842
+ def close_raw_coroutines(values: tuple[Any, ...]) -> None:
843
+ seen: set[int] = set()
844
+ for value in values:
845
+ if inspect.iscoroutine(value) and id(value) not in seen:
846
+ seen.add(id(value))
847
+ try:
848
+ value.close()
849
+ except BaseException:
850
+ _LOG.exception("Failed to close rejected coroutine supplied to Task.when_all/when_any")
851
+
852
+ # Validate everything before scheduling anything, avoiding partial starts
853
+ # when the caller accidentally mixes a coroutine with a non-awaitable.
854
+ invalid = [item for item in items if not isinstance(item, Task) and not inspect.isawaitable(item)]
855
+ if invalid:
856
+ close_raw_coroutines(items)
857
+ raise TypeError(
858
+ "Task.when_all/when_any accepts Task instances or awaitables; "
859
+ f"got {type(invalid[0]).__name__}"
860
+ )
861
+
862
+ # The same one-shot coroutine/awaitable cannot safely be scheduled twice.
863
+ seen_awaitables: set[int] = set()
864
+ for item in items:
865
+ if isinstance(item, Task):
866
+ continue
867
+ identity = id(item)
868
+ if identity in seen_awaitables:
869
+ close_raw_coroutines(items)
870
+ raise ValueError("the same awaitable instance was supplied more than once")
871
+ seen_awaitables.add(identity)
872
+
873
+ normalized: list[Task[Any]] = []
874
+ index = 0
875
+ try:
876
+ for index, item in enumerate(items):
877
+ if isinstance(item, Task):
878
+ normalized.append(item)
879
+ else:
880
+ # Task.run takes ownership of the awaitable and schedules it.
881
+ normalized.append(Task.run(item))
882
+ except BaseException:
883
+ # The failed Task.run closes its own coroutine when possible. Close any
884
+ # later coroutine inputs which have not been handed to the runtime.
885
+ close_raw_coroutines(tuple(items[index + 1 :]))
886
+ raise
887
+ return tuple(normalized)
888
+
889
+
890
+ async def _when_all_coroutine(tasks: tuple[Task[Any], ...]) -> list[Any]:
891
+ await _WaitMany(_get_runtime(), tasks, "all")
892
+ failures = tuple(task._exception for task in tasks if task._state == Task.FAILED)
893
+ if len(failures) == 1:
894
+ assert failures[0] is not None
895
+ raise failures[0]
896
+ if len(failures) > 1:
897
+ raise AggregateError(cast(tuple[Exception, ...], failures))
898
+ if any(task._state == Task.CANCELLED for task in tasks):
899
+ raise OperationCanceledError("One or more tasks passed to Task.when_all() were cancelled")
900
+ return [task._result for task in tasks]
901
+
902
+
903
+ async def _when_any_coroutine(tasks: tuple[Task[Any], ...]) -> Task[Any]:
904
+ winner = await _WaitMany(_get_runtime(), tasks, "any")
905
+ return cast(Task[Any], winner)
906
+
907
+
908
+ class IO:
909
+ """Selector-backed awaitables for nonblocking sockets."""
910
+
911
+ @staticmethod
912
+ def wait_readable(
913
+ sock: socket.socket, *, token: CancellationToken | None = None
914
+ ) -> _WaitIO:
915
+ if token is not None and not isinstance(token, CancellationToken):
916
+ raise TypeError("token must be a CancellationToken")
917
+ return _WaitIO(_get_runtime(), sock, "read", token)
918
+
919
+ @staticmethod
920
+ def wait_writable(
921
+ sock: socket.socket, *, token: CancellationToken | None = None
922
+ ) -> _WaitIO:
923
+ if token is not None and not isinstance(token, CancellationToken):
924
+ raise TypeError("token must be a CancellationToken")
925
+ return _WaitIO(_get_runtime(), sock, "write", token)
926
+
927
+
928
+ class _Runtime:
929
+ """Private singleton scheduler: ready queue + timers + selector + workers."""
930
+
931
+ def __init__(self, *, max_workers: int, max_pending_background: int) -> None:
932
+ self._selector = selectors.DefaultSelector()
933
+ self._wake_reader, self._wake_writer = socket.socketpair()
934
+ self._wake_reader.setblocking(False)
935
+ self._wake_writer.setblocking(False)
936
+ self._selector.register(self._wake_reader, selectors.EVENT_READ, data=_WAKE)
937
+ self._ready: deque[Task[Any]] = deque()
938
+ self._tasks: set[Task[Any]] = set()
939
+ self._timers: list[tuple[float, int, Task[Any], _Sleep]] = []
940
+ self._timer_sequence = count()
941
+ self._io_waiters: dict[socket.socket, dict[str, set[Task[Any]]]] = {}
942
+ self._commands: queue.SimpleQueue[Callable[[], None]] = queue.SimpleQueue()
943
+ self._max_workers = max_workers
944
+ self._max_pending_background = max_pending_background
945
+ self._pending_background: deque[tuple[Task[Any], _BackgroundCall]] = deque()
946
+ self._background_active = 0 # scheduler-thread-owned submitted call count
947
+ self._executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="tinytask-io")
948
+ self._inflight_lock = threading.Condition()
949
+ self._inflight_futures: set[Future[Any]] = set()
950
+ self._current_task: Task[Any] | None = None
951
+ self._thread_id: int | None = None
952
+ self._start_lock = threading.Lock()
953
+ self._started = threading.Event()
954
+ self._state_lock = threading.RLock()
955
+ self._shutdown_requested = False
956
+ self._shutdown_deadline: float | None = None
957
+ self._stopping = False
958
+ self._closed = False
959
+ self._shutdown_tasks_cancelled = 0
960
+ self._thread: threading.Thread | None = None
961
+ self._exception_handler: Callable[[dict[str, Any]], Any] | None = None
962
+
963
+ @property
964
+ def closed(self) -> bool:
965
+ with self._state_lock:
966
+ return self._closed
967
+
968
+ def _ensure_started(self) -> None:
969
+ if self._started.is_set():
970
+ return
971
+ with self._start_lock:
972
+ if self._started.is_set():
973
+ return
974
+ with self._state_lock:
975
+ if self._closed or self._shutdown_requested:
976
+ raise RuntimeShutdownError("TinyTask runtime is shutting down")
977
+ self._thread = threading.Thread(target=self._loop, name="TinyTask-runtime", daemon=True)
978
+ self._thread.start()
979
+ self._started.wait()
980
+
981
+ def is_runtime_thread(self) -> bool:
982
+ return self._thread_id == threading.get_ident()
983
+
984
+ def set_exception_handler(self, handler: Callable[[dict[str, Any]], Any] | None) -> None:
985
+ self._exception_handler = handler
986
+
987
+ def _report_exception(self, context: dict[str, Any], error: BaseException) -> None:
988
+ handler = self._exception_handler
989
+ if handler is not None:
990
+ try:
991
+ handler({**context, "exception": error})
992
+ return
993
+ except Exception:
994
+ _LOG.exception("TinyTask exception handler itself failed")
995
+ message = str(context.get("message", "Unexpected TinyTask runtime error"))
996
+ _LOG.error("%s", message, exc_info=(type(error), error, error.__traceback__))
997
+
998
+ def _post(self, callback: Callable[[], None], *, internal: bool = False) -> None:
999
+ self._ensure_started()
1000
+ execute_inline = False
1001
+ with self._state_lock:
1002
+ if self._closed or self._stopping or (self._shutdown_requested and not internal):
1003
+ raise RuntimeShutdownError("TinyTask runtime is shutting down")
1004
+ if self.is_runtime_thread():
1005
+ execute_inline = True
1006
+ else:
1007
+ self._commands.put(callback)
1008
+ if execute_inline:
1009
+ callback()
1010
+ return
1011
+ try:
1012
+ self._wake_writer.send(b"\x00")
1013
+ except (BlockingIOError, OSError):
1014
+ # A full wake-up socket already guarantees selector readability.
1015
+ pass
1016
+
1017
+ def _submit(self, task: Task[Any], work: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None:
1018
+ try:
1019
+ self._post(lambda: self._start_target(task, work, args, kwargs))
1020
+ except BaseException:
1021
+ if inspect.iscoroutine(work):
1022
+ work.close()
1023
+ raise
1024
+
1025
+ def _submit_coroutine(
1026
+ self,
1027
+ task: Task[Any],
1028
+ coroutine: Coroutine[Any, Any, Any],
1029
+ *,
1030
+ allow_during_shutdown: bool = False,
1031
+ ) -> None:
1032
+ try:
1033
+ self._post(
1034
+ lambda: self._start_coroutine(task, coroutine),
1035
+ internal=allow_during_shutdown,
1036
+ )
1037
+ except BaseException:
1038
+ coroutine.close()
1039
+ raise
1040
+
1041
+ def _start_target(
1042
+ self, task: Task[Any], work: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
1043
+ ) -> None:
1044
+ if task.done():
1045
+ if inspect.iscoroutine(work):
1046
+ work.close()
1047
+ return
1048
+ if task._token.is_cancellation_requested:
1049
+ if inspect.iscoroutine(work):
1050
+ work.close()
1051
+ task._finish(Task.CANCELLED, exception=OperationCanceledError(f"{task.name} cancelled before start"))
1052
+ return
1053
+ try:
1054
+ call_kwargs = dict(kwargs)
1055
+ if task._pass_token:
1056
+ call_kwargs["token"] = task._token
1057
+ if inspect.isawaitable(work):
1058
+ coroutine = work if inspect.iscoroutine(work) else _await_any(work)
1059
+ elif _is_async_callable(work):
1060
+ coroutine = _invoke_async(work, args, call_kwargs)
1061
+ else:
1062
+ coroutine = _invoke_background(self, work, args, call_kwargs)
1063
+ self._start_coroutine(task, coroutine)
1064
+ except BaseException as error:
1065
+ self._tasks.add(task)
1066
+ task._finish(Task.FAILED, exception=error)
1067
+
1068
+ def _start_coroutine(self, task: Task[Any], coroutine: Coroutine[Any, Any, Any]) -> None:
1069
+ if task.done():
1070
+ coroutine.close()
1071
+ return
1072
+ task._coroutine = coroutine
1073
+ self._tasks.add(task)
1074
+ task._bind_token()
1075
+ self._schedule(task)
1076
+
1077
+ def _schedule(
1078
+ self, task: Task[Any], value: Any = None, exc: BaseException | None = None
1079
+ ) -> None:
1080
+ if task.done():
1081
+ return
1082
+ if task._queued:
1083
+ if exc is not None:
1084
+ task._next_exception = exc
1085
+ task._next_value = None
1086
+ elif task._next_exception is None:
1087
+ task._next_value = value
1088
+ return
1089
+ task._queued = True
1090
+ task._next_value = value
1091
+ task._next_exception = exc
1092
+ self._ready.append(task)
1093
+
1094
+ def _task_completed(self, task: Task[Any]) -> None:
1095
+ self._tasks.discard(task)
1096
+ if self._shutdown_requested and task._state == Task.CANCELLED:
1097
+ self._shutdown_tasks_cancelled += 1
1098
+ self._detach_wait(task)
1099
+ for waiter in list(task._awaiters):
1100
+ operation = waiter._waiting_on
1101
+ if isinstance(operation, _WaitTask) and operation.task is task:
1102
+ self._detach_wait(waiter)
1103
+ if task._state == Task.SUCCEEDED:
1104
+ self._schedule(waiter, value=task._result)
1105
+ elif task._state == Task.CANCELLED:
1106
+ self._schedule(waiter, exc=OperationCanceledError(f"{task.name} was cancelled"))
1107
+ else:
1108
+ assert task._exception is not None
1109
+ self._schedule(waiter, exc=task._exception)
1110
+ elif isinstance(operation, _WaitMany) and task in operation.tasks:
1111
+ if operation.mode == "any":
1112
+ self._detach_wait(waiter)
1113
+ self._schedule(waiter, value=task)
1114
+ elif all(child.done() for child in operation.tasks):
1115
+ self._detach_wait(waiter)
1116
+ self._schedule(waiter, value=operation.tasks)
1117
+ task._awaiters.clear()
1118
+
1119
+ def _handle_yielded(self, task: Task[Any], yielded: Any) -> None:
1120
+ if not isinstance(yielded, _Operation):
1121
+ raise RuntimeError(
1122
+ f"{task.name} yielded unsupported object {yielded!r}; "
1123
+ "this runtime only supports TinyTask awaitables"
1124
+ )
1125
+ if yielded.runtime is not self:
1126
+ raise RuntimeError("awaitable belongs to a different runtime")
1127
+ if yielded.token.is_cancellation_requested:
1128
+ self._schedule(task, exc=OperationCanceledError("Operation token was cancelled"))
1129
+ return
1130
+
1131
+ if isinstance(yielded, _YieldNow):
1132
+ self._schedule(task)
1133
+ elif isinstance(yielded, _Sleep):
1134
+ if yielded.delay == 0:
1135
+ self._schedule(task)
1136
+ else:
1137
+ task._waiting_on = yielded
1138
+ heapq.heappush(self._timers, (yielded.deadline, next(self._timer_sequence), task, yielded))
1139
+ self._register_operation_token(task, yielded)
1140
+ elif isinstance(yielded, _WaitIO):
1141
+ self._arm_io(task, yielded)
1142
+ self._register_operation_token(task, yielded)
1143
+ elif isinstance(yielded, _WaitMany):
1144
+ if yielded.mode == "any":
1145
+ # Deterministic tie-breaking for tasks already completed when
1146
+ # the operation is first observed: first completed in input order.
1147
+ winner = next((child for child in yielded.tasks if child.done()), None)
1148
+ if winner is not None:
1149
+ self._schedule(task, value=winner)
1150
+ return
1151
+ elif all(child.done() for child in yielded.tasks):
1152
+ self._schedule(task, value=yielded.tasks)
1153
+ return
1154
+
1155
+ pending = [child for child in yielded.tasks if not child.done()]
1156
+ if any(child._runtime is not self for child in pending):
1157
+ raise RuntimeError("cannot combine pending tasks from different runtimes")
1158
+ task._waiting_on = yielded
1159
+ for child in pending:
1160
+ child._awaiters.add(task)
1161
+ elif isinstance(yielded, _WaitTask):
1162
+ target = yielded.task
1163
+ if target is task:
1164
+ raise RuntimeError("a task cannot await itself")
1165
+ if target.done():
1166
+ if target._state == Task.SUCCEEDED:
1167
+ self._schedule(task, value=target._result)
1168
+ elif target._state == Task.CANCELLED:
1169
+ self._schedule(task, exc=OperationCanceledError(f"{target.name} was cancelled"))
1170
+ else:
1171
+ assert target._exception is not None
1172
+ self._schedule(task, exc=target._exception)
1173
+ else:
1174
+ task._waiting_on = yielded
1175
+ target._awaiters.add(task)
1176
+ self._register_operation_token(task, yielded)
1177
+ elif isinstance(yielded, _BackgroundCall):
1178
+ task._waiting_on = yielded
1179
+ if self._background_active < self._max_workers and not self._pending_background:
1180
+ self._start_background(task, yielded)
1181
+ elif len(self._pending_background) < self._max_pending_background:
1182
+ self._pending_background.append((task, yielded))
1183
+ self._register_operation_token(task, yielded)
1184
+ else:
1185
+ task._waiting_on = None
1186
+ raise BackgroundQueueFull(
1187
+ f"background queue is full (max_workers={self._max_workers}, "
1188
+ f"max_pending_background={self._max_pending_background})"
1189
+ )
1190
+ else:
1191
+ raise RuntimeError(f"unsupported operation type: {type(yielded).__name__}")
1192
+
1193
+ def _register_operation_token(self, task: Task[Any], operation: _Operation) -> None:
1194
+ token = operation.token
1195
+ if not token.can_be_cancelled:
1196
+ return
1197
+
1198
+ def on_cancel() -> None:
1199
+ try:
1200
+ self._post(lambda: self._cancel_operation_wait(task, operation), internal=True)
1201
+ except RuntimeShutdownError:
1202
+ pass
1203
+
1204
+ operation._token_registration = token.register(on_cancel)
1205
+
1206
+ def _dispose_operation_token(self, operation: _Operation) -> None:
1207
+ registration = operation._token_registration
1208
+ operation._token_registration = None
1209
+ if registration is not None:
1210
+ registration.dispose()
1211
+
1212
+ def _cancel_operation_wait(self, task: Task[Any], operation: _Operation) -> None:
1213
+ if task.done() or task._waiting_on is not operation:
1214
+ return
1215
+ self._detach_wait(task)
1216
+ self._schedule(task, exc=OperationCanceledError("Operation token was cancelled"))
1217
+
1218
+ def _start_background(self, task: Task[Any], operation: _BackgroundCall) -> None:
1219
+ if task.done() or task._waiting_on is not operation or self._shutdown_requested:
1220
+ return
1221
+ self._background_active += 1
1222
+ operation._counted_active = True
1223
+ try:
1224
+ future = self._executor.submit(operation.function, *operation.args, **operation.kwargs)
1225
+ except BaseException as error:
1226
+ self._background_active -= 1
1227
+ operation._counted_active = False
1228
+ task._waiting_on = None
1229
+ self._dispose_operation_token(operation)
1230
+ self._schedule(task, exc=error)
1231
+ self._pump_background_queue()
1232
+ return
1233
+ operation.future = future
1234
+ with self._inflight_lock:
1235
+ self._inflight_futures.add(future)
1236
+ future.add_done_callback(
1237
+ lambda completed, task=task, op=operation: self._background_worker_done(task, op, completed)
1238
+ )
1239
+
1240
+ def _background_worker_done(
1241
+ self, task: Task[Any], operation: _BackgroundCall, future: Future[Any]
1242
+ ) -> None:
1243
+ with self._inflight_lock:
1244
+ self._inflight_futures.discard(future)
1245
+ self._inflight_lock.notify_all()
1246
+ try:
1247
+ self._post(lambda: self._background_finished(task, operation, future), internal=True)
1248
+ except RuntimeShutdownError:
1249
+ # The future's terminal state is still recorded in _inflight_futures;
1250
+ # after shutdown there is no scheduler left to resume a task.
1251
+ pass
1252
+
1253
+ def _background_finished(
1254
+ self, task: Task[Any], operation: _BackgroundCall, future: Future[Any]
1255
+ ) -> None:
1256
+ if operation._counted_active:
1257
+ operation._counted_active = False
1258
+ self._background_active = max(0, self._background_active - 1)
1259
+ if not task.done() and task._waiting_on is operation:
1260
+ task._waiting_on = None
1261
+ self._dispose_operation_token(operation)
1262
+ try:
1263
+ value = future.result()
1264
+ except BaseException as error:
1265
+ self._schedule(task, exc=error)
1266
+ else:
1267
+ self._schedule(task, value=value)
1268
+ self._pump_background_queue()
1269
+
1270
+ def _pump_background_queue(self) -> None:
1271
+ if self._shutdown_requested:
1272
+ return
1273
+ while self._background_active < self._max_workers and self._pending_background:
1274
+ task, operation = self._pending_background.popleft()
1275
+ if task.done() or task._waiting_on is not operation:
1276
+ continue
1277
+ self._dispose_operation_token(operation)
1278
+ self._start_background(task, operation)
1279
+
1280
+ def _arm_io(self, task: Task[Any], operation: _WaitIO) -> None:
1281
+ if operation.sock.fileno() < 0:
1282
+ raise ValueError("cannot wait on a closed socket")
1283
+ task._waiting_on = operation
1284
+ sides = self._io_waiters.setdefault(operation.sock, {"read": set(), "write": set()})
1285
+ sides[operation.mode].add(task)
1286
+ self._refresh_socket(operation.sock)
1287
+
1288
+ def _refresh_socket(self, sock: socket.socket) -> None:
1289
+ sides = self._io_waiters.get(sock)
1290
+ mask = 0
1291
+ if sides:
1292
+ if sides["read"]:
1293
+ mask |= selectors.EVENT_READ
1294
+ if sides["write"]:
1295
+ mask |= selectors.EVENT_WRITE
1296
+ try:
1297
+ self._selector.get_key(sock)
1298
+ registered = True
1299
+ except (KeyError, ValueError, OSError):
1300
+ registered = False
1301
+ try:
1302
+ if mask:
1303
+ if registered:
1304
+ self._selector.modify(sock, mask, data=None)
1305
+ else:
1306
+ self._selector.register(sock, mask, data=None)
1307
+ else:
1308
+ if registered:
1309
+ try:
1310
+ self._selector.unregister(sock)
1311
+ except (KeyError, ValueError, OSError):
1312
+ pass
1313
+ self._io_waiters.pop(sock, None)
1314
+ except (KeyError, ValueError, OSError) as error:
1315
+ # A socket may be closed while tasks are waiting on it. Fail those
1316
+ # waits instead of allowing a selector error to kill the scheduler.
1317
+ sides = self._io_waiters.pop(sock, None) or {"read": set(), "write": set()}
1318
+ affected = set(sides["read"]) | set(sides["write"])
1319
+ for task in affected:
1320
+ operation = task._waiting_on
1321
+ if isinstance(operation, _WaitIO) and operation.sock is sock:
1322
+ task._waiting_on = None
1323
+ self._dispose_operation_token(operation)
1324
+ self._schedule(task, exc=error)
1325
+ try:
1326
+ self._selector.unregister(sock)
1327
+ except (KeyError, ValueError, OSError):
1328
+ pass
1329
+
1330
+ def _detach_wait(self, task: Task[Any]) -> None:
1331
+ operation = task._waiting_on
1332
+ if operation is None:
1333
+ return
1334
+ task._waiting_on = None
1335
+ self._dispose_operation_token(operation)
1336
+ if isinstance(operation, _WaitIO):
1337
+ sides = self._io_waiters.get(operation.sock)
1338
+ if sides is not None:
1339
+ sides[operation.mode].discard(task)
1340
+ self._refresh_socket(operation.sock)
1341
+ elif isinstance(operation, _WaitTask):
1342
+ operation.task._awaiters.discard(task)
1343
+ elif isinstance(operation, _WaitMany):
1344
+ for target in operation.tasks:
1345
+ target._awaiters.discard(task)
1346
+ elif isinstance(operation, _BackgroundCall):
1347
+ if operation.future is None:
1348
+ try:
1349
+ self._pending_background.remove((task, operation))
1350
+ except ValueError:
1351
+ pass
1352
+ else:
1353
+ # This cancels queued executor work where possible. If already
1354
+ # running, the worker continues and remains tracked until it exits.
1355
+ operation.future.cancel()
1356
+ # Timer entries are lazily discarded when they reach the heap's top.
1357
+
1358
+ def _drain_wakeup(self) -> None:
1359
+ while True:
1360
+ try:
1361
+ if not self._wake_reader.recv(4096):
1362
+ break
1363
+ except BlockingIOError:
1364
+ break
1365
+ except OSError:
1366
+ break
1367
+ while True:
1368
+ try:
1369
+ callback = self._commands.get_nowait()
1370
+ except queue.Empty:
1371
+ break
1372
+ try:
1373
+ callback()
1374
+ except BaseException as error:
1375
+ self._report_exception(
1376
+ {"message": "Unhandled exception in TinyTask scheduler callback", "callback": callback},
1377
+ error,
1378
+ )
1379
+
1380
+ def _process_events(self, events: list[tuple[selectors.SelectorKey, int]]) -> None:
1381
+ for key, mask in events:
1382
+ if key.data is _WAKE:
1383
+ self._drain_wakeup()
1384
+ continue
1385
+ sock: Any = key.fileobj
1386
+ sides = self._io_waiters.get(sock)
1387
+ if sides is None:
1388
+ continue
1389
+ to_wake: set[Task[Any]] = set()
1390
+ if mask & selectors.EVENT_READ:
1391
+ to_wake.update(sides["read"])
1392
+ sides["read"].clear()
1393
+ if mask & selectors.EVENT_WRITE:
1394
+ to_wake.update(sides["write"])
1395
+ sides["write"].clear()
1396
+ for task in to_wake:
1397
+ operation = task._waiting_on
1398
+ if isinstance(operation, _WaitIO) and operation.sock is sock:
1399
+ task._waiting_on = None
1400
+ self._dispose_operation_token(operation)
1401
+ self._schedule(task)
1402
+ self._refresh_socket(sock)
1403
+
1404
+ def _expire_timers(self) -> None:
1405
+ now = time.monotonic()
1406
+ while self._timers and self._timers[0][0] <= now:
1407
+ _, _, task, operation = heapq.heappop(self._timers)
1408
+ if task.done() or task._waiting_on is not operation:
1409
+ continue
1410
+ task._waiting_on = None
1411
+ self._dispose_operation_token(operation)
1412
+ self._schedule(task)
1413
+
1414
+ def _next_timeout(self) -> float | None:
1415
+ while self._timers:
1416
+ deadline, _, task, operation = self._timers[0]
1417
+ if task.done() or task._waiting_on is not operation:
1418
+ heapq.heappop(self._timers)
1419
+ else:
1420
+ return max(0.0, deadline - time.monotonic())
1421
+ return None
1422
+
1423
+ def _begin_shutdown(self) -> None:
1424
+ self._shutdown_deadline = time.monotonic() + self._shutdown_grace
1425
+ for task in list(self._tasks):
1426
+ task._cancel_on_runtime()
1427
+ # Anything waiting in our bounded background queue is removed as its
1428
+ # task receives cancellation. Do not admit more background work at shutdown.
1429
+
1430
+ def _loop(self) -> None:
1431
+ self._thread_id = threading.get_ident()
1432
+ self._started.set()
1433
+ try:
1434
+ while not self._stopping:
1435
+ self._drain_wakeup()
1436
+ if self._shutdown_requested:
1437
+ if not self._tasks:
1438
+ self._stopping = True
1439
+ elif self._shutdown_deadline is not None and time.monotonic() >= self._shutdown_deadline:
1440
+ self._stopping = True
1441
+ while self._ready and not self._stopping:
1442
+ task = self._ready.popleft()
1443
+ task._queued = False
1444
+ task._step()
1445
+ self._drain_wakeup()
1446
+ if self._shutdown_requested:
1447
+ if not self._tasks:
1448
+ self._stopping = True
1449
+ elif self._shutdown_deadline is not None and time.monotonic() >= self._shutdown_deadline:
1450
+ self._stopping = True
1451
+ self._expire_timers()
1452
+ if self._stopping:
1453
+ break
1454
+ if self._ready:
1455
+ continue
1456
+ try:
1457
+ events = self._selector.select(self._next_timeout())
1458
+ except (OSError, ValueError) as error:
1459
+ if self._stopping or self.closed:
1460
+ break
1461
+ self._report_exception(
1462
+ {"message": "Selector failed; stopping TinyTask runtime"}, error
1463
+ )
1464
+ self._stopping = True
1465
+ break
1466
+ try:
1467
+ self._process_events(events)
1468
+ except BaseException as error:
1469
+ self._report_exception(
1470
+ {"message": "Unhandled error while processing selector events"}, error
1471
+ )
1472
+ self._stopping = True
1473
+ except BaseException as error:
1474
+ self._report_exception({"message": "Fatal error in TinyTask scheduler loop"}, error)
1475
+ self._stopping = True
1476
+ finally:
1477
+ # Stop accepting work before releasing scheduler resources. The
1478
+ # closed flag is published after cleanup, so new runtimes do not
1479
+ # overlap this runtime's resource finalization.
1480
+ tasks = list(self._tasks)
1481
+ for task in tasks:
1482
+ self._detach_wait(task)
1483
+ for task in tasks:
1484
+ coroutine = task._coroutine
1485
+ if coroutine is not None:
1486
+ try:
1487
+ coroutine.close()
1488
+ except BaseException as error:
1489
+ self._report_exception(
1490
+ {"message": f"Error closing coroutine for {task.name}"}, error
1491
+ )
1492
+ task._finish(Task.CANCELLED, exception=OperationCanceledError("runtime shutting down"))
1493
+ # Remove scheduler-owned pending calls; they have never entered a worker.
1494
+ self._pending_background.clear()
1495
+ try:
1496
+ self._selector.close()
1497
+ except BaseException as error:
1498
+ self._report_exception({"message": "Error closing selector"}, error)
1499
+ for wake_socket in (self._wake_reader, self._wake_writer):
1500
+ try:
1501
+ wake_socket.close()
1502
+ except OSError:
1503
+ pass
1504
+ with self._state_lock:
1505
+ self._closed = True
1506
+
1507
+ def shutdown(self, *, grace_timeout: float, worker_timeout: float) -> ShutdownReport:
1508
+ if self.is_runtime_thread():
1509
+ raise RuntimeError("Task.shutdown() cannot block from the runtime thread")
1510
+ with self._state_lock:
1511
+ if self._closed:
1512
+ thread = self._thread
1513
+ else:
1514
+ self._shutdown_requested = True
1515
+ self._shutdown_grace = grace_timeout
1516
+ self._commands.put(self._begin_shutdown)
1517
+ thread = self._thread
1518
+ try:
1519
+ self._wake_writer.send(b"\x00")
1520
+ except (BlockingIOError, OSError):
1521
+ pass
1522
+ if thread is not None and thread.is_alive():
1523
+ # Grace timeout is enforced by the loop; the small allowance lets the
1524
+ # loop execute its finalizer after the deadline.
1525
+ thread.join(timeout=grace_timeout + 1.0)
1526
+ runtime_stopped = thread is None or not thread.is_alive()
1527
+ if not runtime_stopped:
1528
+ self._report_exception(
1529
+ {"message": "TinyTask scheduler thread did not stop before timeout"},
1530
+ TimeoutError("scheduler shutdown timed out"),
1531
+ )
1532
+
1533
+ # Cancel executor-queued calls. Already-running functions are not killable.
1534
+ self._executor.shutdown(wait=False, cancel_futures=True)
1535
+ worker_deadline = time.monotonic() + worker_timeout
1536
+ with self._inflight_lock:
1537
+ while self._inflight_futures:
1538
+ remaining = worker_deadline - time.monotonic()
1539
+ if remaining <= 0:
1540
+ break
1541
+ self._inflight_lock.wait(remaining)
1542
+ workers_still_running = len(self._inflight_futures)
1543
+ if workers_still_running == 0:
1544
+ # It is now safe to join worker threads without an unbounded wait.
1545
+ self._executor.shutdown(wait=True, cancel_futures=True)
1546
+ return ShutdownReport(
1547
+ runtime_stopped=runtime_stopped,
1548
+ tasks_cancelled=self._shutdown_tasks_cancelled,
1549
+ workers_still_running=workers_still_running,
1550
+ timed_out=(not runtime_stopped or workers_still_running > 0),
1551
+ )
1552
+
1553
+
1554
+ def _is_async_callable(value: Any) -> bool:
1555
+ if inspect.iscoroutinefunction(value):
1556
+ return True
1557
+ call = getattr(value, "__call__", None)
1558
+ return inspect.iscoroutinefunction(call)
1559
+
1560
+
1561
+ _runtime_lock = threading.RLock()
1562
+ _runtime: _Runtime | None = None
1563
+ _runtime_config: dict[str, int] = {
1564
+ "max_workers": min(32, (os.cpu_count() or 1) + 4),
1565
+ "max_pending_background": 128,
1566
+ }
1567
+
1568
+
1569
+ def _get_runtime() -> _Runtime:
1570
+ global _runtime
1571
+ with _runtime_lock:
1572
+ if _runtime is None or _runtime.closed:
1573
+ _runtime = _Runtime(**_runtime_config)
1574
+ runtime = _runtime
1575
+ runtime._ensure_started()
1576
+ return runtime
1577
+
1578
+
1579
+ def _shutdown_at_exit() -> None:
1580
+ try:
1581
+ Task.shutdown(grace_timeout=0.2, worker_timeout=0.2)
1582
+ except BaseException:
1583
+ _LOG.exception("Error during TinyTask interpreter shutdown")
1584
+
1585
+
1586
+ atexit.register(_shutdown_at_exit)
1587
+
1588
+ def task(async_function: Callable[..., Any]) -> Callable[..., Task[Any]]:
1589
+ @functools.wraps(async_function)
1590
+ def wrapper(*args: Any, **kwargs: Any) -> Task[Any]:
1591
+ return Task.run(async_function, *args, **kwargs)
1592
+
1593
+ return wrapper
1594
+
1595
+
1596
+ def coro_task(async_function: Callable[..., Any]) -> Callable[..., Task[Any]]:
1597
+ """
1598
+ | Wraps an async coroutine function to be compatible with Task.run(coro)
1599
+
1600
+ Use with caution as this transfers execution control to an isolated asyncio loop runner.
1601
+ """
1602
+ import asyncio
1603
+ @functools.wraps(async_function)
1604
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
1605
+ with asyncio.Runner() as runner:
1606
+ # Call the function with its arguments to generate the Coroutine object instance
1607
+ coro = async_function(*args, **kwargs)
1608
+ return runner.run(coro)
1609
+
1610
+ return task(wrapper)