class SinglePartitionPublisher( Publisher, ConnectionReinitializer[PublishRequest, PublishResponse], BatchTester[PubSubMessage], ): _initial: InitialPublishRequest _batching_settings: BatchSettings _connection: RetryingConnection[PublishRequest, PublishResponse] _batcher: SerialBatcher[PubSubMessage, Cursor] _outstanding_writes: List[List[WorkItem[PubSubMessage, Cursor]]] _receiver: Optional[asyncio.Future] _flusher: Optional[asyncio.Future] def __init__( self, initial: InitialPublishRequest, batching_settings: BatchSettings, factory: ConnectionFactory[PublishRequest, PublishResponse], ): self._initial = initial self._batching_settings = batching_settings self._connection = RetryingConnection(factory, self) self._batcher = SerialBatcher(self) self._outstanding_writes = [] self._receiver = None self._flusher = None @property def _partition(self) -> Partition: return Partition(self._initial.partition) 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 wait_ignore_errors(self._receiver) self._receiver = None if self._flusher: self._flusher.cancel() await wait_ignore_errors(self._flusher) self._flusher = None def _handle_response(self, response: PublishResponse): if "message_response" not in response: self._connection.fail( FailedPrecondition( "Received an invalid subsequent response on the publish stream." ) ) if not self._outstanding_writes: self._connection.fail( FailedPrecondition( "Received an publish response on the stream with no outstanding publishes." ) ) next_offset: Cursor = response.message_response.start_cursor.offset batch: List[WorkItem[PubSubMessage]] = self._outstanding_writes.pop(0) for item in batch: item.response_future.set_result(Cursor(offset=next_offset)) next_offset += 1 async def _receive_loop(self): while True: response = await self._connection.read() self._handle_response(response) async def _flush_loop(self): while True: await asyncio.sleep(self._batching_settings.max_latency) await self._flush() async def __aexit__(self, exc_type, exc_val, exc_tb): if self._connection.error(): self._fail_if_retrying_failed() else: await self._flush() await self._stop_loopers() await self._connection.__aexit__(exc_type, exc_val, exc_tb) def _fail_if_retrying_failed(self): if self._connection.error(): for batch in self._outstanding_writes: for item in batch: item.response_future.set_exception(self._connection.error()) async def _flush(self): batch = self._batcher.flush() if not batch: return self._outstanding_writes.append(batch) aggregate = PublishRequest() aggregate.message_publish_request.messages = [item.request for item in batch] try: await self._connection.write(aggregate) except GoogleAPICallError as e: _LOGGER.debug(f"Failed publish on stream: {e}") self._fail_if_retrying_failed() async def publish(self, message: PubSubMessage) -> MessageMetadata: cursor_future = self._batcher.add(message) if self._batcher.should_flush(): await self._flush() return MessageMetadata(self._partition, await cursor_future) async def reinitialize( self, connection: Connection[PublishRequest, PublishResponse] ): await self._stop_loopers() await connection.write(PublishRequest(initial_request=self._initial)) response = await connection.read() if "initial_response" not in response: self._connection.fail( FailedPrecondition( "Received an invalid initial response on the publish stream." ) ) for batch in self._outstanding_writes: aggregate = PublishRequest() aggregate.message_publish_request.messages = [ item.request for item in batch ] await connection.write(aggregate) self._start_loopers() def test(self, requests: Iterable[PubSubMessage]) -> bool: request_count = 0 byte_count = 0 for req in requests: request_count += 1 byte_count += PubSubMessage.pb(req).ByteSize() return (request_count >= _MAX_MESSAGES) or (byte_count >= _MAX_BYTES)
class CommitterImpl( Committer, ConnectionReinitializer[StreamingCommitCursorRequest, StreamingCommitCursorResponse], ): _initial: InitialCommitCursorRequest _flush_seconds: float _connection: RetryingConnection[StreamingCommitCursorRequest, StreamingCommitCursorResponse] _batcher: SerialBatcher[Cursor, None] _outstanding_commits: List[List[WorkItem[Cursor, None]]] _receiver: Optional[asyncio.Future] _flusher: Optional[asyncio.Future] _empty: asyncio.Event def __init__( self, initial: InitialCommitCursorRequest, flush_seconds: float, factory: ConnectionFactory[StreamingCommitCursorRequest, StreamingCommitCursorResponse], ): self._initial = initial self._flush_seconds = flush_seconds self._connection = RetryingConnection(factory, self) self._batcher = SerialBatcher() self._outstanding_commits = [] self._receiver = None self._flusher = None self._empty = asyncio.Event() self._empty.set() 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 wait_ignore_errors(self._receiver) self._receiver = None if self._flusher: self._flusher.cancel() await wait_ignore_errors(self._flusher) self._flusher = None def _handle_response(self, response: StreamingCommitCursorResponse): if "commit" not in response: self._connection.fail( FailedPrecondition( "Received an invalid subsequent response on the commit stream." )) if response.commit.acknowledged_commits > len( self._outstanding_commits): self._connection.fail( FailedPrecondition( "Received a commit response on the stream with no outstanding commits." )) for _ in range(response.commit.acknowledged_commits): batch = self._outstanding_commits.pop(0) for item in batch: item.response_future.set_result(None) if len(self._outstanding_commits) == 0: self._empty.set() async def _receive_loop(self): while True: response = await self._connection.read() self._handle_response(response) async def _flush_loop(self): while True: await asyncio.sleep(self._flush_seconds) await self._flush() async def __aexit__(self, exc_type, exc_val, exc_tb): await self._stop_loopers() if self._connection.error(): self._fail_if_retrying_failed() else: await self._flush() await self._connection.__aexit__(exc_type, exc_val, exc_tb) def _fail_if_retrying_failed(self): if self._connection.error(): for batch in self._outstanding_commits: for item in batch: item.response_future.set_exception( self._connection.error()) async def _flush(self): batch = self._batcher.flush() if not batch: return self._outstanding_commits.append(batch) self._empty.clear() req = StreamingCommitCursorRequest() req.commit.cursor = batch[-1].request try: await self._connection.write(req) except GoogleAPICallError as e: _LOGGER.debug(f"Failed commit on stream: {e}") self._fail_if_retrying_failed() async def wait_until_empty(self): await self._flush() await self._connection.await_unless_failed(self._empty.wait()) async def commit(self, cursor: Cursor) -> None: future = self._batcher.add(cursor) await future async def reinitialize( self, connection: Connection[StreamingCommitCursorRequest, StreamingCommitCursorResponse], last_error: Optional[GoogleAPICallError], ): await self._stop_loopers() await connection.write( StreamingCommitCursorRequest(initial=self._initial)) response = await connection.read() if "initial" not in response: self._connection.fail( FailedPrecondition( "Received an invalid initial response on the publish stream." )) if self._outstanding_commits: # Roll up outstanding commits rollup: List[WorkItem[Cursor, None]] = [] for batch in self._outstanding_commits: for item in batch: rollup.append(item) self._outstanding_commits = [rollup] req = StreamingCommitCursorRequest() req.commit.cursor = rollup[-1].request await connection.write(req) self._start_loopers()