def open(self, callback):
        """Begin consuming messages.

        Args:
            callback (Callable[None, google.cloud.pubsub_v1.message.Messages]):
                A callback that will be called for each message received on the
                stream.
        """
        if self.is_active:
            raise ValueError('This manager is already open.')

        if self._closed:
            raise ValueError(
                'This manager has been closed and can not be re-used.')

        self._callback = functools.partial(_wrap_callback_errors, callback)

        # Start the thread to pass the requests.
        self._dispatcher = dispatcher.Dispatcher(self, self._scheduler.queue)
        self._dispatcher.start()

        # Start consuming messages.
        self._rpc = bidi.ResumableBidiRpc(
            start_rpc=self._client.api.streaming_pull,
            initial_request=self._get_initial_request,
            should_recover=self._should_recover)
        self._rpc.add_done_callback(self._on_rpc_done)
        self._consumer = bidi.BackgroundConsumer(
            self._rpc, self._on_response)
        self._consumer.start()

        # Start the lease maintainer thread.
        self._leaser = leaser.Leaser(self)
        self._leaser.start()
Example #2
0
def test_modify_ack_deadline_splitting_large_payload():
    manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
                                   instance=True)
    dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)

    items = [
        # use realistic lengths for ACK IDs (max 176 bytes)
        requests.ModAckRequest(ack_id=str(i).zfill(176), seconds=60)
        for i in range(5001)
    ]
    dispatcher_.modify_ack_deadline(items)

    calls = manager.send.call_args_list
    assert len(calls) == 3

    all_ack_ids = {item.ack_id for item in items}
    sent_ack_ids = collections.Counter()

    for call in calls:
        message = call.args[0]
        assert message._pb.ByteSize() <= 524288  # server-side limit (2**19)
        sent_ack_ids.update(message.modify_deadline_ack_ids)

    assert set(
        sent_ack_ids) == all_ack_ids  # all messages should have been MODACK-ed
    assert sent_ack_ids.most_common(
        1)[0][1] == 1  # each message MODACK-ed exactly once
Example #3
0
    def open(self, callback, on_callback_error):
        """Begin consuming messages.

        Args:
            callback (Callable[None, google.cloud.pubsub_v1.message.Message]):
                A callback that will be called for each message received on the
                stream.
            on_callback_error (Callable[Exception]):
                A callable that will be called if an exception is raised in
                the provided `callback`.
        """
        if self.is_active:
            raise ValueError("This manager is already open.")

        if self._closed:
            raise ValueError("This manager has been closed and can not be re-used.")

        self._callback = functools.partial(
            _wrap_callback_errors, callback, on_callback_error
        )

        # Create the RPC
        stream_ack_deadline_seconds = self.ack_histogram.percentile(99)

        get_initial_request = functools.partial(
            self._get_initial_request, stream_ack_deadline_seconds
        )
        self._rpc = bidi.ResumableBidiRpc(
            start_rpc=self._client.api.streaming_pull,
            initial_request=get_initial_request,
            should_recover=self._should_recover,
            should_terminate=self._should_terminate,
            throttle_reopen=True,
        )
        self._rpc.add_done_callback(self._on_rpc_done)

        _LOGGER.debug(
            "Creating a stream, default ACK deadline set to {} seconds.".format(
                stream_ack_deadline_seconds
            )
        )

        # Create references to threads
        self._dispatcher = dispatcher.Dispatcher(self, self._scheduler.queue)
        self._consumer = bidi.BackgroundConsumer(self._rpc, self._on_response)
        self._leaser = leaser.Leaser(self)
        self._heartbeater = heartbeater.Heartbeater(self)

        # Start the thread to pass the requests.
        self._dispatcher.start()

        # Start consuming messages.
        self._consumer.start()

        # Start the lease maintainer thread.
        self._leaser.start()

        # Start the stream heartbeater thread.
        self._heartbeater.start()
def test_dispatch_callback(item, method):
    subscriber_ = mock.create_autospec(subscriber.Subscriber, instance=True)
    dispatcher_ = dispatcher.Dispatcher(mock.sentinel.queue, subscriber_)

    items = [item]
    dispatcher_.dispatch_callback(items)

    getattr(subscriber_, method).assert_called_once_with([item])
def test_dispatch_callback_inactive():
    subscriber_ = mock.create_autospec(subscriber.Subscriber, instance=True)
    subscriber_.is_active = False
    dispatcher_ = dispatcher.Dispatcher(mock.sentinel.queue, subscriber_)

    dispatcher_.dispatch_callback([requests.AckRequest(0, 0, 0)])

    subscriber_.ack.assert_not_called()
def test_start_already_started(thread):
    subscriber_ = mock.create_autospec(subscriber.Subscriber, instance=True)
    dispatcher_ = dispatcher.Dispatcher(mock.sentinel.queue, subscriber_)
    dispatcher_._thread = mock.sentinel.thread

    with pytest.raises(ValueError):
        dispatcher_.start()

    thread.assert_not_called()
Example #7
0
def test_dispatch_callback_inactive():
    manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
                                   instance=True)
    manager.is_active = False
    dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)

    dispatcher_.dispatch_callback([requests.AckRequest(0, 0, 0, "")])

    manager.send.assert_not_called()
Example #8
0
def test_drop():
    manager = mock.create_autospec(
        streaming_pull_manager.StreamingPullManager, instance=True)
    dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)

    items = [requests.DropRequest(ack_id='ack_id_string', byte_size=10)]
    dispatcher_.drop(items)

    manager.leaser.remove.assert_called_once_with(items)
    manager.maybe_resume_consumer.assert_called_once()
Example #9
0
def test_start_already_started(thread):
    manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
                                   instance=True)
    dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
    dispatcher_._thread = mock.sentinel.thread

    with pytest.raises(ValueError):
        dispatcher_.start()

    thread.assert_not_called()
Example #10
0
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()
Example #11
0
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]))
Example #12
0
def test_stop():
    queue_ = queue.Queue()
    dispatcher_ = dispatcher.Dispatcher(mock.sentinel.manager, queue_)
    thread = mock.create_autospec(threading.Thread, instance=True)
    dispatcher_._thread = thread

    dispatcher_.stop()

    assert queue_.get() is helper_threads.STOP
    thread.join.assert_called_once()
    assert dispatcher_._thread is None
Example #13
0
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])
Example #14
0
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_start(thread):
    subscriber_ = mock.create_autospec(subscriber.Subscriber, instance=True)
    dispatcher_ = dispatcher.Dispatcher(mock.sentinel.queue, subscriber_)

    dispatcher_.start()

    thread.assert_called_once_with(
        name=dispatcher._CALLBACK_WORKER_NAME, target=mock.ANY)

    thread.return_value.start.assert_called_once()

    assert dispatcher_._thread is not None
Example #16
0
def test_start(thread):
    manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
                                   instance=True)
    dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)

    dispatcher_.start()

    thread.assert_called_once_with(name=dispatcher._CALLBACK_WORKER_NAME,
                                   target=mock.ANY)

    thread.return_value.start.assert_called_once()

    assert dispatcher_._thread is not None
Example #17
0
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_drop_unordered_messages():
    manager = mock.create_autospec(
        streaming_pull_manager.StreamingPullManager, instance=True
    )
    dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)

    items = [
        requests.DropRequest(ack_id="ack_id_string", byte_size=10, ordering_key="")
    ]
    dispatcher_.drop(items)

    manager.leaser.remove.assert_called_once_with(items)
    assert list(manager.activate_ordering_keys.call_args.args[0]) == []
    manager.maybe_resume_consumer.assert_called_once()
Example #19
0
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_stop_no_join():
    dispatcher_ = dispatcher.Dispatcher(
        mock.sentinel.queue, mock.sentinel.subscriber)

    dispatcher_.stop()
Example #21
0
def test_stop_no_join():
    dispatcher_ = dispatcher.Dispatcher(mock.sentinel.manager,
                                        mock.sentinel.queue)

    dispatcher_.stop()
Example #22
0
    def open(self, callback, on_callback_error):
        """Begin consuming messages.

        Args:
            callback (Callable[None, google.cloud.pubsub_v1.message.Message]):
                A callback that will be called for each message received on the
                stream.
            on_callback_error (Callable[Exception]):
                A callable that will be called if an exception is raised in
                the provided `callback`.
        """
        if self.is_active:
            raise ValueError("This manager is already open.")

        if self._closed:
            raise ValueError(
                "This manager has been closed and can not be re-used.")

        self._callback = functools.partial(_wrap_callback_errors, callback,
                                           on_callback_error)

        # Create the RPC

        # We must use a fixed value for the ACK deadline, as we cannot read it
        # from the subscription. The latter would require `pubsub.subscriptions.get`
        # permission, which is not granted to the default subscriber role
        # `roles/pubsub.subscriber`.
        # See also https://github.com/googleapis/google-cloud-python/issues/9339
        #
        # When dynamic lease management is enabled for the "on hold" messages,
        # the default stream ACK deadline should again be set based on the
        # historic ACK timing data, i.e. `self.ack_histogram.percentile(99)`.
        stream_ack_deadline_seconds = _DEFAULT_STREAM_ACK_DEADLINE

        get_initial_request = functools.partial(self._get_initial_request,
                                                stream_ack_deadline_seconds)
        self._rpc = bidi.ResumableBidiRpc(
            start_rpc=self._client.api.streaming_pull,
            initial_request=get_initial_request,
            should_recover=self._should_recover,
            should_terminate=self._should_terminate,
            throttle_reopen=True,
        )
        self._rpc.add_done_callback(self._on_rpc_done)

        _LOGGER.debug(
            "Creating a stream, default ACK deadline set to {} seconds.".
            format(stream_ack_deadline_seconds))

        # Create references to threads
        self._dispatcher = dispatcher.Dispatcher(self, self._scheduler.queue)
        self._consumer = bidi.BackgroundConsumer(self._rpc, self._on_response)
        self._leaser = leaser.Leaser(self)
        self._heartbeater = heartbeater.Heartbeater(self)

        # Start the thread to pass the requests.
        self._dispatcher.start()

        # Start consuming messages.
        self._consumer.start()

        # Start the lease maintainer thread.
        self._leaser.start()

        # Start the stream heartbeater thread.
        self._heartbeater.start()