def test_lease_load_and_pause():
manager = make_manager(
flow_control=types.FlowControl(max_messages=10, max_bytes=1000))
manager._leaser = leaser.Leaser(manager)
manager._consumer = mock.create_autospec(
bidi.BackgroundConsumer, instance=True)
manager._consumer.is_paused = False
# This should mean that our messages count is at 10%, and our bytes
# are at 15%; load should return the higher (0.15), and shouldn't cause
# the consumer to pause.
manager.leaser.add([requests.LeaseRequest(ack_id='one', byte_size=150)])
assert manager.load == 0.15
manager.maybe_pause_consumer()
manager._consumer.pause.assert_not_called()
# After this message is added, the messages should be higher at 20%
# (versus 16% for bytes).
manager.leaser.add([requests.LeaseRequest(ack_id='two', byte_size=10)])
assert manager.load == 0.2
# Returning a number above 100% is fine, and it should cause this to pause.
manager.leaser.add([requests.LeaseRequest(ack_id='three', byte_size=1000)])
assert manager.load == 1.16
manager.maybe_pause_consumer()
manager._consumer.pause.assert_called_once()
def test_drop_and_resume():
manager = make_manager(
flow_control=types.FlowControl(max_messages=10, max_bytes=1000))
manager._leaser = leaser.Leaser(manager)
manager._consumer = mock.create_autospec(
bidi.BackgroundConsumer, instance=True)
manager._consumer.is_paused = True
# Add several messages until we're over the load threshold.
manager.leaser.add([
requests.LeaseRequest(ack_id='one', byte_size=750),
requests.LeaseRequest(ack_id='two', byte_size=250)])
assert manager.load == 1.0
# Trying to resume now should have no effect as we're over the threshold.
manager.maybe_resume_consumer()
manager._consumer.resume.assert_not_called()
# Drop the 200 byte message, which should put us under the resume
# threshold.
manager.leaser.remove([
requests.DropRequest(ack_id='two', byte_size=250)])
manager.maybe_resume_consumer()
manager._consumer.resume.assert_called_once()
def test_maintain_leases_outdated_items(sleep, time):
manager = create_manager()
make_sleep_mark_manager_as_inactive(sleep, manager)
leaser_ = leaser.Leaser(manager)
# Add these items at the beginning of the timeline
time.return_value = 0
leaser_.add([
requests.LeaseRequest(ack_id='ack1', byte_size=50)])
# Add another item at towards end of the timeline
time.return_value = manager.flow_control.max_lease_duration - 1
leaser_.add([
requests.LeaseRequest(ack_id='ack2', byte_size=50)])
# Now make sure time reports that we are at the end of our timeline.
time.return_value = manager.flow_control.max_lease_duration + 1
leaser_.maintain_leases()
# Only ack2 should be renewed. ack1 should've been dropped
manager.dispatcher.modify_ack_deadline.assert_called_once_with([
requests.ModAckRequest(
ack_id='ack2',
seconds=10,
)
])
manager.dispatcher.drop.assert_called_once_with([
requests.DropRequest(ack_id='ack1', byte_size=50)
])
sleep.assert_called()
def test_add_already_managed(caplog):
caplog.set_level(logging.DEBUG)
leaser_ = leaser.Leaser(mock.sentinel.manager)
leaser_.add([requests.LeaseRequest(ack_id="ack1", byte_size=50)])
leaser_.add([requests.LeaseRequest(ack_id="ack1", byte_size=50)])
assert "already lease managed" in caplog.text
def test_add_already_managed(caplog):
caplog.set_level(logging.DEBUG)
leaser_ = leaser.Leaser(mock.sentinel.subscriber)
leaser_.add([requests.LeaseRequest(ack_id='ack1', byte_size=50)])
leaser_.add([requests.LeaseRequest(ack_id='ack1', byte_size=50)])
assert 'already lease managed' in caplog.text
def test_add_and_remove():
leaser_ = leaser.Leaser(mock.sentinel.manager)
leaser_.add([requests.LeaseRequest(ack_id="ack1", byte_size=50)])
leaser_.add([requests.LeaseRequest(ack_id="ack2", byte_size=25)])
assert leaser_.message_count == 2
assert set(leaser_.ack_ids) == set(["ack1", "ack2"])
assert leaser_.bytes == 75
leaser_.remove([requests.DropRequest(ack_id="ack1", byte_size=50)])
assert leaser_.message_count == 1
assert set(leaser_.ack_ids) == set(["ack2"])
assert leaser_.bytes == 25
def test_add_and_remove():
leaser_ = leaser.Leaser(mock.sentinel.subscriber)
leaser_.add([requests.LeaseRequest(ack_id='ack1', byte_size=50)])
leaser_.add([requests.LeaseRequest(ack_id='ack2', byte_size=25)])
assert leaser_.message_count == 2
assert set(leaser_.ack_ids) == set(['ack1', 'ack2'])
assert leaser_.bytes == 75
leaser_.remove([requests.DropRequest(ack_id='ack1', byte_size=50)])
assert leaser_.message_count == 1
assert set(leaser_.ack_ids) == set(['ack2'])
assert leaser_.bytes == 25
def _maybe_release_messages(self):
"""Release (some of) the held messages if the current load allows for it.
The method tries to release as many messages as the current leaser load
would allow. Each released message is added to the lease management,
and the user callback is scheduled for it.
If there are currently no messageges on hold, or if the leaser is
already overloaded, this method is effectively a no-op.
The method assumes the caller has acquired the ``_pause_resume_lock``.
"""
while True:
if self.load >= 1.0:
break # already overloaded
try:
msg = self._messages_on_hold.get_nowait()
except queue.Empty:
break
self.leaser.add(
[requests.LeaseRequest(ack_id=msg.ack_id, byte_size=msg.size)]
)
_LOGGER.debug(
"Released held message to leaser, scheduling callback for it, "
"still on hold %s.",
self._messages_on_hold.qsize(),
)
self._scheduler.schedule(self._callback, msg)
def test_lease():
msg = create_message(b"foo", ack_id="bogus_ack_id")
with mock.patch.object(msg._request_queue, "put") as put:
msg.lease()
put.assert_called_once_with(
requests.LeaseRequest(ack_id="bogus_ack_id", byte_size=30))
check_call_types(put, requests.LeaseRequest)
def test__maybe_release_messages_below_overload():
manager = make_manager(
flow_control=types.FlowControl(max_messages=10, max_bytes=1000))
manager._callback = mock.sentinel.callback
# init leaser message count to 8 to leave room for 2 more messages
_leaser = manager._leaser = mock.create_autospec(leaser.Leaser)
fake_leaser_add(_leaser, init_msg_count=8, assumed_msg_size=25)
_leaser.add = mock.Mock(wraps=_leaser.add) # to spy on calls
messages = [
mock.create_autospec(message.Message,
instance=True,
ack_id="ack_foo",
size=11),
mock.create_autospec(message.Message,
instance=True,
ack_id="ack_bar",
size=22),
mock.create_autospec(message.Message,
instance=True,
ack_id="ack_baz",
size=33),
]
for msg in messages:
manager._messages_on_hold.put(msg)
# the actual call of MUT
manager._maybe_release_messages()
assert manager._messages_on_hold.qsize() == 1
msg = manager._messages_on_hold.get_nowait()
assert msg.ack_id == "ack_baz"
assert len(_leaser.add.mock_calls) == 2
expected_calls = [
mock.call([requests.LeaseRequest(ack_id="ack_foo", byte_size=11)]),
mock.call([requests.LeaseRequest(ack_id="ack_bar", byte_size=22)]),
]
_leaser.add.assert_has_calls(expected_calls)
schedule_calls = manager._scheduler.schedule.mock_calls
assert len(schedule_calls) == 2
for _, call_args, _ in schedule_calls:
assert call_args[0] == mock.sentinel.callback
assert isinstance(call_args[1], message.Message)
assert call_args[1].ack_id in ("ack_foo", "ack_bar")
def _on_response(self, response):
"""Process all received Pub/Sub messages.
For each message, send a modified acknowledgment request to the
server. This prevents expiration of the message due to buffering by
gRPC or proxy/firewall. This makes the server and client expiration
timer closer to each other thus preventing the message being
redelivered multiple times.
After the messages have all had their ack deadline updated, execute
the callback for each message using the executor.
"""
if response is None:
_LOGGER.debug(
"Response callback invoked with None, likely due to a "
"transport shutdown."
)
return
# IMPORTANT: Circumvent the wrapper class and operate on the raw underlying
# protobuf message to significantly gain on attribute access performance.
received_messages = response._pb.received_messages
_LOGGER.debug(
"Processing %s received message(s), currently on hold %s (bytes %s).",
len(received_messages),
self._messages_on_hold.size,
self._on_hold_bytes,
)
# Immediately (i.e. without waiting for the auto lease management)
# modack the messages we received, as this tells the server that we've
# received them.
items = [
requests.ModAckRequest(message.ack_id, self._ack_histogram.percentile(99))
for message in received_messages
]
self._dispatcher.modify_ack_deadline(items)
with self._pause_resume_lock:
for received_message in received_messages:
message = google.cloud.pubsub_v1.subscriber.message.Message(
received_message.message,
received_message.ack_id,
received_message.delivery_attempt,
self._scheduler.queue,
)
self._messages_on_hold.put(message)
self._on_hold_bytes += message.size
req = requests.LeaseRequest(
ack_id=message.ack_id,
byte_size=message.size,
ordering_key=message.ordering_key,
)
self.leaser.add([req])
self._maybe_release_messages()
self.maybe_pause_consumer()
def lease(self):
"""Inform the policy to lease this message continually.
.. note::
This method is called by the constructor, and you should never
need to call it manually.
"""
self._request_queue.put(
requests.LeaseRequest(ack_id=self._ack_id, byte_size=self.size))
def _on_response(self, response):
"""Process all received Pub/Sub messages.
For each message, send a modified acknowledgment request to the
server. This prevents expiration of the message due to buffering by
gRPC or proxy/firewall. This makes the server and client expiration
timer closer to each other thus preventing the message being
redelivered multiple times.
After the messages have all had their ack deadline updated, execute
the callback for each message using the executor.
"""
_LOGGER.debug(
"Processing %s received message(s), currenty on hold %s (bytes %s).",
len(response.received_messages),
self._messages_on_hold.qsize(),
self._on_hold_bytes,
)
# Immediately (i.e. without waiting for the auto lease management)
# modack the messages we received, as this tells the server that we've
# received them.
items = [
requests.ModAckRequest(message.ack_id, self._ack_histogram.percentile(99))
for message in response.received_messages
]
self._dispatcher.modify_ack_deadline(items)
invoke_callbacks_for = []
for received_message in response.received_messages:
message = google.cloud.pubsub_v1.subscriber.message.Message(
received_message.message, received_message.ack_id, self._scheduler.queue
)
# Making a decision based on the load, and modifying the data that
# affects the load -> needs a lock, as that state can be modified
# by different threads.
with self._pause_resume_lock:
if self.load < _MAX_LOAD:
invoke_callbacks_for.append(message)
else:
self._messages_on_hold.put(message)
self._on_hold_bytes += message.size
req = requests.LeaseRequest(ack_id=message.ack_id, byte_size=message.size)
self.leaser.add([req])
self.maybe_pause_consumer()
_LOGGER.debug(
"Scheduling callbacks for %s new messages, new total on hold %s (bytes %s).",
len(invoke_callbacks_for),
self._messages_on_hold.qsize(),
self._on_hold_bytes,
)
for msg in invoke_callbacks_for:
self._scheduler.schedule(self._callback, msg)
def test_lease():
manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [requests.LeaseRequest(ack_id="ack_id_string", byte_size=10)]
dispatcher_.lease(items)
manager.leaser.add.assert_called_once_with(items)
manager.maybe_pause_consumer.assert_called_once()
def test_remove_negative_bytes(caplog):
caplog.set_level(logging.DEBUG)
leaser_ = leaser.Leaser(mock.sentinel.manager)
leaser_.add([requests.LeaseRequest(ack_id="ack1", byte_size=50)])
leaser_.remove([requests.DropRequest(ack_id="ack1", byte_size=75)])
assert leaser_.bytes == 0
assert "unexpectedly negative" in caplog.text
def test_maintain_leases_outdated_items(time):
manager = create_manager()
leaser_ = leaser.Leaser(manager)
make_sleep_mark_event_as_done(leaser_)
# Add and start expiry timer at the beginning of the timeline.
time.return_value = 0
leaser_.add(
[requests.LeaseRequest(ack_id="ack1", byte_size=50, ordering_key="")])
leaser_.start_lease_expiry_timer(["ack1"])
# Add a message but don't start the lease expiry timer.
leaser_.add(
[requests.LeaseRequest(ack_id="ack2", byte_size=50, ordering_key="")])
# Add a message and start expiry timer towards the end of the timeline.
time.return_value = manager.flow_control.max_lease_duration - 1
leaser_.add(
[requests.LeaseRequest(ack_id="ack3", byte_size=50, ordering_key="")])
leaser_.start_lease_expiry_timer(["ack3"])
# Add a message towards the end of the timeline, but DO NOT start expiry
# timer.
leaser_.add(
[requests.LeaseRequest(ack_id="ack4", byte_size=50, ordering_key="")])
# Now make sure time reports that we are past the end of our timeline.
time.return_value = manager.flow_control.max_lease_duration + 1
leaser_.maintain_leases()
# ack2, ack3, and ack4 should be renewed. ack1 should've been dropped
modacks = manager.dispatcher.modify_ack_deadline.call_args.args[0]
expected = [
requests.ModAckRequest(ack_id="ack2", seconds=10),
requests.ModAckRequest(ack_id="ack3", seconds=10),
requests.ModAckRequest(ack_id="ack4", seconds=10),
]
# Use sorting to allow for ordering variance.
assert sorted(modacks) == sorted(expected)
manager.dispatcher.drop.assert_called_once_with(
[requests.DropRequest(ack_id="ack1", byte_size=50, ordering_key="")])
def test_maintain_leases_ack_ids():
manager = create_manager()
leaser_ = leaser.Leaser(manager)
make_sleep_mark_manager_as_inactive(leaser_)
leaser_.add([requests.LeaseRequest(ack_id="my ack id", byte_size=50)])
leaser_.maintain_leases()
manager.dispatcher.modify_ack_deadline.assert_called_once_with(
[requests.ModAckRequest(ack_id="my ack id", seconds=10)])
def test_lease():
msg = create_message(b'foo', ack_id='bogus_ack_id')
with mock.patch.object(msg._request_queue, 'put') as put:
msg.lease()
put.assert_called_once_with(
requests.LeaseRequest(
ack_id='bogus_ack_id',
byte_size=30,
))
check_call_types(put, requests.LeaseRequest)
def lease(self):
"""Inform the policy to lease this message continually.
.. note::
By default this method is called by the constructor, and you should
never need to call it manually, unless the
:class:`~.pubsub_v1.subscriber.message.Message` instance was
created with ``autolease=False``.
"""
self._request_queue.put(
requests.LeaseRequest(ack_id=self._ack_id, byte_size=self.size))
def _on_response(self, response):
"""Process all received Pub/Sub messages.
For each message, send a modified acknowledgment request to the
server. This prevents expiration of the message due to buffering by
gRPC or proxy/firewall. This makes the server and client expiration
timer closer to each other thus preventing the message being
redelivered multiple times.
After the messages have all had their ack deadline updated, execute
the callback for each message using the executor.
"""
_LOGGER.debug(
"Processing %s received message(s), currenty on hold %s.",
len(response.received_messages),
self._messages_on_hold.qsize(),
)
# Immediately modack the messages we received, as this tells the server
# that we've received them.
items = [
requests.ModAckRequest(message.ack_id, self._ack_histogram.percentile(99))
for message in response.received_messages
]
self._dispatcher.modify_ack_deadline(items)
invoke_callbacks_for = []
for received_message in response.received_messages:
message = google.cloud.pubsub_v1.subscriber.message.Message(
received_message.message,
received_message.ack_id,
self._scheduler.queue,
autolease=False,
)
if self.load < 1.0:
req = requests.LeaseRequest(
ack_id=message.ack_id, byte_size=message.size
)
self.leaser.add([req])
invoke_callbacks_for.append(message)
self.maybe_pause_consumer()
else:
self._messages_on_hold.put(message)
_LOGGER.debug(
"Scheduling callbacks for %s new messages, new total on hold %s.",
len(invoke_callbacks_for),
self._messages_on_hold.qsize(),
)
for msg in invoke_callbacks_for:
self._scheduler.schedule(self._callback, msg)
def test_maintain_leases_ack_ids(sleep):
subscriber_ = create_subscriber()
make_sleep_mark_subscriber_as_inactive(sleep, subscriber_)
leaser_ = leaser.Leaser(subscriber_)
leaser_.add([requests.LeaseRequest(ack_id='my ack id', byte_size=50)])
leaser_.maintain_leases()
subscriber_.modify_ack_deadline.assert_called_once_with(
[requests.ModAckRequest(
ack_id='my ack id',
seconds=10,
)])
sleep.assert_called()
def test_maintain_leases_inactive_manager(caplog):
caplog.set_level(logging.INFO)
manager = create_manager()
manager.is_active = False
leaser_ = leaser.Leaser(manager)
make_sleep_mark_event_as_done(leaser_)
leaser_.add([
requests.LeaseRequest(ack_id="my_ack_ID",
byte_size=42,
ordering_key="")
])
leaser_.maintain_leases()
# Leases should still be maintained even if the manager is inactive.
manager.dispatcher.modify_ack_deadline.assert_called()
assert "exiting" in caplog.text
def lease(self):
"""Inform the policy to lease this message continually.
.. note::
By default this method is called by the constructor, and you should
never need to call it manually, unless the
:class:`~.pubsub_v1.subscriber.message.Message` instance was
created with ``autolease=False``.
.. deprecated:: 0.44.0
Will be removed in future versions.
"""
warnings.warn(
"lease() is deprecated since 0.44.0, and will be removed in future versions.",
category=DeprecationWarning,
)
self._request_queue.put(
requests.LeaseRequest(ack_id=self._ack_id, byte_size=self.size))
from google.cloud.pubsub_v1.subscriber._protocol import dispatcher
from google.cloud.pubsub_v1.subscriber._protocol import helper_threads
from google.cloud.pubsub_v1.subscriber._protocol import requests
from google.cloud.pubsub_v1.subscriber._protocol import streaming_pull_manager
from google.pubsub_v1 import types as gapic_types
import mock
import pytest
@pytest.mark.parametrize(
"item,method_name",
[
(requests.AckRequest("0", 0, 0, ""), "ack"),
(requests.DropRequest("0", 0, ""), "drop"),
(requests.LeaseRequest("0", 0, ""), "lease"),
(requests.ModAckRequest("0", 0), "modify_ack_deadline"),
(requests.NackRequest("0", 0, ""), "nack"),
],
)
def test_dispatch_callback_active_manager(item, method_name):
manager = mock.create_autospec(
streaming_pull_manager.StreamingPullManager, instance=True
)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [item]
with mock.patch.object(dispatcher_, method_name) as method:
dispatcher_.dispatch_callback(items)
from google.cloud.pubsub_v1.subscriber._protocol import dispatcher
from google.cloud.pubsub_v1.subscriber._protocol import helper_threads
from google.cloud.pubsub_v1.subscriber._protocol import requests
from google.cloud.pubsub_v1.subscriber._protocol import streaming_pull_manager
from google.pubsub_v1 import types as gapic_types
import mock
import pytest
@pytest.mark.parametrize(
"item,method_name",
[
(requests.AckRequest(0, 0, 0, ""), "ack"),
(requests.DropRequest(0, 0, ""), "drop"),
(requests.LeaseRequest(0, 0, ""), "lease"),
(requests.ModAckRequest(0, 0), "modify_ack_deadline"),
(requests.NackRequest(0, 0, ""), "nack"),
],
)
def test_dispatch_callback(item, method_name):
manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [item]
with mock.patch.object(dispatcher_, method_name) as method:
dispatcher_.dispatch_callback(items)
method.assert_called_once_with([item])
from google.cloud.pubsub_v1 import types
from google.cloud.pubsub_v1.subscriber._protocol import dispatcher
from google.cloud.pubsub_v1.subscriber._protocol import helper_threads
from google.cloud.pubsub_v1.subscriber._protocol import requests
from google.cloud.pubsub_v1.subscriber._protocol import streaming_pull_manager
import mock
from six.moves import queue
import pytest
@pytest.mark.parametrize('item,method_name', [
(requests.AckRequest(0, 0, 0), 'ack'),
(requests.DropRequest(0, 0), 'drop'),
(requests.LeaseRequest(0, 0), 'lease'),
(requests.ModAckRequest(0, 0), 'modify_ack_deadline'),
(requests.NackRequest(0, 0), 'nack')
])
def test_dispatch_callback(item, method_name):
manager = mock.create_autospec(
streaming_pull_manager.StreamingPullManager, instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [item]
with mock.patch.object(dispatcher_, method_name) as method:
dispatcher_.dispatch_callback(items)
method.assert_called_once_with([item])
