Skip to content

Commit

Permalink
add dispatch.worker and dispatch.batch
Browse files Browse the repository at this point in the history
Signed-off-by: Achille Roussel <[email protected]>
  • Loading branch information
achille-roussel committed May 12, 2024
1 parent 1988b74 commit 723e844
Showing 1 changed file with 30 additions and 2 deletions.
32 changes: 30 additions & 2 deletions src/dispatch/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,10 @@
from __future__ import annotations

import os
import threading
from concurrent import futures
from http.server import ThreadingHTTPServer
from typing import Any, Callable, Coroutine, Optional, TypeVar, overload
from typing import Any, Callable, Coroutine, List, Optional, TypeVar, overload
from urllib.parse import urlsplit

from typing_extensions import ParamSpec, TypeAlias
Expand All @@ -31,6 +32,7 @@
"Status",
"all",
"any",
"batch",
"call",
"function",
"gather",
Expand All @@ -44,7 +46,8 @@
T = TypeVar("T")

_registry: Optional[Registry] = None

_workers: List[Callable[None, None]] = []
_threads: List[threading.Thread] = []

def default_registry():
global _registry
Expand Down Expand Up @@ -89,10 +92,35 @@ def run(init: Optional[Callable[P, None]] = None, *args: P.args, **kwargs: P.kwa
parsed_url = urlsplit("//" + address)
server_address = (parsed_url.hostname or "", parsed_url.port or 0)
server = ThreadingHTTPServer(server_address, Dispatch(default_registry()))

for worker in _workers:
def entrypoint():
try:
worker()
finally:
server.shutdown()
_threads.append(threading.Thread(target=entrypoint))

for thread in _threads:
thread.start()

try:
if init is not None:
init(*args, **kwargs)
server.serve_forever()
finally:
server.shutdown()
server.server_close()

for thread in _threads:
thread.join()

def batch() -> Batch:
"""Create a new batch object."""
return default_registry().batch()


def worker(fn: Callable[None, None]) -> Callable[None, None]:
"""Decorator declaring workers that will be started when dipatch.run is called."""
_workers.append(fn)
return fn

0 comments on commit 723e844

Please sign in to comment.