class SinglePublisherImpl(SinglePublisher):
_managed_loop: ManagedEventLoop
_underlying: AsyncSinglePublisher
def __init__(self, underlying: AsyncSinglePublisher):
super().__init__()
self._managed_loop = ManagedEventLoop("PublisherLoopThread")
self._underlying = underlying
def publish(self,
data: bytes,
ordering_key: str = "",
**attrs: Mapping[str, str]) -> "Future[str]":
return self._managed_loop.submit(
self._underlying.publish(data=data,
ordering_key=ordering_key,
**attrs))
def __enter__(self):
self._managed_loop.__enter__()
self._managed_loop.submit(self._underlying.__aenter__()).result()
return self
def __exit__(self, __exc_type, __exc_value, __traceback):
self._managed_loop.submit(
self._underlying.__aexit__(__exc_type, __exc_value,
__traceback)).result()
self._managed_loop.__exit__(__exc_type, __exc_value, __traceback)
class SubscriberImpl(ContextManager, StreamingPullManager):
_underlying: AsyncSingleSubscriber
_callback: MessageCallback
_unowned_executor: ThreadPoolExecutor
_event_loop: ManagedEventLoop
_poller_future: concurrent.futures.Future
_close_lock: threading.Lock
_failure: Optional[GoogleAPICallError]
_close_callback: Optional[CloseCallback]
_closed: bool
def __init__(
self,
underlying: AsyncSingleSubscriber,
callback: MessageCallback,
unowned_executor: ThreadPoolExecutor,
):
self._underlying = underlying
self._callback = callback
self._unowned_executor = unowned_executor
self._event_loop = ManagedEventLoop("SubscriberLoopThread")
self._close_lock = threading.Lock()
self._failure = None
self._close_callback = None
self._closed = False
def add_close_callback(self, close_callback: CloseCallback):
"""
A close callback must be set exactly once by the StreamingPullFuture managing this subscriber.
This two-phase init model is made necessary by the requirements of StreamingPullFuture.
"""
with self._close_lock:
assert self._close_callback is None
self._close_callback = close_callback
def close(self):
with self._close_lock:
if self._closed:
return
self._closed = True
self.__exit__(None, None, None)
def _fail(self, error: GoogleAPICallError):
self._failure = error
self.close()
async def _poller(self):
try:
while True:
message = await self._underlying.read()
self._unowned_executor.submit(self._callback, message)
except GoogleAPICallError as e: # noqa: F841 Flake8 thinks e is unused
self._unowned_executor.submit(lambda: self._fail(e)) # noqa: F821
def __enter__(self):
assert self._close_callback is not None
self._event_loop.__enter__()
self._event_loop.submit(self._underlying.__aenter__()).result()
self._poller_future = self._event_loop.submit(self._poller())
return self
def __exit__(self, exc_type, exc_value, traceback):
try:
self._poller_future.cancel()
self._poller_future.result()
except concurrent.futures.CancelledError:
pass
self._event_loop.submit(
self._underlying.__aexit__(exc_type, exc_value,
traceback)).result()
self._event_loop.__exit__(exc_type, exc_value, traceback)
assert self._close_callback is not None
self._close_callback(self, self._failure)