def test_send_unary():
manager = make_manager()
manager._UNARY_REQUESTS = True
manager.send(
types.StreamingPullRequest(
ack_ids=["ack_id1", "ack_id2"],
modify_deadline_ack_ids=["ack_id3", "ack_id4", "ack_id5"],
modify_deadline_seconds=[10, 20, 20],
))
manager._client.acknowledge.assert_called_once_with(
subscription=manager._subscription, ack_ids=["ack_id1", "ack_id2"])
manager._client.modify_ack_deadline.assert_has_calls(
[
mock.call(
subscription=manager._subscription,
ack_ids=["ack_id3"],
ack_deadline_seconds=10,
),
mock.call(
subscription=manager._subscription,
ack_ids=["ack_id4", "ack_id5"],
ack_deadline_seconds=20,
),
],
any_order=True,
)
def test_maintain_leases_outdated_items(sleep, time):
policy = create_policy()
policy._consumer._stopped.clear()
# Add these items at the beginning of the timeline
time.return_value = 0
policy.lease([base.LeaseRequest(ack_id='ack1', byte_size=50)])
# Add another item at towards end of the timeline
time.return_value = policy.flow_control.max_lease_duration - 1
policy.lease([base.LeaseRequest(ack_id='ack2', byte_size=50)])
# Now make sure time reports that we are at the end of our timeline.
time.return_value = policy.flow_control.max_lease_duration + 1
# Mock the sleep object.
def trigger_inactive(seconds):
assert 0 < seconds < 10
policy._consumer._stopped.set()
sleep.side_effect = trigger_inactive
# Also mock the consumer, which sends the request.
with mock.patch.object(policy._consumer, 'send_request') as send:
policy.maintain_leases()
# Only ack2 should be renewed. ack1 should've been dropped
send.assert_called_once_with(
types.StreamingPullRequest(
modify_deadline_ack_ids=['ack2'],
modify_deadline_seconds=[10],
))
assert len(policy.leased_messages) == 1
sleep.assert_called()
def _get_initial_request(self):
"""Return the initial request for the RPC.
This defines the initial request that must always be sent to Pub/Sub
immediately upon opening the subscription.
Returns:
google.cloud.pubsub_v1.types.StreamingPullRequest: A request
suitable for being the first request on the stream (and not
suitable for any other purpose).
"""
# Any ack IDs that are under lease management need to have their
# deadline extended immediately.
if self._leaser is not None:
# Explicitly copy the list, as it could be modified by another
# thread.
lease_ids = list(self._leaser.ack_ids)
else:
lease_ids = []
# Put the request together.
request = types.StreamingPullRequest(
modify_deadline_ack_ids=list(lease_ids),
modify_deadline_seconds=[self.ack_deadline] * len(lease_ids),
stream_ack_deadline_seconds=self.ack_histogram.percentile(99),
subscription=self._subscription,
)
# Return the initial request.
return request
def ack(self, items):
"""Acknowledge the given messages.
Args:
items(Sequence[AckRequest]): The items to acknowledge.
"""
# If we got timing information, add it to the histogram.
for item in items:
time_to_ack = item.time_to_ack
if time_to_ack is not None:
self.histogram.add(int(time_to_ack))
ack_ids = [item.ack_id for item in items]
if self._consumer.active:
# Send the request to ack the message.
request = types.StreamingPullRequest(ack_ids=ack_ids)
self._consumer.send_request(request)
else:
# If the consumer is inactive, then queue the ack_ids here; it
# will be acked as part of the initial request when the consumer
# is started again.
self._ack_on_resume.update(ack_ids)
# Remove the message from lease management.
self.drop(items)
def get_initial_request(self):
"""Return the initial request.
This defines the initial request that must always be sent to Pub/Sub
immediately upon opening the subscription.
Returns:
google.cloud.pubsub_v1.types.StreamingPullRequest: A request
suitable for being the first request on the stream (and not
suitable for any other purpose).
"""
# Any ack IDs that are under lease management and not being acked
# need to have their deadline extended immediately.
lease_ids = set(self.leased_messages.keys())
# Exclude any IDs that we're about to ack.
lease_ids = lease_ids.difference(self._ack_on_resume)
# Put the request together.
request = types.StreamingPullRequest(
ack_ids=list(self._ack_on_resume),
modify_deadline_ack_ids=list(lease_ids),
modify_deadline_seconds=[self.ack_deadline] * len(lease_ids),
stream_ack_deadline_seconds=self.histogram.percentile(99),
subscription=self.subscription,
)
# Clear the ack_ids set.
self._ack_on_resume.clear()
# Return the initial request.
return request
def ack(self, ack_id, time_to_ack=None, byte_size=None):
"""Acknowledge the message corresponding to the given ack_id.
Args:
ack_id (str): The ack ID.
time_to_ack (int): The time it took to ack the message, measured
from when it was received from the subscription. This is used
to improve the automatic ack timing.
byte_size (int): The size of the PubSub message, in bytes.
"""
# If we got timing information, add it to the histogram.
if time_to_ack is not None:
self.histogram.add(int(time_to_ack))
# Send the request to ack the message.
# However, if the consumer is inactive, then queue the ack_id here
# instead; it will be acked as part of the initial request when the
# consumer is started again.
if self._consumer.active:
request = types.StreamingPullRequest(ack_ids=[ack_id])
self._consumer.send_request(request)
else:
self._ack_on_resume.add(ack_id)
# Remove the message from lease management.
self.drop(ack_id=ack_id, byte_size=byte_size)
def test_ack_no_time():
policy = create_policy()
policy._consumer.stopped.clear()
with mock.patch.object(policy._consumer, 'send_request') as send_request:
policy.ack('ack_id_string')
send_request.assert_called_once_with(
types.StreamingPullRequest(ack_ids=['ack_id_string'], ))
assert len(policy.histogram) == 0
def test_heartbeat():
manager = make_manager()
manager._rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
manager._rpc.is_active = True
manager.heartbeat()
manager._rpc.send.assert_called_once_with(types.StreamingPullRequest())
def heartbeat(self):
"""Sends an empty request over the streaming pull RPC.
This always sends over the stream, regardless of if
``self._UNARY_REQUESTS`` is set or not.
"""
if self._rpc is not None and self._rpc.is_active:
self._rpc.send(types.StreamingPullRequest())
def test_modify_ack_deadline():
policy = create_policy()
with mock.patch.object(policy._consumer, 'send_request') as send_request:
policy.modify_ack_deadline('ack_id_string', 60)
send_request.assert_called_once_with(types.StreamingPullRequest(
modify_deadline_ack_ids=['ack_id_string'],
modify_deadline_seconds=[60],
))
def test_send_unary_empty():
manager = make_manager()
manager._UNARY_REQUESTS = True
manager.send(types.StreamingPullRequest())
manager._client.acknowledge.assert_not_called()
manager._client.modify_ack_deadline.assert_not_called()
def test_ack():
policy = create_policy()
policy._consumer.active = True
with mock.patch.object(policy._consumer, 'send_request') as send_request:
policy.ack('ack_id_string', 20)
send_request.assert_called_once_with(
types.StreamingPullRequest(ack_ids=['ack_id_string'], ))
assert len(policy.histogram) == 1
assert 20 in policy.histogram
def test_modify_ack_deadline():
manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [requests.ModAckRequest(ack_id="ack_id_string", seconds=60)]
dispatcher_.modify_ack_deadline(items)
manager.send.assert_called_once_with(
types.StreamingPullRequest(modify_deadline_ack_ids=["ack_id_string"],
modify_deadline_seconds=[60]))
def maintain_leases(self):
"""Maintain all of the leases being managed by the policy.
This method modifies the ack deadline for all of the managed
ack IDs, then waits for most of that time (but with jitter), and
then calls itself.
.. warning::
This method blocks, and generally should be run in a separate
thread or process.
Additionally, you should not have to call this method yourself,
unless you are implementing your own policy. If you are
implementing your own policy, you _should_ call this method
in an appropriate form of subprocess.
"""
while True:
# Sanity check: Should this infinite loop quit?
if not self._consumer.active:
_LOGGER.debug('Consumer inactive, ending lease maintenance.')
return
# Determine the appropriate duration for the lease. This is
# based off of how long previous messages have taken to ack, with
# a sensible default and within the ranges allowed by Pub/Sub.
p99 = self.histogram.percentile(99)
_LOGGER.debug('The current p99 value is %d seconds.', p99)
# Create a streaming pull request.
# We do not actually call `modify_ack_deadline` over and over
# because it is more efficient to make a single request.
ack_ids = list(self.managed_ack_ids)
_LOGGER.debug('Renewing lease for %d ack IDs.', len(ack_ids))
if ack_ids:
request = types.StreamingPullRequest(
modify_deadline_ack_ids=ack_ids,
modify_deadline_seconds=[p99] * len(ack_ids),
)
# NOTE: This may not work as expected if ``consumer.active``
# has changed since we checked it. An implementation
# without any sort of race condition would require a
# way for ``send_request`` to fail when the consumer
# is inactive.
self._consumer.send_request(request)
# Now wait an appropriate period of time and do this again.
#
# We determine the appropriate period of time based on a random
# period between 0 seconds and 90% of the lease. This use of
# jitter (http://bit.ly/2s2ekL7) helps decrease contention in cases
# where there are many clients.
snooze = random.uniform(0.0, p99 * 0.9)
_LOGGER.debug('Snoozing lease management for %f seconds.', snooze)
time.sleep(snooze)
def test_nack():
manager = mock.create_autospec(
streaming_pull_manager.StreamingPullManager, instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [requests.NackRequest(ack_id='ack_id_string', byte_size=10)]
dispatcher_.nack(items)
manager.send.assert_called_once_with(types.StreamingPullRequest(
modify_deadline_ack_ids=['ack_id_string'],
modify_deadline_seconds=[0],
))
def test_send_unary_api_call_error(caplog):
caplog.set_level(logging.DEBUG)
manager = make_manager()
manager._UNARY_REQUESTS = True
error = exceptions.GoogleAPICallError("The front fell off")
manager._client.acknowledge.side_effect = error
manager.send(types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))
assert "The front fell off" in caplog.text
def modify_ack_deadline(self, ack_id, seconds):
"""Modify the ack deadline for the given ack_id.
Args:
ack_id (str): The ack ID
seconds (int): The number of seconds to set the new deadline to.
"""
request = types.StreamingPullRequest(
modify_deadline_ack_ids=[ack_id],
modify_deadline_seconds=[seconds],
)
self._consumer.send_request(request)
def modify_ack_deadline(self, items):
"""Modify the ack deadline for the given messages.
Args:
items(Sequence[ModAckRequest]): The items to modify.
"""
ack_ids = [item.ack_id for item in items]
seconds = [item.seconds for item in items]
request = types.StreamingPullRequest(modify_deadline_ack_ids=ack_ids,
modify_deadline_seconds=seconds)
self._manager.send(request)
def test_ack():
policy = create_policy()
policy._consumer._stopped.clear()
with mock.patch.object(policy._consumer, 'send_request') as send_request:
policy.ack([
base.AckRequest(ack_id='ack_id_string',
time_to_ack=20,
byte_size=0)
])
send_request.assert_called_once_with(
types.StreamingPullRequest(ack_ids=['ack_id_string'], ))
assert len(policy.histogram) == 1
assert 20 in policy.histogram
def test_ack_no_time():
manager = mock.create_autospec(
streaming_pull_manager.StreamingPullManager, instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [requests.AckRequest(
ack_id='ack_id_string', byte_size=0, time_to_ack=None)]
dispatcher_.ack(items)
manager.send.assert_called_once_with(types.StreamingPullRequest(
ack_ids=['ack_id_string'],
))
manager.ack_histogram.add.assert_not_called()
def test_send_unary_retry_error(caplog):
caplog.set_level(logging.DEBUG)
manager, _, _, _, _, _ = make_running_manager()
manager._UNARY_REQUESTS = True
error = exceptions.RetryError("Too long a transient error",
cause=Exception("Out of time!"))
manager._client.acknowledge.side_effect = error
with pytest.raises(exceptions.RetryError):
manager.send(
types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))
assert "RetryError while sending unary RPC" in caplog.text
assert "signaled streaming pull manager shutdown" in caplog.text
def get_initial_request(self, ack_queue=False):
"""Return the initial request.
This defines the initial request that must always be sent to Pub/Sub
immediately upon opening the subscription.
Args:
ack_queue (bool): Whether to include any acks that were sent
while the connection was paused.
Returns:
~.pubsub_v1.types.StreamingPullRequest: A request suitable
for being the first request on the stream (and not suitable
for any other purpose).
.. note::
If ``ack_queue`` is set to True, this includes the ack_ids, but
also clears the internal set.
This means that calls to :meth:`get_initial_request` with
``ack_queue`` set to True are not idempotent.
"""
# Any ack IDs that are under lease management and not being acked
# need to have their deadline extended immediately.
ack_ids = set()
lease_ids = self.managed_ack_ids
if ack_queue:
ack_ids = self._ack_on_resume
lease_ids = lease_ids.difference(ack_ids)
# Put the request together.
request = types.StreamingPullRequest(
ack_ids=list(ack_ids),
modify_deadline_ack_ids=list(lease_ids),
modify_deadline_seconds=[self.ack_deadline] * len(lease_ids),
stream_ack_deadline_seconds=self.histogram.percentile(99),
subscription=self.subscription,
)
# Clear the ack_ids set.
# Note: If `ack_queue` is False, this just ends up being a no-op,
# since the set is just an empty set.
ack_ids.clear()
# Return the initial request.
return request
def ack(self, items):
"""Acknowledge the given messages.
Args:
items(Sequence[AckRequest]): The items to acknowledge.
"""
# If we got timing information, add it to the histogram.
for item in items:
time_to_ack = item.time_to_ack
if time_to_ack is not None:
self._manager.ack_histogram.add(time_to_ack)
ack_ids = [item.ack_id for item in items]
request = types.StreamingPullRequest(ack_ids=ack_ids)
self._manager.send(request)
# Remove the message from lease management.
self.drop(items)
def test_ack():
manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [
requests.AckRequest(ack_id="ack_id_string",
byte_size=0,
time_to_ack=20)
]
dispatcher_.ack(items)
manager.send.assert_called_once_with(
types.StreamingPullRequest(ack_ids=["ack_id_string"]))
manager.leaser.remove.assert_called_once_with(items)
manager.maybe_resume_consumer.assert_called_once()
manager.ack_histogram.add.assert_called_once_with(20)
def test_request_generator_thread():
consumer = create_consumer()
generator = consumer._request_generator_thread()
# The first request that comes from the request generator thread
# should always be the initial request.
initial_request = next(generator)
assert initial_request.subscription == 'sub_name_e'
assert initial_request.stream_ack_deadline_seconds == 10
# Subsequent requests correspond to items placed in the request queue.
consumer.send_request(types.StreamingPullRequest(ack_ids=['i']))
request = next(generator)
assert request.ack_ids == ['i']
# The poison pill should stop the loop.
consumer.send_request(_helper_threads.STOP)
with pytest.raises(StopIteration):
next(generator)
def modify_ack_deadline(self, items):
"""Modify the ack deadline for the given messages.
Args:
items(Sequence[ModAckRequest]): The items to modify.
"""
# We must potentially split the request into multiple smaller requests
# to avoid the server-side max request size limit.
ack_ids = (item.ack_id for item in items)
seconds = (item.seconds for item in items)
total_chunks = int(math.ceil(len(items) / _ACK_IDS_BATCH_SIZE))
for _ in range(total_chunks):
request = types.StreamingPullRequest(
modify_deadline_ack_ids=itertools.islice(
ack_ids, _ACK_IDS_BATCH_SIZE),
modify_deadline_seconds=itertools.islice(
seconds, _ACK_IDS_BATCH_SIZE),
)
self._manager.send(request)
def test_maintain_leases_ack_ids():
policy = create_policy()
policy._consumer.active = True
policy.lease('my ack id', 50)
# Mock the sleep object.
with mock.patch.object(time, 'sleep', autospec=True) as sleep:
def trigger_inactive(seconds):
assert 0 < seconds < 10
policy._consumer.active = False
sleep.side_effect = trigger_inactive
# Also mock the consumer, which sends the request.
with mock.patch.object(policy._consumer, 'send_request') as send:
policy.maintain_leases()
send.assert_called_once_with(types.StreamingPullRequest(
modify_deadline_ack_ids=['my ack id'],
modify_deadline_seconds=[10],
))
sleep.assert_called()
def _get_initial_request(self, stream_ack_deadline_seconds):
"""Return the initial request for the RPC.
This defines the initial request that must always be sent to Pub/Sub
immediately upon opening the subscription.
Args:
stream_ack_deadline_seconds (int):
The default message acknowledge deadline for the stream.
Returns:
google.cloud.pubsub_v1.types.StreamingPullRequest: A request
suitable for being the first request on the stream (and not
suitable for any other purpose).
"""
# Any ack IDs that are under lease management need to have their
# deadline extended immediately.
if self._leaser is not None:
# Explicitly copy the list, as it could be modified by another
# thread.
lease_ids = list(self._leaser.ack_ids)
else:
lease_ids = []
# Put the request together.
request = types.StreamingPullRequest(
modify_deadline_ack_ids=list(lease_ids),
modify_deadline_seconds=[self.ack_deadline] * len(lease_ids),
stream_ack_deadline_seconds=stream_ack_deadline_seconds,
subscription=self._subscription,
client_id=self._client_id,
max_outstanding_messages=self._flow_control.max_messages,
max_outstanding_bytes=self._flow_control.max_bytes,
)
# Return the initial request.
return request
def ack(self, items):
"""Acknowledge the given messages.
Args:
items(Sequence[AckRequest]): The items to acknowledge.
"""
# If we got timing information, add it to the histogram.
for item in items:
time_to_ack = item.time_to_ack
if time_to_ack is not None:
self._manager.ack_histogram.add(time_to_ack)
# We must potentially split the request into multiple smaller requests
# to avoid the server-side max request size limit.
ack_ids = (item.ack_id for item in items)
total_chunks = int(math.ceil(len(items) / _ACK_IDS_BATCH_SIZE))
for _ in range(total_chunks):
request = types.StreamingPullRequest(
ack_ids=itertools.islice(ack_ids, _ACK_IDS_BATCH_SIZE))
self._manager.send(request)
# Remove the message from lease management.
self.drop(items)
def test_send_request():
consumer = _consumer.Consumer()
request = types.StreamingPullRequest(subscription='foo')
with mock.patch.object(queue.Queue, 'put') as put:
consumer.send_request(request)
put.assert_called_once_with(request)
