def test_ordered_messages_two_keys_interleaved():
moh = messages_on_hold.MessagesOnHold()
# Put message with "key1".
msg1 = make_message(ack_id="ack1", ordering_key="key1")
moh.put(msg1)
assert moh.size == 1
# Put message with another key: "key2".
msg2 = make_message(ack_id="ack2", ordering_key="key2")
moh.put(msg2)
assert moh.size == 2
# Put second message with "key1".
msg3 = make_message(ack_id="ack3", ordering_key="key1")
moh.put(msg3)
assert moh.size == 3
# Get first message for "key1"
assert moh.get() == msg1
assert moh.size == 2
# Get another message. msg2 with "key2" is next in line in the queue.
assert moh.get() is msg2
assert moh.size == 1
# Activate "key1". Clean up state for "key1" because another message with
# the same key (msg3) hasn't been sorted into _pending_ordered_messages yet
# through a call to get().
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert not callback_tracker.called
assert moh.size == 1
# Get another message. msg3 is next in line in the queue.
assert moh.get() is msg3
assert moh.size == 0
# Activate "key2" to mark msg2 as complete. Release no messages because
# there are none left for that key. State for "key2" should be cleaned up.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key2"], callback_tracker)
assert not callback_tracker.called
assert moh.size == 0
# Activate "key1" to mark msg3 as complete. Release no messages because
# there are none left for that key. State for "key1" should be cleaned up.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert not callback_tracker.called
# Check that clean-up happened.
assert moh.size == 0
assert len(moh._messages_on_hold) == 0
assert len(moh._pending_ordered_messages) == 0
# No messages left.
assert moh.get() is None
assert moh.size == 0
def __init__(
self,
client,
subscription,
flow_control=types.FlowControl(),
scheduler=None,
use_legacy_flow_control=False,
await_callbacks_on_shutdown=False,
):
self._client = client
self._subscription = subscription
self._flow_control = flow_control
self._use_legacy_flow_control = use_legacy_flow_control
self._await_callbacks_on_shutdown = await_callbacks_on_shutdown
self._ack_histogram = histogram.Histogram()
self._last_histogram_size = 0
self._ack_deadline = 10
self._rpc = None
self._callback = None
self._closing = threading.Lock()
self._closed = False
self._close_callbacks = []
self._regular_shutdown_thread = None # Created on intentional shutdown.
# Generate a random client id tied to this object. All streaming pull
# connections (initial and re-connects) will then use the same client
# id. Doing so lets the server establish affinity even across stream
# disconncetions.
self._client_id = str(uuid.uuid4())
if scheduler is None:
self._scheduler = (
google.cloud.pubsub_v1.subscriber.scheduler.ThreadScheduler())
else:
self._scheduler = scheduler
# A collection for the messages that have been received from the server,
# but not yet sent to the user callback.
self._messages_on_hold = messages_on_hold.MessagesOnHold()
# The total number of bytes consumed by the messages currently on hold
self._on_hold_bytes = 0
# A lock ensuring that pausing / resuming the consumer are both atomic
# operations that cannot be executed concurrently. Needed for properly
# syncing these operations with the current leaser load. Additionally,
# the lock is used to protect modifications of internal data that
# affects the load computation, i.e. the count and size of the messages
# currently on hold.
self._pause_resume_lock = threading.Lock()
# The threads created in ``.open()``.
self._dispatcher = None
self._leaser = None
self._consumer = None
self._heartbeater = None
def test_ordered_messages_two_keys():
moh = messages_on_hold.MessagesOnHold()
# Put message with "key1".
msg1 = make_message(ack_id="ack1", ordering_key="key1")
moh.put(msg1)
assert moh.size == 1
# Put second message with "key1".
msg2 = make_message(ack_id="ack2", ordering_key="key1")
moh.put(msg2)
assert moh.size == 2
# Put message with another key: "key2".
msg3 = make_message(ack_id="ack3", ordering_key="key2")
moh.put(msg3)
assert moh.size == 3
# Get first message for "key1"
assert moh.get() == msg1
assert moh.size == 2
# Get another message. Still waiting on the previously-sent message for
# "key1", so release msg3 with key "key2".
assert moh.get() is msg3
assert moh.size == 1
# Activate "key1" to release the second message with that key.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert callback_tracker.called
assert callback_tracker.message == msg2
assert moh.size == 0
# Activate "key2" and release no messages because there are none left for
# that key. State for "key2" should be cleaned up.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key2"], callback_tracker)
assert not callback_tracker.called
assert moh.size == 0
# Activate "key1" again to mark msg2 as complete. Since there are no other
# messages for that key, clean up state for both keys.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert not callback_tracker.called
# Check that clean-up happened.
assert moh.size == 0
assert len(moh._messages_on_hold) == 0
assert len(moh._pending_ordered_messages) == 0
# No messages left.
assert moh.get() is None
assert moh.size == 0
def test_ordered_and_unordered_messages_interleaved():
moh = messages_on_hold.MessagesOnHold()
# Put message with "key1".
msg1 = make_message(ack_id="ack1", ordering_key="key1")
moh.put(msg1)
assert moh.size == 1
# Put another message "key1"
msg2 = make_message(ack_id="ack2", ordering_key="key1")
moh.put(msg2)
assert moh.size == 2
# Put a message with no ordering key.
msg3 = make_message(ack_id="ack3", ordering_key="")
moh.put(msg3)
assert moh.size == 3
# Get first message for "key1"
assert moh.get() == msg1
assert moh.size == 2
# Get another message. msg2 will be skipped because another message with the
# same key (msg1) is in flight.
assert moh.get() is msg3
assert moh.size == 1
# Activate "key1". Send msg2, the next in line for the same ordering key.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert callback_tracker.called
assert callback_tracker.message == msg2
assert moh.size == 0
# No more messages left.
assert moh.get() is None
# Activate "key1" to mark msg2 as complete. Release no messages because
# there are none left for that key. State for "key1" should be cleaned up.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert not callback_tracker.called
# Check that clean-up happened.
assert moh.size == 0
assert len(moh._messages_on_hold) == 0
assert len(moh._pending_ordered_messages) == 0
# No messages left.
assert moh.get() is None
assert moh.size == 0
def test_put_and_get_unordered_messages():
moh = messages_on_hold.MessagesOnHold()
msg1 = make_message(ack_id="ack1", ordering_key="")
moh.put(msg1)
assert moh.size == 1
msg2 = make_message(ack_id="ack2", ordering_key="")
moh.put(msg2)
assert moh.size == 2
assert moh.get() == msg1
assert moh.size == 1
assert moh.get() == msg2
assert moh.size == 0
assert moh.get() is None
def __init__(self,
client,
subscription,
flow_control=types.FlowControl(),
scheduler=None):
self._client = client
self._subscription = subscription
self._flow_control = flow_control
self._ack_histogram = histogram.Histogram()
self._last_histogram_size = 0
self._ack_deadline = 10
self._rpc = None
self._callback = None
self._closing = threading.Lock()
self._closed = False
self._close_callbacks = []
if scheduler is None:
self._scheduler = (
google.cloud.pubsub_v1.subscriber.scheduler.ThreadScheduler())
else:
self._scheduler = scheduler
# A collection for the messages that have been received from the server,
# but not yet sent to the user callback.
self._messages_on_hold = messages_on_hold.MessagesOnHold()
# The total number of bytes consumed by the messages currently on hold
self._on_hold_bytes = 0
# A lock ensuring that pausing / resuming the consumer are both atomic
# operations that cannot be executed concurrently. Needed for properly
# syncing these operations with the current leaser load. Additionally,
# the lock is used to protect modifications of internal data that
# affects the load computation, i.e. the count and size of the messages
# currently on hold.
self._pause_resume_lock = threading.Lock()
# The threads created in ``.open()``.
self._dispatcher = None
self._leaser = None
self._consumer = None
self._heartbeater = None
def test_ordered_messages_one_key():
moh = messages_on_hold.MessagesOnHold()
msg1 = make_message(ack_id="ack1", ordering_key="key1")
moh.put(msg1)
assert moh.size == 1
msg2 = make_message(ack_id="ack2", ordering_key="key1")
moh.put(msg2)
assert moh.size == 2
# Get first message for "key1"
assert moh.get() == msg1
assert moh.size == 1
# Still waiting on the previously-sent message for "key1", and there are no
# other messages, so return None.
assert moh.get() is None
assert moh.size == 1
# Activate "key1" to release the second message with that key.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert callback_tracker.called
assert callback_tracker.message == msg2
assert moh.size == 0
assert len(moh._pending_ordered_messages) == 1
# Activate "key1" again. There are no other messages for that key, so clean
# up state for that key.
callback_tracker = ScheduleMessageCallbackTracker()
moh.activate_ordering_keys(["key1"], callback_tracker)
assert not callback_tracker.called
# Check that clean-up happened.
assert moh.size == 0
assert len(moh._messages_on_hold) == 0
assert len(moh._pending_ordered_messages) == 0
# No messages left.
assert moh.get() is None
assert moh.size == 0
def test_init():
moh = messages_on_hold.MessagesOnHold()
assert moh.size == 0
assert moh.get() is None
