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.

Hints from issue comments:
@liuyunnnn: Thank you for the detailed report. I'll try to reproduce this and fix accordingly. 
In the meanwhile, I'm wondering why you choose to cancel first and then wait for result here
```
future.cancel()
future.result(timeout=60)
```
 instead of a common use case as the [sample](https://cloud.google.com/pubsub/docs/pull#client_library_code_samples) shows. This could be a mitigation for the issue you are experiencing.@hson2: > Thank you for the detailed report. I'll try to reproduce this and fix accordingly. 
> 
> In the meanwhile, I'm wondering why you choose to cancel first and then wait for result here
> 
> ```
> 
> future.cancel()
> 
> future.result(timeout=60)
> 
> ```
> 
>  instead of a common use case as the [sample](https://cloud.google.com/pubsub/docs/pull#client_library_code_samples) shows. This could be a mitigation for the issue you are experiencing.

Because I wanted to subscribe messages until certain timestamp (which turns out wrong idea 😅) and make callback to signal event if it subscribed all messages by using threading.Event object.

If I call result first without timeout then it will wait indefinitely and if I set timeout then it will raise TimeoutError at second call of result. I wanted to just cancel it first to make it avoid raising error.
@chase-peach: We are also seeing this issue.

Python version: 3.11
pip version: 23.3.1
google-cloud-pubsub version: 2.14.1

We are running in Google Kubernetes Engine (GKE) and listening to the following two signals in our worker for scaling:

```python
import signal
signal.signal(signal.SIGINT, self.exit_gracefully)
signal.signal(signal.SIGTERM, self.exit_gracefully)
```

With a `exit_gracefully` implemented as follows:

```python
def exit_gracefully(self) -> None:
    """Stop the async worker."""
    self.future.cancel()  # Request shutdown
    try:
        self.future.result(timeout=60 * 5)  # Block until the shutdown is complete (or up to 5min).
    except TimeoutError:
        self.logger.warning("Stop timeout reached, Pod will die.")
    except Exception as e:
        self.logger.error(f"Error while shutting down worker: {e}")
```

with the following setup:

```python
self.subscriber = pubsub_v1.SubscriberClient()

flow_control = pubsub_v1.types.FlowControl(max_messages=self.pub_sub_flow_control_max_messages)
executor = futures.ThreadPoolExecutor(
    max_workers=self.pub_sub_thread_pool_executor_max_workers, thread_name_prefix="TPE-PeachWorker"
)
scheduler = pubsub_v1.subscriber.scheduler.ThreadScheduler(executor=executor)

self.future = self.subscriber.subscribe(
    self.subscription_path,
    callback=self.pubsub_callback,
    flow_control=flow_control,
    scheduler=scheduler,
    await_callbacks_on_shutdown=True,
)
```

We do our main work in a separate thread from the main thread to allow python to receive signals:

```python
# Execute the wait on the future for streaming pull in a different thread
# because Python's main thread is the only one that will receive signals such as
# SIGTERM for when a pod is requested to die so it can be replaced with a new one.
thread = Thread(target=self.process, args=(), daemon=True)
thread.start()
```

With the `process` method implemented as follows:

```python
def process(self):
    with self.subscriber:
        try:
            self.future.result()
        except Cancelled:
            self.logger.info("Cancel request received.")
```

We are seeing the same scheduler assertion error that @hson2 saw:

```
AssertionError: null
  File "threading.py", line 982, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/local/lib/python3.11/site-packages/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py", line 921, in _shutdown
    assert self._scheduler is not None
```

However, we are also seeing another error immediately before:

```
AttributeError: 'NoneType' object has no attribute 'nack'
  File "threading.py", line 982, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/local/lib/python3.11/site-packages/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py", line 948, in _shutdown
    msg.nack()
```

@liuyunnnn Could you expand on the preference to not call cancel first in a scenario like this?

Thank you for your help :)!


Solution 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 b8531db17..c01dd7f2e 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 @@ def _on_response(self, response: gapic_types.StreamingPullResponse) -> None:
         )
 
         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 278f3e88e..d4ce2cfdb 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