def __init__( self, initial: InitialPartitionAssignmentRequest, factory: ConnectionFactory[PartitionAssignmentRequest, PartitionAssignment], ): self._initial = initial self._connection = RetryingConnection(factory, self) self._outstanding_assignment = False self._receiver = None self._new_assignment = asyncio.Queue(maxsize=1)
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
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) self._outstanding_commits = [] self._receiver = None self._flusher = None
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
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._next_to_commit = None self._outstanding_commits = [] self._receiver = None self._flusher = None self._empty = asyncio.Event() self._empty.set()
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 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 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: 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): while True: response = await self._connection.read() self._handle_response(response) 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): while True: await asyncio.sleep(self._token_flush_seconds) await self._try_send_tokens() 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()
class AssignerImpl(Assigner, ConnectionReinitializer[PartitionAssignmentRequest, PartitionAssignment]): _initial: InitialPartitionAssignmentRequest _connection: RetryingConnection[PartitionAssignmentRequest, PartitionAssignment] _outstanding_assignment: bool _receiver: Optional[asyncio.Future] # A queue that may only hold one element with the next assignment. _new_assignment: "asyncio.Queue[Set[Partition]]" def __init__( self, initial: InitialPartitionAssignmentRequest, factory: ConnectionFactory[PartitionAssignmentRequest, PartitionAssignment], ): self._initial = initial self._connection = RetryingConnection(factory, self) self._outstanding_assignment = False self._receiver = None self._new_assignment = asyncio.Queue(maxsize=1) async def __aenter__(self): await self._connection.__aenter__() return self def _start_receiver(self): assert self._receiver is None self._receiver = asyncio.ensure_future(self._receive_loop()) async def _stop_receiver(self): if self._receiver: self._receiver.cancel() await wait_ignore_errors(self._receiver) self._receiver = None async def _receive_loop(self): while True: response = await self._connection.read() if self._outstanding_assignment or not self._new_assignment.empty( ): self._connection.fail( FailedPrecondition( "Received a duplicate assignment on the stream while one was outstanding." )) return self._outstanding_assignment = True partitions = set() for partition in response.partitions: partitions.add(Partition(partition)) self._new_assignment.put_nowait(partitions) async def __aexit__(self, exc_type, exc_val, exc_tb): await self._stop_receiver() await self._connection.__aexit__(exc_type, exc_val, exc_tb) async def reinitialize( self, connection: Connection[PartitionAssignmentRequest, PartitionAssignment], last_error: Optional[GoogleAPICallError], ): self._outstanding_assignment = False while not self._new_assignment.empty(): self._new_assignment.get_nowait() await self._stop_receiver() await connection.write( PartitionAssignmentRequest(initial=self._initial)) self._start_receiver() async def get_assignment(self) -> Set[Partition]: if self._outstanding_assignment: try: await self._connection.write( PartitionAssignmentRequest(ack=PartitionAssignmentAck())) self._outstanding_assignment = False except GoogleAPICallError as e: # If there is a failure to ack, keep going. The stream likely restarted. _LOGGER.debug( f"Assignment ack attempt failed due to stream failure: {e}" ) return await self._connection.await_unless_failed( self._new_assignment.get())
class CommitterImpl( Committer, ConnectionReinitializer[StreamingCommitCursorRequest, StreamingCommitCursorResponse], ): _initial: InitialCommitCursorRequest _flush_seconds: float _connection: RetryingConnection[StreamingCommitCursorRequest, StreamingCommitCursorResponse] _next_to_commit: Optional[Cursor] _outstanding_commits: List[Cursor] _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._next_to_commit = None 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): self._outstanding_commits.pop(0) 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 not self._connection.error(): await self._flush() await self._connection.__aexit__(exc_type, exc_val, exc_tb) async def _flush(self): if self._next_to_commit is None: return req = StreamingCommitCursorRequest() req.commit.cursor = self._next_to_commit self._outstanding_commits.append(self._next_to_commit) self._next_to_commit = None self._empty.clear() try: await self._connection.write(req) except GoogleAPICallError as e: _LOGGER.debug(f"Failed commit on stream: {e}") async def wait_until_empty(self): await self._flush() await self._connection.await_unless_failed(self._empty.wait()) def commit(self, cursor: Cursor) -> None: if self._connection.error(): raise self._connection.error() self._next_to_commit = cursor @overrides async def stop_processing(self, error: GoogleAPICallError): await self._stop_loopers() @overrides async def reinitialize( self, connection: Connection[StreamingCommitCursorRequest, StreamingCommitCursorResponse], ): 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._next_to_commit is None: if self._outstanding_commits: self._next_to_commit = self._outstanding_commits[-1] self._outstanding_commits = [] self._start_loopers()
def retrying_connection(connection_factory, reinitializer): return RetryingConnection(connection_factory, reinitializer)
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()
class SubscriberImpl( Subscriber, ConnectionReinitializer[SubscribeRequest, SubscribeResponse] ): _base_initial: InitialSubscribeRequest _token_flush_seconds: float _connection: RetryingConnection[SubscribeRequest, SubscribeResponse] _reset_handler: SubscriberResetHandler _outstanding_flow_control: FlowControlBatcher _last_received_offset: Optional[int] _message_queue: "asyncio.Queue[List[SequencedMessage.meta.pb]]" _receiver: Optional[asyncio.Future] _flusher: Optional[asyncio.Future] def __init__( self, base_initial: InitialSubscribeRequest, token_flush_seconds: float, factory: ConnectionFactory[SubscribeRequest, SubscribeResponse], reset_handler: SubscriberResetHandler, ): self._base_initial = base_initial self._token_flush_seconds = token_flush_seconds self._connection = RetryingConnection(factory, self) self._reset_handler = reset_handler 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 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: SubscribeResponse): if "messages" not in response: self._connection.fail( FailedPrecondition( "Received an invalid subsequent response on the subscribe stream." ) ) return # Workaround for incredibly slow proto-plus-python accesses messages = list(response.messages.messages._pb) self._outstanding_flow_control.on_messages(messages) for message in 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 # queue is unbounded. self._message_queue.put_nowait(messages) async def _receive_loop(self): while True: response = await self._connection.read() self._handle_response(response) 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): while True: await asyncio.sleep(self._token_flush_seconds) await self._try_send_tokens() async def __aexit__(self, exc_type, exc_val, exc_tb): await self._stop_loopers() await self._connection.__aexit__(exc_type, exc_val, exc_tb) @overrides async def stop_processing(self, error: GoogleAPICallError): await self._stop_loopers() if is_reset_signal(error): # Discard undelivered messages and refill flow control tokens. while not self._message_queue.empty(): batch: List[SequencedMessage.meta.pb] = self._message_queue.get_nowait() allowed_bytes = sum(message.size_bytes for message in batch) self._outstanding_flow_control.add( FlowControlRequest( allowed_messages=len(batch), allowed_bytes=allowed_bytes, ) ) await self._reset_handler.handle_reset() self._last_received_offset = None @overrides async def reinitialize( self, connection: Connection[SubscribeRequest, SubscribeResponse] ): initial = deepcopy(self._base_initial) if self._last_received_offset is not None: initial.initial_location = SeekRequest( cursor=Cursor(offset=self._last_received_offset + 1) ) else: initial.initial_location = SeekRequest( named_target=SeekRequest.NamedTarget.COMMITTED_CURSOR ) await connection.write(SubscribeRequest(initial=initial)) response = await connection.read() if "initial" not in response: self._connection.fail( FailedPrecondition( "Received an invalid initial 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._start_loopers() async def read(self) -> List[SequencedMessage.meta.pb]: return await self._connection.await_unless_failed(self._message_queue.get()) def allow_flow(self, request: FlowControlRequest): self._outstanding_flow_control.add(request)