Title: Workload Identity stops working when using gevent monkey patching
Tags: python,google-cloud-platform,google-cloud-pubsub,gevent,workload-identity

Question: I have an application which includes a Flask API server and a worker to process messages from PubSub. These run as separate containers on separate pods in Kubernetes.
I've migrated to using Workload Identity, previously I'd mount the service account's key file and set GOOGLE_APPLICATION_CREDENTIALS. However, the call to PubSub throws an error when using Workload Identity.
A key factor appears to be the call to monkey.patch_all() from gevent.
Below is a reproducible example when run on a container using Workload Identity:
from gevent import monkey
monkey.patch_all()

from google.cloud import pubsub_v1
client = pubsub_v1.SubscriberClient()
resp = client.pull(request={"subscription": "projects/abc/subscriptions/xyz", "max_messages": 1, "return_immediately": True})

Which results in:
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/venv/lib/python3.8/site-packages/google/cloud/pubsub_v1/_gapic.py", line 40, in <lambda>
    fx = lambda self, *a, **kw: wrapped_fx(self.api, *a, **kw)  # noqa
  File "/opt/venv/lib/python3.8/site-packages/google/pubsub_v1/services/subscriber/client.py", line 1131, in pull
    response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,)
  File "/opt/venv/lib/python3.8/site-packages/google/api_core/gapic_v1/method.py", line 154, in __call__
    return wrapped_func(*args, **kwargs)
  File "/opt/venv/lib/python3.8/site-packages/google/api_core/retry.py", line 283, in retry_wrapped_func
    return retry_target(
  File "/opt/venv/lib/python3.8/site-packages/google/api_core/retry.py", line 190, in retry_target
    return target()
  File "/opt/venv/lib/python3.8/site-packages/google/api_core/grpc_helpers.py", line 72, in error_remapped_callable
    return callable_(*args, **kwargs)
  File "/opt/venv/lib/python3.8/site-packages/grpc/_channel.py", line 944, in __call__
    state, call, = self._blocking(request, timeout, metadata, credentials,
  File "/opt/venv/lib/python3.8/site-packages/grpc/_channel.py", line 933, in _blocking
    event = call.next_event()
  File "src/python/grpcio/grpc/_cython/_cygrpc/channel.pyx.pxi", line 338, in grpc._cython.cygrpc.SegregatedCall.next_event
  File "src/python/grpcio/grpc/_cython/_cygrpc/channel.pyx.pxi", line 169, in grpc._cython.cygrpc._next_call_event
  File "src/python/grpcio/grpc/_cython/_cygrpc/channel.pyx.pxi", line 163, in grpc._cython.cygrpc._next_call_event
  File "src/python/grpcio/grpc/_cython/_cygrpc/completion_queue.pyx.pxi", line 63, in grpc._cython.cygrpc._latent_event
  File "src/python/grpcio/grpc/_cython/_cygrpc/credentials.pyx.pxi", line 62, in grpc._cython.cygrpc._get_metadata
RuntimeError: cannot exit context: thread state references a different context object

Any idea why the monkey.patch_all() from gevent breaks this when using Workload Identity and not a key file? Also how could I fix this but keep monkey.patch_all()?

Answer: Good day sir
The problem you're encountering is due to a conflict between gevent and gRPC, which is used by the Google Cloud Pub/Sub client. It looks like the problem is caused by the monkey patching that gevent does.
The problem is caused by the fact that gRPC uses thread-local storage (TLS) to handle context. When gevent monkey patches the threading module, it messes up gRPC's TLS implementation, which is why you're getting an error.
Try doing
monkey.patch_all(thread=False)