wool 0.1rc6__py3-none-any.whl → 0.1rc8__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,4 @@
+from wool._mempool._mempool import MemoryPool
+from wool._mempool._service import MemoryPoolService
+
+__all__ = ["MemoryPool", "MemoryPoolService"]

wool/_mempool/_mempool.py

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import mmap
+import os
+import pathlib
+import shutil
+from contextlib import asynccontextmanager
+from typing import BinaryIO
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+import shortuuid
+
+from wool._mempool._metadata import MetadataMessage
+
+
+class SharedObject:
+    _id: str
+    _mempool: MemoryPool
+    _file: BinaryIO
+    _mmap: mmap.mmap
+    _size: int
+    _md5: bytes
+
+    def __init__(self, id: str, *, mempool: MemoryPool):
+        self._id = id
+        self._mempool = mempool
+        self._file = open(self._path / "dump", "r+b")
+        self._mmap = mmap.mmap(self._file.fileno(), 0)
+        self._size = self.metadata.size
+        self._md5 = self.metadata.md5
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def metadata(self) -> SharedObjectMetadata:
+        return SharedObjectMetadata(self.id, mempool=self._mempool)
+
+    @property
+    def mmap(self) -> mmap.mmap:
+        return self._mmap
+
+    @property
+    def _path(self) -> pathlib.Path:
+        return pathlib.Path(self._mempool.path, self.id)
+
+    def close(self):
+        self.metadata.close()
+        self._mmap.close()
+        self._file.close()
+
+    def refresh(self) -> Self:
+        if self._size != self.metadata.size or self._md5 != self.metadata.md5:
+            self._mmap.close()
+            self._file.close()
+            self._file = open(self._path / "dump", "r+b")
+            self._mmap = mmap.mmap(self._file.fileno(), 0)
+            self._size = self.metadata.size
+            self._md5 = self.metadata.md5
+        return self
+
+
+class SharedObjectMetadata:
+    _id: str
+    _mempool: MemoryPool
+    _file: BinaryIO
+    _mmap: mmap.mmap
+    _instances: dict[str, SharedObjectMetadata] = {}
+
+    def __new__(cls, id: str, *, mempool: MemoryPool):
+        if id in cls._instances:
+            return cls._instances[id]
+        return super().__new__(cls)
+
+    def __init__(self, id: str, mempool: MemoryPool):
+        self._id = id
+        self._mempool = mempool
+        self._file = open(self._path / "meta", "r+b")
+        self._mmap = mmap.mmap(self._file.fileno(), 0)
+        self._instances[id] = self
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def mutable(self) -> bool:
+        return self._metadata.mutable
+
+    @property
+    def size(self) -> int:
+        return self._metadata.size
+
+    @property
+    def md5(self) -> bytes:
+        return self._metadata.md5
+
+    @property
+    def mmap(self) -> mmap.mmap:
+        return self._mmap
+
+    @property
+    def _path(self) -> pathlib.Path:
+        return pathlib.Path(self._mempool.path, self.id)
+
+    @property
+    def _metadata(self) -> MetadataMessage:
+        return MetadataMessage.loads(bytes(self._mmap))
+
+    def close(self):
+        self._mmap.close()
+        self._file.close()
+        del self._instances[self.id]
+
+
+class MemoryPool:
+    _objects: dict[str, SharedObject]
+    _path: pathlib.Path
+
+    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._acquire(f"pid-{os.getpid()}")
+        self._objects = dict()
+
+    def __contains__(self, ref: str) -> bool:
+        return ref in self._objects
+
+    def __del__(self):
+        self._release(f"pid-{os.getpid()}")
+
+    @property
+    def path(self) -> pathlib.Path:
+        return self._path
+
+    async def map(self, ref: str | None = None):
+        if ref is not None and ref not in self._objects:
+            if self._locked(f"delete-{ref}"):
+                raise RuntimeError(
+                    f"Reference {ref} is currently locked for deletion"
+                )
+            async with self._reference_lock(ref):
+                self._map(ref)
+        else:
+            for entry in os.scandir(self._path):
+                if entry.is_dir() and (ref := entry.name) != "locks":
+                    if not self._locked(f"delete-{ref}"):
+                        async with self._reference_lock(ref):
+                            self._map(ref)
+
+    async def put(
+        self, dump: bytes, *, mutable: bool = False, ref: str | None = None
+    ) -> str:
+        ref = ref or str(shortuuid.uuid())
+        async with self._reference_lock(ref):
+            self._put(ref, dump, mutable=mutable, exist_ok=False)
+            return ref
+
+    async def post(self, ref: str, dump: bytes) -> bool:
+        if self._locked(f"delete-{ref}"):
+            raise RuntimeError(
+                f"Reference {ref} is currently locked for deletion"
+            )
+        async with self._reference_lock(ref):
+            if ref not in self._objects:
+                self._map(ref)
+            obj = self._objects[ref]
+            if not obj.metadata.mutable:
+                raise ValueError("Cannot modify an immutable reference")
+            if (size := len(dump)) != obj.metadata.size:
+                try:
+                    obj.mmap.resize(size)
+                    self._post(ref, obj, dump)
+                except SystemError:
+                    self._put(ref, dump, mutable=True, exist_ok=True)
+                return True
+            elif hashlib.md5(dump).digest() != obj.metadata.md5:
+                self._post(ref, obj, dump)
+                return True
+            else:
+                return False
+
+    async def get(self, ref: str) -> bytes:
+        if self._locked(f"delete-{ref}"):
+            raise RuntimeError(
+                f"Reference {ref} is currently locked for deletion"
+            )
+        async with self._reference_lock(ref):
+            if ref not in self._objects:
+                self._map(ref)
+            return bytes(self._objects[ref].refresh().mmap)
+
+    async def delete(self, ref: str):
+        async with self._delete_lock(ref):
+            async with self._reference_lock(ref):
+                if ref not in self._objects:
+                    self._map(ref)
+                self._objects.pop(ref).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)
+
+    def _post(self, ref: str, obj: SharedObject, dump: bytes):
+        if not obj.metadata.mutable:
+            raise ValueError("Cannot modify an immutable reference")
+        metadata = MetadataMessage(
+            ref=ref,
+            mutable=True,
+            size=len(dump),
+            md5=hashlib.md5(dump).digest(),
+        )
+        obj.metadata.mmap[:] = metadata.dumps()
+        obj.metadata.mmap.flush()
+        obj.mmap.seek(0)
+        obj.mmap.write(dump)
+
+    def _map(self, ref: str):
+        obj = self._objects.pop(ref, None)
+        if obj:
+            obj.close()
+        self._objects[ref] = SharedObject(id=ref, mempool=self)
+
+    def _lockpath(self, key: str) -> pathlib.Path:
+        return pathlib.Path(self._lockdir, f"{key}.lock")
+
+    def _acquire(self, key: str) -> bool:
+        try:
+            os.symlink(f"{key}", self._lockpath(key))
+            return True
+        except FileExistsError:
+            return False
+
+    def _release(self, key: str):
+        try:
+            if os.path.islink(lock_path := self._lockpath(key)):
+                os.unlink(lock_path)
+        except FileNotFoundError:
+            pass
+
+    def _locked(self, key: str) -> bool:
+        return os.path.islink(self._lockpath(key))
+
+    @asynccontextmanager
+    async def _reference_lock(self, ref: str):
+        try:
+            while not self._acquire(ref):
+                await asyncio.sleep(0)
+            yield
+        finally:
+            self._release(ref)
+
+    @asynccontextmanager
+    async def _delete_lock(self, ref: str):
+        key = f"delete-{ref}"
+        if not self._acquire(f"delete-{ref}"):
+            raise RuntimeError(
+                f"Reference {ref} is currently locked for deletion"
+            )
+        else:
+            yield
+        self._release(key)

wool/_mempool/_metadata.py

@@ -0,0 +1,39 @@
+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"]