def test_error_if_mesage_would_block_indefinitely():
settings = types.PublishFlowControl(
message_limit=0, # simulate non-sane settings
byte_limit=1,
limit_exceeded_behavior=types.LimitExceededBehavior.BLOCK,
)
flow_controller = FlowController(settings)
msg = grpc_types.PubsubMessage(data=b"xyz")
adding_done = threading.Event()
error_event = threading.Event()
_run_in_daemon(flow_controller.add, [msg],
adding_done,
error_event=error_event)
assert error_event.wait(
timeout=0.1), "No error on adding too large a message."
# Now that we know that an error occurs, we can check its type directly
# without the fear of blocking indefinitely.
flow_controller = FlowController(settings) # we want a fresh controller
with pytest.raises(exceptions.FlowControlLimitError) as error_info:
flow_controller.add(msg)
error_msg = str(error_info.value)
assert "would block forever" in error_msg
assert "messages: 1 / 0" in error_msg
assert "bytes: {} / 1".format(msg._pb.ByteSize()) in error_msg
def test_blocked_messages_are_accepted_in_fifo_order():
settings = types.PublishFlowControl(
message_limit=1,
byte_limit=1_000_000, # Unlimited for practical purposes in the test.
limit_exceeded_behavior=types.LimitExceededBehavior.BLOCK,
)
flow_controller = FlowController(settings)
# It's OK if the message instance is shared, as flow controlelr is only concerned
# with byte sizes and counts, and not with particular message instances.
message = grpc_types.PubsubMessage(data=b"x")
adding_done_events = [threading.Event() for _ in range(10)]
releasing_done_events = [threading.Event() for _ in adding_done_events]
# Add messages. The first one will be accepted, and the rest should queue behind.
for adding_done in adding_done_events:
_run_in_daemon(flow_controller.add, [message], adding_done)
time.sleep(0.1)
if not adding_done_events[0].wait(timeout=0.1): # pragma: NO COVER
pytest.fail("The first message unexpectedly got blocked on adding.")
# For each message, check that it has indeed been added to the flow controller.
# Then release it to make room for the next message in line, and repeat the check.
enumeration = enumerate(zip(adding_done_events, releasing_done_events))
for i, (adding_done, releasing_done) in enumeration:
if not adding_done.wait(timeout=0.1): # pragma: NO COVER
pytest.fail(f"Queued message still blocked on adding (i={i}).")
_run_in_daemon(flow_controller.release, [message], releasing_done)
if not releasing_done.wait(timeout=0.1): # pragma: NO COVER
pytest.fail(f"Queued message was not released in time (i={i}).")
def test_incorrectly_releasing_too_many_messages():
settings = types.PublishFlowControl(
message_limit=1,
byte_limit=150,
limit_exceeded_behavior=types.LimitExceededBehavior.ERROR,
)
flow_controller = FlowController(settings)
msg1 = grpc_types.PubsubMessage(data=b"x" * 100)
msg2 = grpc_types.PubsubMessage(data=b"y" * 100)
msg3 = grpc_types.PubsubMessage(data=b"z" * 100)
# Releasing a message that would make the load negative should result in a warning.
with warnings.catch_warnings(record=True) as warned:
flow_controller.release(msg1)
assert len(warned) == 1
assert issubclass(warned[0].category, RuntimeWarning)
warning_msg = str(warned[0].message)
assert "never added or already released" in warning_msg
# Incorrectly removing a message does not mess up internal stats, we can
# still only add a single message at a time to this flow.
flow_controller.add(msg2)
with pytest.raises(exceptions.FlowControlLimitError) as error:
flow_controller.add(msg3)
error_msg = str(error.value)
assert "messages: 2 / 1" in error_msg
total_size = msg2._pb.ByteSize() + msg3._pb.ByteSize()
expected_size_info = "bytes: {} / 150".format(total_size)
assert expected_size_info in error_msg
def test_threads_posting_large_messages_do_not_starve():
settings = types.PublishFlowControl(
message_limit=100,
byte_limit=110,
limit_exceeded_behavior=types.LimitExceededBehavior.BLOCK,
)
flow_controller = FlowController(settings)
large_msg = types.PubsubMessage(data=b"x" * 100) # close to entire byte limit
adding_initial_done = threading.Event()
adding_large_done = threading.Event()
adding_busy_done = threading.Event()
releasing_busy_done = threading.Event()
releasing_large_done = threading.Event()
# Occupy some of the flow capacity, then try to add a large message. Releasing
# enough messages should eventually allow the large message to come through, even
# if more messages are added after it (those should wait for the large message).
initial_messages = [types.PubsubMessage(data=b"x" * 10)] * 5
_run_in_daemon(flow_controller, "add", initial_messages, adding_initial_done)
assert adding_initial_done.wait(timeout=0.1)
_run_in_daemon(flow_controller, "add", [large_msg], adding_large_done)
# Continuously keep adding more messages after the large one.
messages = [types.PubsubMessage(data=b"x" * 10)] * 10
_run_in_daemon(flow_controller, "add", messages, adding_busy_done, action_pause=0.1)
# At the same time, gradually keep releasing the messages - the freeed up
# capacity should be consumed by the large message, not the other small messages
# being added after it.
_run_in_daemon(
flow_controller, "release", messages, releasing_busy_done, action_pause=0.1
)
# Sanity check - releasing should have completed by now.
if not releasing_busy_done.wait(timeout=1.1):
pytest.fail("Releasing messages blocked or errored.")
# Enough messages released, the large message should have come through in
# the meantime.
if not adding_large_done.wait(timeout=0.1):
pytest.fail("A thread adding a large message starved.")
if adding_busy_done.wait(timeout=0.1):
pytest.fail("Adding multiple small messages did not block.")
# Releasing the large message should unblock adding the remaining "busy" messages
# that have not been added yet.
_run_in_daemon(flow_controller, "release", [large_msg], releasing_large_done)
if not releasing_large_done.wait(timeout=0.1):
pytest.fail("Releasing a message blocked or errored.")
if not adding_busy_done.wait(timeout=1.0):
pytest.fail("Adding messages blocked or errored.")
def test_overflow_no_error_on_ignore():
settings = types.PublishFlowControl(
message_limit=1,
byte_limit=2,
limit_exceeded_behavior=types.LimitExceededBehavior.IGNORE,
)
flow_controller = FlowController(settings)
# there should be no overflow errors
flow_controller.add(grpc_types.PubsubMessage(data=b"foo"))
flow_controller.add(grpc_types.PubsubMessage(data=b"bar"))
def test_no_overflow_no_error():
settings = types.PublishFlowControl(
message_limit=100,
byte_limit=10000,
limit_exceeded_behavior=types.LimitExceededBehavior.ERROR,
)
flow_controller = FlowController(settings)
# there should be no errors
for data in (b"foo", b"bar", b"baz"):
msg = grpc_types.PubsubMessage(data=data)
flow_controller.add(msg)
def test_message_count_overflow_error():
settings = types.PublishFlowControl(
message_limit=1,
byte_limit=10000,
limit_exceeded_behavior=types.LimitExceededBehavior.ERROR,
)
flow_controller = FlowController(settings)
flow_controller.add(grpc_types.PubsubMessage(data=b"foo"))
with pytest.raises(exceptions.FlowControlLimitError) as error:
flow_controller.add(grpc_types.PubsubMessage(data=b"bar"))
assert "messages: 2 / 1" in str(error.value)
def test_warning_on_internal_reservation_stats_error_when_unblocking():
settings = types.PublishFlowControl(
message_limit=1,
byte_limit=150,
limit_exceeded_behavior=types.LimitExceededBehavior.BLOCK,
)
flow_controller = FlowController(settings)
msg1 = grpc_types.PubsubMessage(data=b"x" * 100)
msg2 = grpc_types.PubsubMessage(data=b"y" * 100)
# If there is a concurrency bug in FlowController, we do not want to block
# the main thread running the tests, thus we delegate all add/release
# operations to daemon threads and check the outcome (blocked/not blocked)
# through Events.
adding_1_done = threading.Event()
adding_2_done = threading.Event()
releasing_1_done = threading.Event()
# Adding a message with free capacity should not block.
_run_in_daemon(flow_controller.add, [msg1], adding_1_done)
if not adding_1_done.wait(timeout=0.1):
pytest.fail( # pragma: NO COVER
"Adding a message with enough flow capacity blocked or errored.")
# Adding messages when there is not enough capacity should block, even if
# added through multiple threads.
_run_in_daemon(flow_controller.add, [msg2], adding_2_done)
if adding_2_done.wait(timeout=0.1):
pytest.fail(
"Adding a message on overflow did not block.") # pragma: NO COVER
# Intentionally corrupt internal stats
reservation = next(iter(flow_controller._waiting.values()), None)
assert reservation is not None, "No messages blocked by flow controller."
reservation.bytes_reserved = reservation.bytes_needed + 1
with warnings.catch_warnings(record=True) as warned:
_run_in_daemon(flow_controller.release, [msg1], releasing_1_done)
if not releasing_1_done.wait(timeout=0.1):
pytest.fail(
"Releasing a message blocked or errored.") # pragma: NO COVER
matches = [
warning for warning in warned if warning.category is RuntimeWarning
]
assert len(matches) == 1
assert "too many bytes reserved" in str(matches[0].message).lower()
def test_publish_error_exceeding_flow_control_limits(creds):
publisher_options = types.PublisherOptions(
flow_control=types.PublishFlowControl(
message_limit=10,
byte_limit=150,
limit_exceeded_behavior=types.LimitExceededBehavior.ERROR,
))
client = publisher.Client(credentials=creds,
publisher_options=publisher_options)
mock_batch = mock.Mock(spec=client._batch_class)
topic = "topic/path"
client._set_batch(topic, mock_batch)
future1 = client.publish(topic, b"a" * 100)
future2 = client.publish(topic, b"b" * 100)
future1.result() # no error, still within flow control limits
with pytest.raises(exceptions.FlowControlLimitError):
future2.result()
def test_no_error_on_moderate_message_flow():
settings = types.PublishFlowControl(
message_limit=2,
byte_limit=250,
limit_exceeded_behavior=types.LimitExceededBehavior.ERROR,
)
flow_controller = FlowController(settings)
msg1 = grpc_types.PubsubMessage(data=b"x" * 100)
msg2 = grpc_types.PubsubMessage(data=b"y" * 100)
msg3 = grpc_types.PubsubMessage(data=b"z" * 100)
# The flow control settings will accept two in-flight messages, but not three.
# If releasing messages works correctly, the sequence below will not raise errors.
flow_controller.add(msg1)
flow_controller.add(msg2)
flow_controller.release(msg1)
flow_controller.add(msg3)
flow_controller.release(msg2)
flow_controller.release(msg3)
def test_byte_size_overflow_error():
settings = types.PublishFlowControl(
message_limit=10000,
byte_limit=199,
limit_exceeded_behavior=types.LimitExceededBehavior.ERROR,
)
flow_controller = FlowController(settings)
# Since the message data itself occupies 100 bytes, it means that both
# messages combined will exceed the imposed byte limit of 199, but a single
# message will not (the message size overhead is way lower than data size).
msg1 = grpc_types.PubsubMessage(data=b"x" * 100)
msg2 = grpc_types.PubsubMessage(data=b"y" * 100)
flow_controller.add(msg1)
with pytest.raises(exceptions.FlowControlLimitError) as error:
flow_controller.add(msg2)
total_size = msg1._pb.ByteSize() + msg2._pb.ByteSize()
expected_info = "bytes: {} / 199".format(total_size)
assert expected_info in str(error.value)
def test_rejected_messages_do_not_increase_total_load():
settings = types.PublishFlowControl(
message_limit=1,
byte_limit=150,
limit_exceeded_behavior=types.LimitExceededBehavior.ERROR,
)
flow_controller = FlowController(settings)
msg1 = grpc_types.PubsubMessage(data=b"x" * 100)
msg2 = grpc_types.PubsubMessage(data=b"y" * 100)
flow_controller.add(msg1)
for _ in range(5):
with pytest.raises(exceptions.FlowControlLimitError):
flow_controller.add(grpc_types.PubsubMessage(data=b"z" * 100))
# After releasing a message we should again be able to add another one, despite
# previously trying to add a lot of other messages.
flow_controller.release(msg1)
flow_controller.add(msg2)
def test_blocking_on_overflow_until_free_capacity():
settings = types.PublishFlowControl(
message_limit=1,
byte_limit=150,
limit_exceeded_behavior=types.LimitExceededBehavior.BLOCK,
)
flow_controller = FlowController(settings)
msg1 = grpc_types.PubsubMessage(data=b"x" * 100)
msg2 = grpc_types.PubsubMessage(data=b"y" * 100)
msg3 = grpc_types.PubsubMessage(data=b"z" * 100)
msg4 = grpc_types.PubsubMessage(data=b"w" * 100)
# If there is a concurrency bug in FlowController, we do not want to block
# the main thread running the tests, thus we delegate all add/release
# operations to daemon threads and check the outcome (blocked/not blocked)
# through Events.
adding_1_done = threading.Event()
adding_2_done = threading.Event()
adding_3_done = threading.Event()
adding_4_done = threading.Event()
releasing_1_done = threading.Event()
releasing_x_done = threading.Event()
# Adding a message with free capacity should not block.
_run_in_daemon(flow_controller.add, [msg1], adding_1_done)
if not adding_1_done.wait(timeout=0.1):
pytest.fail( # pragma: NO COVER
"Adding a message with enough flow capacity blocked or errored.")
# Adding messages when there is not enough capacity should block, even if
# added through multiple threads.
_run_in_daemon(flow_controller.add, [msg2], adding_2_done)
if adding_2_done.wait(timeout=0.1):
pytest.fail(
"Adding a message on overflow did not block.") # pragma: NO COVER
_run_in_daemon(flow_controller.add, [msg3], adding_3_done)
if adding_3_done.wait(timeout=0.1):
pytest.fail(
"Adding a message on overflow did not block.") # pragma: NO COVER
_run_in_daemon(flow_controller.add, [msg4], adding_4_done)
if adding_4_done.wait(timeout=0.1):
pytest.fail(
"Adding a message on overflow did not block.") # pragma: NO COVER
# After releasing one message, there should be room for a new message, which
# should result in unblocking one of the waiting threads.
_run_in_daemon(flow_controller.release, [msg1], releasing_1_done)
if not releasing_1_done.wait(timeout=0.1):
pytest.fail(
"Releasing a message blocked or errored.") # pragma: NO COVER
done_status = [
adding_2_done.wait(timeout=0.1),
adding_3_done.wait(timeout=0.1),
adding_4_done.wait(timeout=0.1),
]
# In sum() we use the fact that True==1 and False==0, and that Event.wait()
# returns False only if it times out, i.e. its internal flag has not been set.
done_count = sum(done_status)
assert done_count == 1, "Exactly one thread should have been unblocked."
# Release another message and verify that yet another thread gets unblocked.
added_msg = [msg2, msg3, msg4][done_status.index(True)]
_run_in_daemon(flow_controller.release, [added_msg], releasing_x_done)
if not releasing_x_done.wait(timeout=0.1):
pytest.fail(
"Releasing messages blocked or errored.") # pragma: NO COVER
released_count = sum((
adding_2_done.wait(timeout=0.1),
adding_3_done.wait(timeout=0.1),
adding_4_done.wait(timeout=0.1),
))
assert released_count == 2, "Exactly two threads should have been unblocked."
