Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Implement Subscriber, which handles flow control and batch mess…
…age processing. (#16) * feat: Implement Subscriber, which handles flow control and batch message processing. Also ensure all asynchronous loopers are torn down when their underlying objects are.
- Loading branch information
1 parent
0a09bb3
commit 697df4a
Showing
7 changed files
with
473 additions
and
5 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,28 @@ | ||
| from abc import abstractmethod | ||
| from typing import AsyncContextManager | ||
| from google.cloud.pubsublite_v1.types import SequencedMessage, FlowControlRequest | ||
|
|
||
|
|
||
| class Subscriber(AsyncContextManager): | ||
| """ | ||
| A Pub/Sub Lite asynchronous wire protocol subscriber. | ||
| """ | ||
| @abstractmethod | ||
| async def read(self) -> SequencedMessage: | ||
| """ | ||
| Read the next message off of the stream. | ||
| Returns: | ||
| The next message. | ||
| Raises: | ||
| GoogleAPICallError: On a permanent error. | ||
| """ | ||
| raise NotImplementedError() | ||
|
|
||
| @abstractmethod | ||
| async def allow_flow(self, request: FlowControlRequest): | ||
| """ | ||
| Allow an additional amount of messages and bytes to be sent to this client. | ||
| """ | ||
| raise NotImplementedError() |
135 changes: 135 additions & 0 deletions
135
google/cloud/pubsublite/internal/wire/subscriber_impl.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,135 @@ | ||
| import asyncio | ||
| from typing import Optional | ||
|
|
||
| from google.api_core.exceptions import GoogleAPICallError, FailedPrecondition | ||
|
|
||
| from google.cloud.pubsublite.internal.wire.connection import Request, Connection, Response, ConnectionFactory | ||
| from google.cloud.pubsublite.internal.wire.connection_reinitializer import ConnectionReinitializer | ||
| from google.cloud.pubsublite.internal.wire.flow_control_batcher import FlowControlBatcher | ||
| from google.cloud.pubsublite.internal.wire.retrying_connection import RetryingConnection | ||
| from google.cloud.pubsublite.internal.wire.subscriber import Subscriber | ||
| from google.cloud.pubsublite_v1 import SubscribeRequest, SubscribeResponse, FlowControlRequest, SequencedMessage, \ | ||
| InitialSubscribeRequest, SeekRequest, Cursor | ||
|
|
||
|
|
||
| class SubscriberImpl(Subscriber, ConnectionReinitializer[SubscribeRequest, SubscribeResponse]): | ||
| _initial: InitialSubscribeRequest | ||
| _token_flush_seconds: float | ||
| _connection: RetryingConnection[SubscribeRequest, SubscribeResponse] | ||
|
|
||
| _outstanding_flow_control: FlowControlBatcher | ||
|
|
||
| _reinitializing: bool | ||
| _last_received_offset: Optional[int] | ||
|
|
||
| _message_queue: 'asyncio.Queue[SequencedMessage]' | ||
|
|
||
| _receiver: Optional[asyncio.Future] | ||
| _flusher: Optional[asyncio.Future] | ||
|
|
||
| def __init__(self, initial: InitialSubscribeRequest, token_flush_seconds: float, | ||
| factory: ConnectionFactory[SubscribeRequest, SubscribeResponse]): | ||
| self._initial = initial | ||
| self._token_flush_seconds = token_flush_seconds | ||
| self._connection = RetryingConnection(factory, self) | ||
| self._outstanding_flow_control = FlowControlBatcher() | ||
| self._reinitializing = False | ||
| self._last_received_offset = None | ||
| self._message_queue = asyncio.Queue() | ||
| self._receiver = None | ||
| self._flusher = None | ||
|
|
||
| async def __aenter__(self): | ||
| await self._connection.__aenter__() | ||
| return self | ||
|
|
||
| def _start_loopers(self): | ||
| assert self._receiver is None | ||
| assert self._flusher is None | ||
| self._receiver = asyncio.ensure_future(self._receive_loop()) | ||
| self._flusher = asyncio.ensure_future(self._flush_loop()) | ||
|
|
||
| async def _stop_loopers(self): | ||
| if self._receiver: | ||
| self._receiver.cancel() | ||
| await self._receiver | ||
| self._receiver = None | ||
| if self._flusher: | ||
| self._flusher.cancel() | ||
| await self._flusher | ||
| self._flusher = None | ||
|
|
||
| def _handle_response(self, response: SubscribeResponse): | ||
| if "messages" not in response: | ||
| self._connection.fail(FailedPrecondition("Received an invalid subsequent response on the subscribe stream.")) | ||
| return | ||
| self._outstanding_flow_control.on_messages(response.messages.messages) | ||
| for message in response.messages.messages: | ||
| if self._last_received_offset is not None and message.cursor.offset <= self._last_received_offset: | ||
| self._connection.fail(FailedPrecondition( | ||
| "Received an invalid out of order message from the server. Message is {}, previous last received is {}.".format( | ||
| message.cursor.offset, self._last_received_offset))) | ||
| return | ||
| self._last_received_offset = message.cursor.offset | ||
| for message in response.messages.messages: | ||
| # queue is unbounded. | ||
| self._message_queue.put_nowait(message) | ||
|
|
||
| async def _receive_loop(self): | ||
| try: | ||
| while True: | ||
| response = await self._connection.read() | ||
| self._handle_response(response) | ||
| except (asyncio.CancelledError, GoogleAPICallError): | ||
| return | ||
|
|
||
| async def _try_send_tokens(self): | ||
| req = self._outstanding_flow_control.release_pending_request() | ||
| if req is None: | ||
| return | ||
| try: | ||
| await self._connection.write(SubscribeRequest(flow_control=req)) | ||
| except GoogleAPICallError: | ||
| # May be transient, in which case these tokens will be resent. | ||
| pass | ||
|
|
||
| async def _flush_loop(self): | ||
| try: | ||
| while True: | ||
| await asyncio.sleep(self._token_flush_seconds) | ||
| await self._try_send_tokens() | ||
| except asyncio.CancelledError: | ||
| return | ||
|
|
||
| async def __aexit__(self, exc_type, exc_val, exc_tb): | ||
| await self._stop_loopers() | ||
| await self._connection.__aexit__(exc_type, exc_val, exc_tb) | ||
|
|
||
| async def reinitialize(self, connection: Connection[SubscribeRequest, SubscribeResponse]): | ||
| self._reinitializing = True | ||
| await self._stop_loopers() | ||
| await connection.write(SubscribeRequest(initial=self._initial)) | ||
| response = await connection.read() | ||
| if "initial" not in response: | ||
| self._connection.fail(FailedPrecondition("Received an invalid initial response on the subscribe stream.")) | ||
| return | ||
| if self._last_received_offset is not None: | ||
| # Perform a seek to get the next message after the one we received. | ||
| await connection.write(SubscribeRequest(seek=SeekRequest(cursor=Cursor(offset=self._last_received_offset + 1)))) | ||
| seek_response = await connection.read() | ||
| if "seek" not in seek_response: | ||
| self._connection.fail(FailedPrecondition("Received an invalid seek response on the subscribe stream.")) | ||
| return | ||
| tokens = self._outstanding_flow_control.request_for_restart() | ||
| if tokens is not None: | ||
| await connection.write(SubscribeRequest(flow_control=tokens)) | ||
| self._reinitializing = False | ||
| self._start_loopers() | ||
|
|
||
| async def read(self) -> SequencedMessage: | ||
| return await self._connection.await_unless_failed(self._message_queue.get()) | ||
|
|
||
| async def allow_flow(self, request: FlowControlRequest): | ||
| self._outstanding_flow_control.add(request) | ||
| if not self._reinitializing and self._outstanding_flow_control.should_expedite(): | ||
| await self._try_send_tokens() |
Oops, something went wrong.