grpc integration causes exception with grpc.aio streaming

[majorgreys]

Summary of problem

Enabling the grpc integration results in a StopAsyncIteration exception when handling streaming responses over a grpc.aio channel.

Which version of dd-trace-py are you using?

2.8.3

Which version of pip are you using?

24.0

Which libraries and their versions are you using?

`pip freeze`

attrs==23.2.0
bytecode==0.15.1
cattrs==23.2.3
ddsketch==3.0.1
ddtrace==2.8.3
Deprecated==1.2.14
envier==0.5.1
grpcio==1.63.0
importlib-metadata==7.0.0
opentelemetry-api==1.24.0
protobuf==3.20.3
six==1.16.0
sqlparse==0.5.0
typing_extensions==4.11.0
wrapt==1.16.0
xmltodict==0.13.0
zipp==3.18.1

How can we reproduce your problem?

1. git clone https://gist.github.com/majorgreys/f342630b8815ccb08b8445e1864c0012 && cd f342630b8815ccb08b8445e1864c0012
2. pip install grpcio==1.63.0 "protobuf>=3.5.0.post1,<4.0dev" ddtrace==2.8.3
3. Start server in terminal: python async_greeter_server.py
4. ddtrace-run python async_greeter_client.py

What is the result that you get?

Traceback (most recent call last):
  File "/Users/tahir.butt/dev/local/f342630b8815ccb08b8445e1864c0012/async_greeter_client.py", line 52, in <module>
    asyncio.run(run())
  File "/Users/tahir.butt/.pyenv/versions/3.11.3/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/Users/tahir.butt/.pyenv/versions/3.11.3/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tahir.butt/.pyenv/versions/3.11.3/lib/python3.11/asyncio/base_events.py", line 653, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/Users/tahir.butt/dev/local/f342630b8815ccb08b8445e1864c0012/venv/lib/python3.11/site-packages/ddtrace/contrib/asyncio/patch.py", line 50, in traced_coro
    return await coro
           ^^^^^^^^^^
  File "/Users/tahir.butt/dev/local/f342630b8815ccb08b8445e1864c0012/async_greeter_client.py", line 42, in run
    response = await hello_stream.read()
               ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tahir.butt/dev/local/f342630b8815ccb08b8445e1864c0012/venv/lib/python3.11/site-packages/grpc/aio/_interceptor.py", line 502, in read
    return await self._response_aiter.asend(None)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
StopAsyncIteration

What is the result that you expected?

Greeter client received from async generator: Hello number 0, you!
Greeter client received from async generator: Hello number 1, you!
Greeter client received from async generator: Hello number 2, you!
Greeter client received from async generator: Hello number 3, you!
Greeter client received from async generator: Hello number 4, you!
Greeter client received from async generator: Hello number 5, you!
Greeter client received from async generator: Hello number 6, you!
Greeter client received from async generator: Hello number 7, you!
Greeter client received from async generator: Hello number 8, you!
Greeter client received from async generator: Hello number 9, you!
Greeter client received from direct read: Hello number 0, you!
Greeter client received from direct read: Hello number 1, you!
Greeter client received from direct read: Hello number 2, you!
Greeter client received from direct read: Hello number 3, you!
Greeter client received from direct read: Hello number 4, you!
Greeter client received from direct read: Hello number 5, you!
Greeter client received from direct read: Hello number 6, you!
Greeter client received from direct read: Hello number 7, you!
Greeter client received from direct read: Hello number 8, you!
Greeter client received from direct read: Hello number 9, you!

[wconti27]

This doesn't seem to be specific to dd-trace-py. I was able to replicate the issue with the following sample app which simply adds an async client interceptor to the repo mentioned above. In theory, this should work (unless I'm doing something wrong adding the client interceptor).

# Copyright 2021 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The Python AsyncIO implementation of the GRPC hellostreamingworld.MultiGreeter client."""

import asyncio
import logging

import grpc
import hellostreamingworld_pb2
import hellostreamingworld_pb2_grpc


async def wrap(call):
    try:
        async for response in call:
            print(f"Received response: {response.message}")
            yield response
    except Exception as exc:
        breakpoint()
        raise exc


class StreamInterceptor(grpc.aio.UnaryStreamClientInterceptor):
    async def intercept_unary_stream(self, continuation, call_details, request):
        print("Before request")
        response_iterator = await continuation(call_details, request)
        print("after request")
        return wrap(response_iterator)


async def run() -> None:
    async with grpc.aio.insecure_channel("localhost:50051", interceptors=[StreamInterceptor()]) as channel:
        stub = hellostreamingworld_pb2_grpc.MultiGreeterStub(channel)

        # # Read from an async generator
        # async for response in stub.sayHello(
        #     hellostreamingworld_pb2.HelloRequest(name="you")
        # ):
        #     print(
        #         "Greeter client received from async generator: "
        #         + response.message
        #     )

        # Direct read from the stub
        hello_stream = stub.sayHello(
            hellostreamingworld_pb2.HelloRequest(name="will")
        )
        while True:
            breakpoint()
            response = await hello_stream.read()
            if response == grpc.aio.EOF:
                break
            print(
                "Greeter client received from direct read: " + response.message
            )


if __name__ == "__main__":
    logging.basicConfig()
    asyncio.run(run())

    # At the end of your program
    trace.get_tracer_provider().shutdown()