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

wool/_mempool/__init__.py

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

wool/_mempool/_client.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Final
+from typing import Protocol
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+import grpc
+
+try:
+    from wool._protobuf.mempool import service_pb2 as pb
+    from wool._protobuf.mempool import service_pb2_grpc as rpc
+except ImportError as e:
+    from wool._protobuf import ProtobufImportError
+
+    raise ProtobufImportError(e) from e
+
+
+class EventHandler(Protocol):
+    @staticmethod
+    def __call__(client: MemoryPoolClient, response: pb.Event) -> None: ...
+
+
+class MemoryPoolClient:
+    _event_stream: Final[grpc.aio.UnaryStreamCall]
+    _listener: asyncio.Task
+    _session: pb.Session | None
+
+    def __init__(self, channel: grpc.aio.Channel, event_handler: EventHandler):
+        self._stub = rpc.MemoryPoolStub(channel)
+        self._event_stream = self._stub.session(pb.SessionRequest())
+        self._event_handler = event_handler
+        self._session = None
+
+    async def __aenter__(self) -> Self:
+        if self._session:
+            raise RuntimeError(
+                "Client session may not be re-entered. "
+                "Use 'async with MemoryPoolClient(...)' only once."
+            )
+        response = await self._event_stream.read()
+        assert isinstance(response, pb.SessionResponse)
+        self._session = response.session
+        self._listener = asyncio.create_task(self._listen())
+        return self
+
+    async def __aexit__(self, *_):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        assert self._listener
+        self._listener.cancel()
+        try:
+            await self._listener
+        except asyncio.CancelledError:
+            pass
+
+    def __del__(self):
+        try:
+            self._event_stream.cancel()
+        except Exception:
+            pass
+
+    @property
+    def id(self) -> str:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        return self._session.id
+
+    async def map(self, ref: str):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.AcquireRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        await self._stub.map(request)
+
+    async def get(self, ref: str) -> bytes:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.GetRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        response: pb.GetResponse = await self._stub.get(request)
+        return response.dump
+
+    async def put(self, dump: bytes, *, mutable: bool = False) -> str:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.PutRequest(
+            session=self._session,
+            dump=dump,
+            mutable=mutable,
+        )
+        response: pb.PutResponse = await self._stub.put(request)
+        return response.reference.id
+
+    async def post(self, ref: str, dump: bytes) -> bool:
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.PostRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+            dump=dump,
+        )
+        response: pb.PostResponse = await self._stub.post(request)
+        return response.updated
+
+    async def acquire(self, ref: str):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.AcquireRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        await self._stub.acquire(request)
+
+    async def release(self, ref: str):
+        if not self._session:
+            raise RuntimeError(
+                "Client session has not been entered. "
+                "Use 'async with MemoryPoolClient(...)'."
+            )
+        request = pb.ReleaseRequest(
+            session=self._session,
+            reference=pb.Reference(id=ref),
+        )
+        await self._stub.release(request)
+
+    async def _listen(self):
+        assert self._event_stream
+        try:
+            while True:
+                if (response := await self._event_stream.read()) is None:
+                    break
+                assert isinstance(response, pb.SessionResponse), (
+                    f"Unexpected event type: {type(response)}"
+                )
+                self._event_handler(self, response.event)
+        except asyncio.CancelledError:
+            self._event_stream.cancel()

wool/_mempool/_mempool.py

@@ -1,34 +1,141 @@
+from __future__ import annotations
+
 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
+from mmap import mmap
+from typing import BinaryIO
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
 
 import shortuuid
 
 from wool._mempool._metadata import MetadataMessage
 
-Metadata = namedtuple(
-    "Metadata",
-    [field.name for field in fields(MetadataMessage)],
-)
 
+class SharedObject:
+    _id: str
+    _mempool: MemoryPool
+    _file: BinaryIO
+    _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(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:
+        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(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
+    _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(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:
+        return self._mmap
+
+    @property
+    def _path(self) -> pathlib.Path:
+        return pathlib.Path(self._mempool.path, self.id)
 
-class MetadataMapping:
-    def __init__(self, mapping) -> None:
-        self._mapping = MappingProxyType(mapping)
+    @property
+    def _metadata(self) -> MetadataMessage:
+        return MetadataMessage.loads(bytes(self._mmap))
 
-    def __getitem__(self, key: str) -> Metadata:
-        return Metadata(**asdict(self._mapping[key]))
+    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)
@@ -36,89 +143,86 @@ class MemoryPool:
             self._path = path
         self._lockdir = self._path / "locks"
         os.makedirs(self._lockdir, exist_ok=True)
-        self._files = {}
-        self._mmaps = {}
-        self._metadata: dict[str, MetadataMessage] = {}
+        self._acquire(f"pid-{os.getpid()}")
+        self._objects = dict()
 
-    @property
-    def metadata(self) -> MetadataMapping:
-        return MetadataMapping(self._metadata)
+    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):
-        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 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:
-        async with self._reflock(ref := ref or str(shortuuid.uuid())):
+        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):
-        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())
+    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"
             )
-            if not metadata.mutable:
+        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)) != metadata.size:
+            if (size := len(dump)) != obj.metadata.size:
                 try:
-                    dumpmap.resize(size)
-                    self._post(ref, dump, metadata)
+                    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() != metadata.md5:
-                self._post(ref, dump, metadata)
+            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:
-        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
+        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)
-            dumpmap.seek(0)
-            return dumpmap.read()
+            return bytes(self._objects[ref].refresh().mmap)
 
     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
+        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,
@@ -145,60 +249,63 @@ class MemoryPool:
             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 _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):
-        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):
+        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"{ref}", self._lockpath(ref))
+            os.symlink(f"{key}", self._lockpath(key))
             return True
         except FileExistsError:
             return False
 
-    def _release(self, ref: str):
+    def _release(self, key: str):
         try:
-            if os.path.islink(lock_path := self._lockpath(ref)):
+            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 _reflock(self, ref: str):
+    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/__init__.py → _metadata.py}

@@ -1,18 +1,14 @@
 from __future__ import annotations
 
-import logging
+from dataclasses import asdict
 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
+    from wool._protobuf.mempool import metadata_pb2 as pb
+except ImportError as e:
+    from wool._protobuf import ProtobufImportError
+
+    raise ProtobufImportError(e) from e
 
 
 @dataclass
@@ -24,7 +20,7 @@ class MetadataMessage:
 
     @classmethod
     def loads(cls, data: bytes) -> MetadataMessage:
-        (metadata := _MetadataMessage()).ParseFromString(data)
+        (metadata := pb.MetadataMessage()).ParseFromString(data)
         return cls(
             ref=metadata.ref,
             mutable=metadata.mutable,
@@ -33,9 +29,7 @@ class MetadataMessage:
         )
 
     def dumps(self) -> bytes:
-        return _MetadataMessage(
-            ref=self.ref, mutable=self.mutable, size=self.size, md5=self.md5
-        ).SerializeToString()
+        return pb.MetadataMessage(**asdict(self)).SerializeToString()
 
 
 __all__ = ["MetadataMessage"]