Repository: googleapis/python-pubsub

Issue:
Possible race condition between `_on_response` and `close`

#### Environment details

  - OS type and version: Ubuntu 22.04
  - Python version: 3.11
  - pip version: 23.0.1
  - `google-cloud-pubsub` version: 2.18.4

#### Steps to reproduce

  1. Run code sample indefinitely
  2. Sometime it shows Assertion error in `_on_response` function.

#### Code example

```python
def on_subscribe(subscription, until=None):
    """Decorator factory that provides subscribed messages to function.
    It handle decorated function as callback. So message should be acked/nacked
    inside decorated function.

    Args:
        subscription (str): Subscription ID. 
            Should be `projects/{PROJECT_ID}/subscriptions/{SUBSCRIPTION_ID}`
        until (datetime.datetime): This function will subscribe messages
            published until this timestamp.
    """

    def _callback_factory(func, finished, subscribe_until, **kwargs):

        def _callback(message):
            """Callback function.
                
            It sends signal if subscribed all messages.
            """
            publish_time = datetime.fromtimestamp(
                message.publish_time.timestamp())
            if subscribe_until and publish_time <= subscribe_until:
                return func(message, **kwargs)
            if subscribe_until and not finished.is_set():
                logging.info('Subscribed all messages published until %s',
                             subscribe_until)
                finished.set()
                message.nack()

        return _callback

    def _wrapper(func):

        @functools.wraps(func)
        def _inner_wrapper(**kwargs):
            # Event variable that is triggered when all messages are subscribed
            all_subscribed = Event()

            callback = _callback_factory(func=func,
                                         finished=all_subscribed,
                                         subscribe_until=subscribe_until,
                                         **kwargs)

            # Ensure closing subscriber for memory leak prevention.
            with pubsub_v1.SubscriberClient() as subscriber:
                future = subscriber.subscribe(
                    subscription=subscription,
                    callback=callback,
                    await_callbacks_on_shutdown=True,
                    flow_control=pubsub_v1.types.FlowControl(max_messages=5000),
                )

                all_subscribed.wait(timeout=60)

                # Wait until future is finished when it's cancelled.
                # If it cancelled by timeout or keyboard interrupt, ignore it.
                try:
                    future.cancel()
                    future.result(timeout=60)
                except (KeyboardInterrupt, TimeoutError):
                    pass
                except Exception as e:
                    logging.error("Error occurs during subscription to %s: %s",
                                  subscription, str(e))

        return _inner_wrapper

    return _wrapper

@on_subscribe(subscription="SUBSCRIPTION")
def callback(message):
    # Do something with message
```

#### Stack trace
```
Traceback (most recent call last):
  File "/layers/google.python.pip/pip/lib/python3.11/site-packages/google/api_core/bidi.py", line 657, in _thread_main
    self._on_response(response)
  File "/layers/google.python.pip/pip/lib/python3.11/site-packages/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py", line 1107, in _on_response
    assert self._scheduler is not None
```

#### Explanation
It's because `future.cancel()` executes `manager.close()` which makes `_scheduler` as `None` and it makes `_on_response` raise AssertionError.

Maybe it has to be protected by threading lock somehow.

Discussion:
conventional-commit-lint-gcf[bot]: <!-- probot comment [1356997]-->
🤖 I detect that the PR title and the commit message differ and there's only one commit. To use the PR title for the commit history, you can use Github's automerge feature with squashing, or use `automerge` label. Good luck human!

 -- conventional-commit-lint bot
https://conventionalcommits.org/
mukund-ananthu: @parthea PTAL.
mukund-ananthu: @hongalex 

Patch:
diff --git a/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py b/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
index b8531db..c01dd7f 100644
--- a/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
+++ b/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
@@ -1104,8 +1104,11 @@ class StreamingPullManager(object):
         )
 
         with self._pause_resume_lock:
-            assert self._scheduler is not None
-            assert self._leaser is not None
+            if self._scheduler is None or self._leaser is None:
+                _LOGGER.debug(
+                    f"self._scheduler={self._scheduler} or self._leaser={self._leaser} is None. Stopping further processing."
+                )
+                return
 
             for received_message in received_messages:
                 if (


Test patch:
diff --git a/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
index 278f3e8..d4ce2cf 100644
--- a/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
+++ b/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
@@ -1375,6 +1375,62 @@ def test_close_blocking_scheduler_shutdown():
     scheduler.shutdown.assert_called_once_with(await_msg_callbacks=True)
 
 
+def test__on_response_none_scheduler():
+    manager, _, _, _, _, _ = make_running_manager()
+
+    manager._callback = mock.sentinel.callback
+    manager._scheduler = None
+    # Set up the messages.
+    response = gapic_types.StreamingPullResponse(
+        received_messages=[
+            gapic_types.ReceivedMessage(
+                ack_id="ack1",
+                message=gapic_types.PubsubMessage(data=b"foo", message_id="1"),
+            ),
+            gapic_types.ReceivedMessage(
+                ack_id="ack2",
+                message=gapic_types.PubsubMessage(data=b"bar", message_id="2"),
+                delivery_attempt=6,
+            ),
+        ]
+    )
+
+    manager._maybe_release_messages = mock.Mock()
+
+    # adjust message bookkeeping in leaser
+    fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=42)
+    manager._on_response(response)
+
+    manager._maybe_release_messages.assert_not_called
+
+
+def test__on_response_none_leaser():
+    manager, _, _, _, _, _ = make_running_manager()
+
+    manager._callback = mock.sentinel.callback
+    manager._leaser = None
+    # Set up the messages.
+    response = gapic_types.StreamingPullResponse(
+        received_messages=[
+            gapic_types.ReceivedMessage(
+                ack_id="ack1",
+                message=gapic_types.PubsubMessage(data=b"foo", message_id="1"),
+            ),
+            gapic_types.ReceivedMessage(
+                ack_id="ack2",
+                message=gapic_types.PubsubMessage(data=b"bar", message_id="2"),
+                delivery_attempt=6,
+            ),
+        ]
+    )
+
+    manager._maybe_release_messages = mock.Mock()
+
+    manager._on_response(response)
+
+    manager._maybe_release_messages.assert_not_called
+
+
 def test_close_nonblocking_scheduler_shutdown():
     manager, _, _, _, _, _ = make_running_manager(await_callbacks_on_shutdown=False)
     scheduler = manager._scheduler