Reapply "ref(producer): Simplify typing, add more tests (#441)" (#448)

This reverts commit 2713592.

Author: untitaker

.github/workflows/docs-pr.yml

@@ -1,6 +1,7 @@
 name: Arroyo Docs
 
 on:
+  pull_request:
   push:
     branches:
       - main
@@ -14,13 +15,14 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v2
         with:
-          python-version: '3.8'
+          python-version: '3.13'
 
       - name: Build docs
         run: |
           pip install virtualenv
           make docs
       - uses: peaceiris/[email protected]
+        if: ${{ (github.ref_name == 'main') }}
         name: Publish to GitHub Pages
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/docs.yml

@@ -4,16 +4,7 @@
 import time
 from abc import ABC, abstractmethod, abstractproperty
 from concurrent.futures import Future
-from typing import (
-    Callable,
-    Generic,
-    Mapping,
-    Optional,
-    Protocol,
-    Sequence,
-    TypeVar,
-    Union,
-)
+from typing import Callable, Generic, Mapping, Optional, Sequence, TypeVar, Union
 
 from arroyo.types import BrokerValue, Partition, Topic, TStrategyPayload
 
@@ -187,32 +178,11 @@ def member_id(self) -> str:
         raise NotImplementedError
 
 
-class ProducerFuture(Protocol, Generic[T]):
-    """
-    An abstract interface for a kind of Future. Stdlib futures are too slow to
-    construct, so we use these.
-    """
-
-    def done(self) -> bool:
-        ...
-
-    def result(self, timeout: float | None = None) -> T:
-        """
-        Return result or raise exception. May block, but does not have to.
-        """
-        ...
-
-    def set_result(self, result: T) -> None:
-        ...
-
-    def set_exception(self, exception: Exception) -> None:
-        ...
-
-
 class SimpleProducerFuture(Generic[T]):
     """
     A stub for concurrent.futures.Future that does not construct any Condition
-    variables, therefore is faster to construct.
+    variables, therefore is faster to construct. However, some methods are
+    missing, and result() in particular is not efficient with timeout > 0.
     """
 
     def __init__(self) -> None:
@@ -232,6 +202,10 @@ def result(self, timeout: float | None = None) -> T:
         # only in tests at most. It is only here for the sake of implementing
         # the contract. If you really need result with timeout>0, you should
         # use the stdlib future.
+        #
+        # If this becomes performance sensitive, we can potentially implement
+        # something more sophisticated such as lazily creating the condition
+        # variable, and synchronizing the creation of that using a global lock.
         while deadline is None or time.time() < deadline:
             if self.result_exception is not None:
                 raise self.result_exception
@@ -248,6 +222,9 @@ def set_exception(self, exception: Exception) -> None:
         self.result_exception = exception
 
 
+ProducerFuture = Union[SimpleProducerFuture[T], Future[T]]
+
+
 class Producer(Generic[TStrategyPayload], ABC):
     @abstractmethod
     def produce(

arroyo/backends/abstract.py

@@ -167,6 +167,15 @@ def test_auto_offset_reset_latest(self, use_simple_futures: bool) -> None:
                 else:
                     raise AssertionError("expected EndOfPartition error")
 
+    @pytest.mark.parametrize("use_simple_futures", [True, False])
+    def test_producer_future_behavior(self, use_simple_futures: bool) -> None:
+        with self.get_topic() as topic:
+            with closing(self.get_producer(use_simple_futures)) as producer:
+                future = producer.produce(topic, next(self.get_payloads()))
+                assert not future.done()
+                assert future.result(5.0)
+                assert future.done()
+
     def test_lenient_offset_reset_latest(self) -> None:
         payload = KafkaPayload(b"a", b"0", [])
         with self.get_topic() as topic: