ack_set_tracker_impl.py

import queue
from collections import deque
from typing import Optional

from google.api_core.exceptions import FailedPrecondition
from google.cloud.pubsublite.cloudpubsub.internal.ack_set_tracker import AckSetTracker
from google.cloud.pubsublite.internal.wire.committer import Committer
from google.cloud.pubsublite_v1 import Cursor


class AckSetTrackerImpl(AckSetTracker):
  _committer: Committer

  _receipts: "deque[int]"
  _acks: "queue.PriorityQueue[int]"

  def __init__(self, committer: Committer):
    super().__init__()
    self._committer = committer
    self._receipts = deque()
    self._acks = queue.PriorityQueue()

  def track(self, offset: int):
    if len(self._receipts) > 0:
      last = self._receipts[0]
      if last >= offset:
        raise FailedPrecondition(f"Tried to track message {offset} which is before last tracked message {last}.")
    self._receipts.append(offset)

  async def ack(self, offset: int):
    # Note: put_nowait is used here and below to ensure that the below logic is executed without yielding
    # to another coroutine in the event loop. The queue is unbounded so it will never throw.
    self._acks.put_nowait(offset)
    prefix_acked_offset: Optional[int] = None
    while len(self._receipts) != 0 and not self._acks.empty():
      receipt = self._receipts.popleft()
      ack = self._acks.get_nowait()
      if receipt == ack:
        prefix_acked_offset = receipt
        continue
      self._receipts.append(receipt)
      self._acks.put(ack)
      break
    if prefix_acked_offset is None:
      return
    # Convert from last acked to first unacked.
    await self._committer.commit(Cursor(offset=prefix_acked_offset+1))

  async def __aenter__(self):
    await self._committer.__aenter__()

  async def __aexit__(self, exc_type, exc_value, traceback):
    await self._committer.__aexit__(exc_type, exc_value, traceback)