feat: port cert rotation streaming

[agrawalradhika-cell]

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

• Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
• Ensure the tests and linter pass
• Code coverage does not decrease (if any source code was changed)
• Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

[gemini-code-assist]

Code Review

This pull request introduces a mechanism for transparent mTLS credential rotation and retries in gRPC channels by implementing a call interceptor (_MTLSCallInterceptor), a refreshing channel wrapper (_MTLSRefreshingChannel), and several retryable call wrappers. Feedback on these changes highlights a critical race condition in _RetryableStreamUnaryFuture between the client thread calling result() and the gRPC callback thread, which can be resolved using a pending retry flag and a condition variable. Additionally, to prevent thread leaks, it is recommended to avoid creating a ThreadPoolExecutor per interceptor instance and instead use a daemon thread for asynchronous retries.

[gemini-code-assist]

There is a critical race condition between the client thread calling result() and the gRPC callback thread executing _on_inner_future_done. When the inner future fails, current_future.result() immediately unblocks and raises RpcError on the client thread. If the background thread has not yet updated self._target_future to the new retried future, the client will receive the RpcError and the transparent retry will fail.

To prevent this, we should introduce a self._retry_pending flag and a threading.Condition variable to block result() and other future methods while a retry is being initiated. We also use a daemon thread instead of the executor to avoid thread leaks.

class _RetryableStreamUnaryFuture(grpc.Call, grpc.Future):
    def __init__(self, continuation, client_call_details, request_iterator, interceptor):
        self._continuation = continuation
        self._client_call_details = client_call_details
        self._replayable_request_iterator = _ReplayableIterator(request_iterator)
        self._interceptor = interceptor
        self._retry_count = 0
        self._done_callbacks = []
        self._target_future = None
        self._lock = threading.Lock()
        self._condition = threading.Condition(self._lock)
        self._retry_pending = False
        self._start_call()

    def _on_inner_future_done(self, inner_future):
        with self._lock:
            if inner_future is not self._target_future:
                return
                
        exc = inner_future.exception()
        if isinstance(exc, grpc.RpcError):
            status_code = exc.code()
            can_replay = self._replayable_request_iterator.can_replay()
            
            if can_replay and self._interceptor._should_retry(status_code, self._retry_count, getattr(self, "_attempt_cert", None)):
                with self._lock:
                    self._retry_count += 1
                    self._retry_pending = True
                
                def async_retry():
                    try:
                        self._interceptor._wrapper.refresh_logic(self._retry_count)
                        time.sleep(random.uniform(0.1, 1.0))
                        self._start_call()
                    finally:
                        with self._lock:
                            self._retry_pending = False
                            self._condition.notify_all()
                    
                t = threading.Thread(target=async_retry, daemon=True)
                t.start()
                return
                
            if getattr(self._interceptor, "_wrapper", None):
                if self._interceptor._should_retry(status_code, 0, getattr(self, "_attempt_cert", None)):
                    self._interceptor._wrapper.refresh_logic(1)

        with self._lock:
            for cb in self._done_callbacks:
                cb(self)

    def _start_call(self):
        self._attempt_cert = self._interceptor._wrapper._cached_cert if self._interceptor._wrapper else None
        req_iter = iter(self._replayable_request_iterator)
        with self._lock:
            self._target_future = self._continuation(self._client_call_details, req_iter)
            self._target_future.add_done_callback(self._on_inner_future_done)

    def result(self, timeout=None):
        deadline = time.time() + timeout if timeout else None
        
        while True:
            with self._lock:
                while self._retry_pending:
                    if deadline:
                        remaining = deadline - time.time()
                        if remaining < 0:
                            raise grpc.FutureTimeoutError()
                        self._condition.wait(timeout=remaining)
                    else:
                        self._condition.wait()
                current_future = self._target_future
                
            try:
                if deadline:
                    remaining = deadline - time.time()
                    if remaining < 0:
                        raise grpc.FutureTimeoutError()
                    return current_future.result(timeout=remaining)
                else:    
                    return current_future.result()
                    
            except grpc.RpcError as e:
                with self._lock:
                    if current_future is not self._target_future or self._retry_pending:
                        continue
                raise e

    def add_done_callback(self, fn):
        with self._lock:
            self._done_callbacks.append(fn)
            if self._target_future.done() and not self._retry_pending:
                exc = self._target_future.exception()
                if not (isinstance(exc, grpc.RpcError) and self._interceptor._should_retry(exc.code(), self._retry_count, getattr(self, "_attempt_cert", None))):
                    fn(self)

    def exception(self, timeout=None):
        try:
            self.result(timeout)
            return None
        except Exception as e:
            return e
            
    def traceback(self, timeout=None):
        try:
            self.result(timeout)
            return None
        except Exception:
            with self._lock:
                return self._target_future.traceback(timeout=timeout)

    def cancel(self): 
        with self._lock: return self._target_future.cancel()
    def cancelled(self): 
        with self._lock: return self._target_future.cancelled()
    def running(self): 
        with self._lock: return self._target_future.running()
    def done(self): 
        with self._lock: return self._target_future.done() and not self._retry_pending
    def code(self): 
        with self._lock: return self._target_future.code()
    def details(self): 
        with self._lock: return self._target_future.details()
    def is_active(self):
        with self._lock: return self._target_future.is_active()
    def time_remaining(self):
        with self._lock: return self._target_future.time_remaining()
    def initial_metadata(self):
        with self._lock: return self._target_future.initial_metadata()
    def trailing_metadata(self):
        with self._lock: return self._target_future.trailing_metadata()
    def add_callback(self, cb):
        with self._lock: return self._target_future.add_callback(cb)

[gemini-code-assist]

Creating a ThreadPoolExecutor per interceptor instance without ever shutting it down leads to thread leaks, especially in environments where channels are frequently created and destroyed (such as unit tests). Since the background thread is only used for a one-off asynchronous retry of failed stream-unary calls, we can use a simple daemon threading.Thread instead.

Suggested change

	def __init__(self):
	self._wrapper = None
	self._max_retries = 2 # Set your desired limit here
	self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=5)
	def __init__(self):
	self._wrapper = None
	self._max_retries = 2 # Set your desired limit here