add a pubsub database implementation

Author: Liam-DeVoe

hypothesis-python/RELEASE.rst

@@ -0,0 +1,5 @@
+RELEASE_TYPE: minor
+
+The :doc:`Hypothesis database <database>` now supports a pub-sub interface to efficiently listen for changes in the database, via ``.add_listener`` and ``.remove_listener``. While all databases that ship with Hypothesis support this interface, implementing it is not required for custom database subclasses. Hypothesis will warn when trying to listen on a database without support.
+
+This feature is currently only used downstream in `hypofuzz <https://github.com/zac-hd/hypofuzz>`_.

hypothesis-python/setup.py

@@ -70,6 +70,7 @@ def local_file(name):
     # We also leave the choice of timezone library to the user, since it
     # might be zoneinfo or pytz depending on version and configuration.
     "django": ["django>=4.2"],
+    "watchdog": ["watchdog>=4.0.0"],
 }
 
 extras["all"] = sorted(set(sum(extras.values(), [])))

hypothesis-python/src/hypothesis/database.py

@@ -24,7 +24,7 @@
 from pathlib import Path, PurePath
 from queue import Queue
 from threading import Thread
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, cast
 from urllib.error import HTTPError, URLError
 from urllib.request import Request, urlopen
 from zipfile import BadZipFile, ZipFile
@@ -47,6 +47,7 @@
     from typing import TypeAlias
 
 StrPathT: "TypeAlias" = Union[str, PathLike[str]]
+ListenerT: "TypeAlias" = Callable[[], Any]
 
 
 def _usable_dir(path: StrPathT) -> bool:
@@ -121,9 +122,13 @@ class ExampleDatabase(metaclass=_EDMeta):
     """An abstract base class for storing examples in Hypothesis' internal format.
 
     An ExampleDatabase maps each ``bytes`` key to many distinct ``bytes``
-    values, like a ``Mapping[bytes, AbstractSet[bytes]]``.
+    values, like a ``Mapping[bytes, set[bytes]]``.
     """
 
+    def __init__(self) -> None:
+        self._listeners: list[ListenerT] = []
+        self._listening = False
+
     @abc.abstractmethod
     def save(self, key: bytes, value: bytes) -> None:
         """Save ``value`` under ``key``.
@@ -159,6 +164,77 @@ def move(self, src: bytes, dest: bytes, value: bytes) -> None:
         self.delete(src, value)
         self.save(dest, value)
 
+    def add_listener(self, f: ListenerT, /) -> None:
+        """Add a change listener."""
+        self._listeners.append(f)
+        self._update_listening()
+
+    def remove_listener(self, f: ListenerT, /) -> None:
+        """
+        Remove a change listener. If the listener is not present, silently do
+        nothing.
+        """
+        self._listeners.remove(f)
+        self._update_listening()
+
+    def clear_listeners(self) -> None:
+        """Remove all change listeners."""
+        self._listeners.clear()
+        self._update_listening()
+
+    def _update_listening(self) -> None:
+        # - start listening if we're moving from zero to some listeners
+        # - stop listening if we're moving from some to zero listeners
+        if not self._listening and self._listeners:
+            self._start_listening()
+            self._listening = True
+        elif self._listening and not self._listeners:
+            self._stop_listening()
+            self._listening = False
+
+    def _broadcast_changed(self) -> None:
+        """
+        Called when a change has been made to the database that would cause
+        .fetch to return something different than it did before. If your database
+        implementation supports change listening, this method should be called
+        whenever an item is added to or deleted from the underlying database store.
+
+        Note that you should not assume you are the only reference to the underlying
+        database store. For example, if two DirectoryBasedExampleDatabase reference
+        the same directory, _broadcast_changed should be called whenever a file is
+        added or removed from the directory, even if that database was not responsible
+        for changing the file.
+        """
+        for listener in self._listeners:
+            listener()
+
+    def _start_listening(self) -> None:
+        """
+        Called when the database adds a change listener, and did not previously
+        have any change listeners. Intended to allow databases to wait to start
+        expensive listening operations until necessary.
+
+        _start_listening and _stop_listening are guaranteed to alternate, so you
+        do not need to handle the case of multiple consecutive _start_listening
+        calls without an intermediate _stop_listening call.
+        """
+        warnings.warn(
+            f"{self.__class__} does not support listening for changes", stacklevel=4
+        )
+
+    def _stop_listening(self) -> None:
+        """
+        Called whenever no change listeners remain on the database.
+
+        _stop_listening and _start_listening are guaranteed to alternate, so you
+        do not need to handle the case of multiple consecutive _stop_listening
+        calls without an intermediate _start_listening call.
+        """
+        warnings.warn(
+            f"{self.__class__} does not support stopping listening for changes",
+            stacklevel=4,
+        )
+
 
 class InMemoryExampleDatabase(ExampleDatabase):
     """A non-persistent example database, implemented in terms of a dict of sets.
@@ -169,6 +245,7 @@ class InMemoryExampleDatabase(ExampleDatabase):
     """
 
     def __init__(self) -> None:
+        super().__init__()
         self.data: dict[bytes, set[bytes]] = {}
 
     def __repr__(self) -> str:
@@ -178,10 +255,31 @@ def fetch(self, key: bytes) -> Iterable[bytes]:
         yield from self.data.get(key, ())
 
     def save(self, key: bytes, value: bytes) -> None:
-        self.data.setdefault(key, set()).add(bytes(value))
+        value = bytes(value)
+        values = self.data.setdefault(key, set())
+        changed = value not in values
+        values.add(value)
+
+        if changed:
+            self._broadcast_changed()
 
     def delete(self, key: bytes, value: bytes) -> None:
-        self.data.get(key, set()).discard(bytes(value))
+        value = bytes(value)
+        values = self.data.get(key, set())
+        changed = value in values
+        values.discard(value)
+
+        if changed:
+            self._broadcast_changed()
+
+    def _start_listening(self) -> None:
+        # declare compatibility with the listener api, but do the actual
+        # implementation in .delete and .save, since we know we are the only
+        # writer to .data.
+        pass
+
+    def _stop_listening(self) -> None:
+        pass
 
 
 def _hash(key: bytes) -> str:
@@ -208,8 +306,12 @@ class DirectoryBasedExampleDatabase(ExampleDatabase):
     """
 
     def __init__(self, path: StrPathT) -> None:
+        super().__init__()
         self.path = Path(path)
         self.keypaths: dict[bytes, Path] = {}
+        # would type as watchdog.observers.Observer | None, but don't want to
+        # mess with a conditional import just for type checking.
+        self._observer: Any = None
 
     def __repr__(self) -> str:
         return f"DirectoryBasedExampleDatabase({self.path!r})"
@@ -278,6 +380,44 @@ def delete(self, key: bytes, value: bytes) -> None:
         except OSError:
             pass
 
+    def _start_listening(self) -> None:
+        try:
+            from watchdog.events import (
+                FileCreatedEvent,
+                FileDeletedEvent,
+                FileSystemEvent,
+                FileSystemEventHandler,
+            )
+            from watchdog.observers import Observer
+        except ImportError:
+            warnings.warn(
+                f"listening for changes in a {self.__class__.__name__} "
+                "requires the watchdog library. To install, run "
+                "`pip install hypothesis[watchdog]`",
+                stacklevel=4,
+            )
+            return
+
+        _broadcast_changed = self._broadcast_changed
+
+        class Handler(FileSystemEventHandler):
+            def on_any_event(self, event: FileSystemEvent) -> None:
+                _broadcast_changed()
+
+        self._observer = Observer()
+        self._observer.schedule(
+            Handler(),
+            self.path,  # type: ignore # upstream type is too narrow (str only)
+            recursive=True,
+            event_filter=[FileCreatedEvent, FileDeletedEvent],
+        )
+        self._observer.start()
+
+    def _stop_listening(self) -> None:
+        self._observer.stop()
+        self._observer.join()
+        self._observer = None
+
 
 class ReadOnlyDatabase(ExampleDatabase):
     """A wrapper to make the given database read-only.
@@ -291,6 +431,7 @@ class ReadOnlyDatabase(ExampleDatabase):
     """
 
     def __init__(self, db: ExampleDatabase) -> None:
+        super().__init__()
         assert isinstance(db, ExampleDatabase)
         self._wrapped = db
 
@@ -306,6 +447,13 @@ def save(self, key: bytes, value: bytes) -> None:
     def delete(self, key: bytes, value: bytes) -> None:
         pass
 
+    def _start_listening(self) -> None:
+        # we're read only, so there are no changes to broadcast.
+        pass
+
+    def _stop_listening(self) -> None:
+        pass
+
 
 class MultiplexedDatabase(ExampleDatabase):
     """A wrapper around multiple databases.
@@ -334,6 +482,7 @@ class MultiplexedDatabase(ExampleDatabase):
     """
 
     def __init__(self, *dbs: ExampleDatabase) -> None:
+        super().__init__()
         assert all(isinstance(db, ExampleDatabase) for db in dbs)
         self._wrapped = dbs
 
@@ -360,6 +509,14 @@ def move(self, src: bytes, dest: bytes, value: bytes) -> None:
         for db in self._wrapped:
             db.move(src, dest, value)
 
+    def _start_listening(self) -> None:
+        for db in self._wrapped:
+            db.add_listener(self._broadcast_changed)
+
+    def _stop_listening(self) -> None:
+        for db in self._wrapped:
+            db.remove_listener(self._broadcast_changed)
+
 
 class GitHubArtifactDatabase(ExampleDatabase):
     """
@@ -439,6 +596,7 @@ def __init__(
         cache_timeout: timedelta = timedelta(days=1),
         path: Optional[StrPathT] = None,
     ):
+        super().__init__()
         self.owner = owner
         self.repo = repo
         self.artifact_name = artifact_name
@@ -699,6 +857,7 @@ class BackgroundWriteDatabase(ExampleDatabase):
     """
 
     def __init__(self, db: ExampleDatabase) -> None:
+        super().__init__()
         self._db = db
         self._queue: Queue[tuple[str, tuple[bytes, ...]]] = Queue()
         self._thread = Thread(target=self._worker, daemon=True)
@@ -735,6 +894,12 @@ def delete(self, key: bytes, value: bytes) -> None:
     def move(self, src: bytes, dest: bytes, value: bytes) -> None:
         self._queue.put(("move", (src, dest, value)))
 
+    def _start_listening(self) -> None:
+        self._db.add_listener(self._broadcast_changed)
+
+    def _stop_listening(self) -> None:
+        self._db.remove_listener(self._broadcast_changed)
+
 
 def _pack_uleb128(value: int) -> bytes:
     """

hypothesis-python/src/hypothesis/extra/redis.py

@@ -11,6 +11,7 @@
 from collections.abc import Iterable
 from contextlib import contextmanager
 from datetime import timedelta
+from typing import Any
 
 from redis import Redis
 
@@ -36,31 +37,42 @@ def __init__(
         *,
         expire_after: timedelta = timedelta(days=8),
         key_prefix: bytes = b"hypothesis-example:",
+        listener_channel: str = "hypothesis-changes",
     ):
+        super().__init__()
         check_type(Redis, redis, "redis")
         check_type(timedelta, expire_after, "expire_after")
         check_type(bytes, key_prefix, "key_prefix")
+        check_type(str, listener_channel, "listener_channel")
         self.redis = redis
         self._expire_after = expire_after
         self._prefix = key_prefix
+        self.listener_channel = listener_channel
+        self._pubsub: Any = None
 
     def __repr__(self) -> str:
         return (
             f"RedisExampleDatabase({self.redis!r}, expire_after={self._expire_after!r})"
         )
 
     @contextmanager
-    def _pipeline(self, *reset_expire_keys, transaction=False, auto_execute=True):
+    def _pipeline(self, *reset_expire_keys, execute_and_publish=True):
         # Context manager to batch updates and expiry reset, reducing TCP roundtrips
-        pipe = self.redis.pipeline(transaction=transaction)
+        pipe = self.redis.pipeline()
         yield pipe
         for key in reset_expire_keys:
             pipe.expire(self._prefix + key, self._expire_after)
-        if auto_execute:
-            pipe.execute()
+        if execute_and_publish:
+            # pipe.execute returns a value for each operation, which includes
+            # whatever we did in the yield as a prefix, and the n operations from
+            # pipe.expire as a suffix. remove that suffix to get just the prefix.
+            values = pipe.execute()[: -len(reset_expire_keys)]
+            # only publish if anything changed
+            if any(value > 0 for value in values):
+                self.redis.publish(self.listener_channel, "change")
 
     def fetch(self, key: bytes) -> Iterable[bytes]:
-        with self._pipeline(key, auto_execute=False) as pipe:
+        with self._pipeline(key, execute_and_publish=False) as pipe:
             pipe.smembers(self._prefix + key)
         yield from pipe.execute()[0]
 
@@ -76,3 +88,19 @@ def move(self, src: bytes, dest: bytes, value: bytes) -> None:
         with self._pipeline(src, dest) as pipe:
             pipe.srem(self._prefix + src, value)
             pipe.sadd(self._prefix + dest, value)
+
+    def _handle_message(self, message: dict) -> None:
+        # other message types include "subscribe" and "unsubscribe". these are
+        # sent to the client, but not to the pubsub channel.
+        assert message["type"] == "message"
+        self._broadcast_changed()
+
+    def _start_listening(self) -> None:
+        self._pubsub = self.redis.pubsub()
+        self._pubsub.subscribe(**{self.listener_channel: self._handle_message})
+
+    def _stop_listening(self) -> None:
+        if self._pubsub is not None:
+            self._pubsub.unsubscribe()
+            self._pubsub.close()
+            self._pubsub = None