Span Issues when merging async genrators

[LukeLalor]

I'm having an issue where I am merging async generators.

Arbitrarily within the generators, more spans are being created. This can get messy and I would like to create a separate span for each generator to group the sub spans.

However, if I merge the generators using aiostream.stream.merge, the span's context gets lost, so the grouping spans disappear deactivate.

I have reduced the issue to a minimal test case that I have posted in stack-overflow.

Does anybody have ideas on what I am doing wrong, if this is a bug, or if there is a workaround I should consider?

[herotomg]

@LukeLalor Did you find a solution for this?

[LukeLalor]

I did not. Since the SO post was removed, I am going to copy it over here for easy repro

My opentelemetry spans are losing their parents when I create them within an async generator that I then merge with another async generator.

I believe the problem is that the span created within stream_fn is getting lost as the merged stream swaps between contexts on the event loop while awaiting.

Now, I know that I could manually restore the context after awaiting in this scenario, butin my real scenario restoring the context manually at every await would be very clunky and error prone.

how can I pass the context through properly when merging the streams?

Here is a minimal repro in the form of a self contained pytest

from collections import defaultdict
from aiostream import stream
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import get_current_span


async def test_recurse_streaming_context():
    trace.set_tracer_provider(TracerProvider())

    async def stream_fn(title):
        with tracer.start_as_current_span(title):
            context = get_current_span().get_span_context()
            yield (title, format(context.span_id, "016x"))
            await asyncio.sleep(0)
            context2 = get_current_span().get_span_context()
            yield (title, format(context2.span_id, "016x"))

    acc = defaultdict(list)
    tracer = trace.get_tracer("test_tracer")
    with tracer.start_as_current_span("test span"):
        async for e in stream.merge(*[stream_fn(f"child_{i}") for i in range(2)]):
            acc[e[0]].append(e[1])

    assert acc["child_0"][0] == acc["child_0"][1]
    # AssertionError: assert 'caa5694e43a514db' == '57adf8fb9e10a596'