wool 0.1rc6__py3-none-any.whl → 0.1rc7__py3-none-any.whl

wool/_manager.py

@@ -1,30 +1,31 @@
 from __future__ import annotations
 
 import logging
-import os
-import signal
 from functools import partial
-from multiprocessing.managers import (
-    BaseManager,
-    DictProxy,
-)
+from multiprocessing.managers import BaseManager
+from multiprocessing.managers import DictProxy
 from queue import Empty
+from threading import Event
 from threading import Lock
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, overload
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Callable
+from typing import TypeVar
+from typing import overload
 from uuid import UUID
 from weakref import WeakValueDictionary
 
-from wool._future import WoolFuture
+from wool._event import TaskEvent
+from wool._future import Future
 from wool._queue import TaskQueue
 from wool._typing import PassthroughDecorator
 
 if TYPE_CHECKING:
-    from wool._task import WoolTask
+    from wool._task import Task
 
 
 C = TypeVar("C", bound=Callable[..., Any])
 
-
 _manager_registry = {}
 
 
@@ -41,6 +42,7 @@ def register(
     fn: C,
     /,
     *,
+    exposed: tuple[str, ...] | None = None,
     proxytype: None = None,
     method_to_typeid: None = None,
 ) -> C: ...
@@ -51,6 +53,7 @@ def register(
     fn: None = None,
     /,
     *,
+    exposed: tuple[str, ...] | None = None,
     proxytype: type | None = None,
     method_to_typeid: dict[str, str] | None = None,
 ) -> PassthroughDecorator[C]: ...
@@ -60,10 +63,13 @@ def register(
     fn: C | None = None,
     /,
     *,
+    exposed: tuple[str, ...] | None = None,
     proxytype: type | None = None,
     method_to_typeid: dict[str, str] | None = None,
 ) -> PassthroughDecorator[C] | C:
     kwargs = {}
+    if exposed is not None:
+        kwargs["exposed"] = exposed
     if proxytype is not None:
         kwargs["proxytype"] = proxytype
     if method_to_typeid is not None:
@@ -83,20 +89,20 @@ class FuturesProxy(DictProxy):
     }
 
 
-_task_queue: TaskQueue[WoolTask] = TaskQueue(1000, None)
+_task_queue: TaskQueue[Task] = TaskQueue(100000, None)
 
 _task_queue_lock = Lock()
 
-_task_futures: WeakValueDictionary[UUID, WoolFuture] = WeakValueDictionary()
+_task_futures: WeakValueDictionary[UUID, Future] = WeakValueDictionary()
 
 
 @register
-def put(task: WoolTask) -> WoolFuture:
+def put(task: Task) -> Future:
     try:
         with queue_lock():
             queue().put(task, block=False)
-            future = futures()[task.id] = WoolFuture()
-            logging.debug(f"Pushed task {task.id} to queue: {task.tag}")
+            future = futures()[task.id] = Future()
+            TaskEvent("task-queued", task=task).emit()
             return future
     except Exception as e:
         logging.exception(e)
@@ -104,7 +110,7 @@ def put(task: WoolTask) -> WoolFuture:
 
 
 @register
-def get() -> WoolTask | Empty | None:
+def get() -> Task | Empty | None:
     try:
         return queue().get(block=False)
     except Empty as e:
@@ -140,9 +146,18 @@ def queue_lock() -> Lock:
     return _task_queue_lock
 
 
-@register
-def stop(wait: bool = True) -> None:
-    os.kill(os.getpid(), signal.SIGINT if wait else signal.SIGTERM)
+_stop_event = Event()
+_wait_event = Event()
+
+
+@register(exposed=("is_set", "set", "clear", "wait"))
+def stopping() -> Event:
+    return _stop_event
+
+
+@register(exposed=("is_set", "set", "clear", "wait"))
+def waiting() -> Event:
+    return _wait_event
 
 
 class ManagerMeta(type):
@@ -162,4 +177,5 @@ class Manager(BaseManager, metaclass=ManagerMeta):
         futures = staticmethod(futures)
         queue = staticmethod(queue)
         queue_lock = staticmethod(queue_lock)
-        stop = staticmethod(stop)
+        stopping = staticmethod(stopping)
+        waiting = staticmethod(waiting)

wool/_mempool/__init__.py

@@ -0,0 +1,3 @@
+from wool._mempool._mempool import MemoryPool
+
+__all__ = ["MemoryPool"]

wool/_mempool/_mempool.py

@@ -0,0 +1,204 @@
+import asyncio
+import hashlib
+import mmap
+import os
+import pathlib
+import shutil
+from collections import namedtuple
+from contextlib import asynccontextmanager
+from dataclasses import asdict
+from dataclasses import fields
+from types import MappingProxyType
+
+import shortuuid
+
+from wool._mempool._metadata import MetadataMessage
+
+Metadata = namedtuple(
+    "Metadata",
+    [field.name for field in fields(MetadataMessage)],
+)
+
+
+class MetadataMapping:
+    def __init__(self, mapping) -> None:
+        self._mapping = MappingProxyType(mapping)
+
+    def __getitem__(self, key: str) -> Metadata:
+        return Metadata(**asdict(self._mapping[key]))
+
+
+class MemoryPool:
+    def __init__(self, path: str | pathlib.Path = pathlib.Path(".mempool")):
+        if isinstance(path, str):
+            self._path = pathlib.Path(path)
+        else:
+            self._path = path
+        self._lockdir = self._path / "locks"
+        os.makedirs(self._lockdir, exist_ok=True)
+        self._files = {}
+        self._mmaps = {}
+        self._metadata: dict[str, MetadataMessage] = {}
+
+    @property
+    def metadata(self) -> MetadataMapping:
+        return MetadataMapping(self._metadata)
+
+    @property
+    def path(self) -> pathlib.Path:
+        return self._path
+
+    async def map(self):
+        for entry in os.scandir(self._path):
+            if entry.is_dir() and (ref := entry.name) != "locks":
+                async with self._reflock(ref):
+                    self._map(ref)
+
+    async def put(
+        self, dump: bytes, *, mutable: bool = False, ref: str | None = None
+    ) -> str:
+        async with self._reflock(ref := ref or str(shortuuid.uuid())):
+            self._put(ref, dump, mutable=mutable, exist_ok=False)
+            return ref
+
+    async def post(self, ref: str, dump: bytes):
+        async with self._reflock(ref):
+            if ref not in self._mmaps:
+                self._map(ref)
+            metamap, dumpmap = self._mmaps[ref]
+            metamap.seek(0)
+            metadata = self._metadata.setdefault(
+                ref, MetadataMessage.loads(metamap.read())
+            )
+            if not metadata.mutable:
+                raise ValueError("Cannot modify an immutable reference")
+            if (size := len(dump)) != metadata.size:
+                try:
+                    dumpmap.resize(size)
+                    self._post(ref, dump, metadata)
+                except SystemError:
+                    self._put(ref, dump, mutable=True, exist_ok=True)
+                return True
+            elif hashlib.md5(dump).digest() != metadata.md5:
+                self._post(ref, dump, metadata)
+                return True
+            else:
+                return False
+
+    async def get(self, ref: str) -> bytes:
+        async with self._reflock(ref):
+            if ref not in self._mmaps:
+                self._map(ref)
+            metamap, dumpmap = self._mmaps[ref]
+            metamap.seek(0)
+            metadata = MetadataMessage.loads(metamap.read())
+            cached_metadata = self._metadata.setdefault(ref, metadata)
+            if metadata.mutable and metadata.size != cached_metadata.size:
+                # Dump size has changed, so we need to re-map it
+                self._map(ref)
+            dumpmap.seek(0)
+            return dumpmap.read()
+
+    async def delete(self, ref: str):
+        async with self._reflock(ref):
+            if ref not in self._mmaps:
+                self._map(ref)
+            metamap, dumpmap = self._mmaps.pop(ref)
+            metafile, dumpfile = self._files.pop(ref)
+            if ref in self._metadata:
+                del self._metadata[ref]
+            if metamap:
+                metamap.close()
+            if dumpmap:
+                dumpmap.close()
+            if metafile:
+                metafile.close()
+            if dumpfile:
+                dumpfile.close()
+            try:
+                shutil.rmtree(self.path / ref)
+            except FileNotFoundError:
+                pass
+
+    def _put(
+        self,
+        ref: str,
+        dump: bytes,
+        *,
+        mutable: bool = False,
+        exist_ok: bool = False,
+    ):
+        metadata = MetadataMessage(
+            ref=ref,
+            mutable=mutable,
+            size=len(dump),
+            md5=hashlib.md5(dump).digest(),
+        )
+
+        refpath = pathlib.Path(self._path, f"{ref}")
+        os.makedirs(refpath, exist_ok=exist_ok)
+
+        with open(refpath / "meta", "wb") as metafile:
+            metafile.write(metadata.dumps())
+
+        with open(refpath / "dump", "wb") as dumpfile:
+            dumpfile.write(dump)
+
+        self._map(ref)
+        self._metadata[ref] = metadata
+
+    def _post(self, ref: str, dump: bytes, metadata: MetadataMessage):
+        metamap, dumpmap = self._mmaps[ref]
+        metadata.size = len(dump)
+        metadata.md5 = hashlib.md5(dump).digest()
+        metamap.seek(0)
+        metamap.write(metadata.dumps())
+        metamap.flush()
+        dumpmap.seek(0)
+        dumpmap.write(dump)
+        dumpmap.flush()
+
+    def _map(self, ref: str):
+        refpath = pathlib.Path(self._path, f"{ref}")
+        metafile = open(refpath / "meta", "r+b")
+        metamap = mmap.mmap(metafile.fileno(), 0, flags=mmap.MAP_SHARED)
+        dumpfile = open(refpath / "dump", "r+b")
+        dumpmap = mmap.mmap(dumpfile.fileno(), 0, flags=mmap.MAP_SHARED)
+        cached_metafile, cached_dumpfile = self._files.pop(ref, (None, None))
+        if cached_metafile:
+            cached_metafile.close()
+        if cached_dumpfile:
+            cached_dumpfile.close()
+        cached_metamap, cached_dumpmap = self._mmaps.pop(ref, (None, None))
+        if cached_metamap is not None and not cached_metamap.closed:
+            cached_metamap.close()
+        if cached_dumpmap is not None and not cached_dumpmap.closed:
+            cached_dumpmap.close()
+        self._files[ref] = (metafile, dumpfile)
+        self._mmaps[ref] = (metamap, dumpmap)
+
+    def _lockpath(self, ref: str):
+        return pathlib.Path(self._lockdir, f"{ref}.lock")
+
+    def _acquire(self, ref: str):
+        try:
+            os.symlink(f"{ref}", self._lockpath(ref))
+            return True
+        except FileExistsError:
+            return False
+
+    def _release(self, ref: str):
+        try:
+            if os.path.islink(lock_path := self._lockpath(ref)):
+                os.unlink(lock_path)
+        except FileNotFoundError:
+            pass
+
+    @asynccontextmanager
+    async def _reflock(self, ref: str):
+        try:
+            while not self._acquire(ref):
+                await asyncio.sleep(0)
+            yield
+        finally:
+            self._release(ref)

wool/_mempool/_metadata/__init__.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+try:
+    from wool._protobuf._mempool._metadata._metadata_pb2 import (
+        _MetadataMessage,
+    )
+except ImportError:
+    logging.error(
+        "Failed to import _MetadataMessage. "
+        "Ensure protocol buffers are compiled."
+    )
+    raise
+
+
+@dataclass
+class MetadataMessage:
+    ref: str
+    mutable: bool
+    size: int
+    md5: bytes
+
+    @classmethod
+    def loads(cls, data: bytes) -> MetadataMessage:
+        (metadata := _MetadataMessage()).ParseFromString(data)
+        return cls(
+            ref=metadata.ref,
+            mutable=metadata.mutable,
+            size=metadata.size,
+            md5=metadata.md5,
+        )
+
+    def dumps(self) -> bytes:
+        return _MetadataMessage(
+            ref=self.ref, mutable=self.mutable, size=self.size, md5=self.md5
+        ).SerializeToString()
+
+
+__all__ = ["MetadataMessage"]