RecordingStream context managers are broken when used with generators

[jleibs]

Problem

Example of problematic code:

import rerun as rr

def my_gen_func(stream, name):
    with stream:
        for i in range(10):
            print(f"{name} {i}")
            rr.log("stream", rr.TextLog(f"{name} {i}"))
            yield

rr.init("rerun_example_leak_context")

stream1 = rr.new_recording("rerun_example_stream1")
stream1.save("stream1.rrd")
stream2 = rr.new_recording("rerun_example_stream2")
stream2.save("stream2.rrd")

gen1 = my_gen_func(stream1, "stream1")
gen2 = my_gen_func(stream2, "stream2")

next(gen1)
next(gen2)
rr.log("stream", rr.TextLog("This should go to the global stream"))
next(gen1)
next(gen1)
next(gen1)

Run this and then open stream2.rrd.
You will see that both global messages and messages that should have been destined for stream 1 leaked into stream 2:

Investigation

The reason this happens is our context-manager tries to establish the context by running:

bindings.set_thread_local_data_recording(...)

But the context manager stays active on the thread after yielding the result, since the context is held around by the generator state. This means state leaks out of the generator into the thread and additionally when resuming the generator, can leak in from whatever context is calling the generator.

[jleibs]

Adding to 0.16 since this is very easy to hit when working while running code inside of an environment like Gradio when using streaming generators.

[jleibs]

Turns out this may not be generically solvable. See https://peps.python.org/pep-0568/ and https://discuss.python.org/t/preventing-yield-inside-certain-context-managers

The two best things we may be able to offer are probably:

• a utility-decorator that can capture and restore stream context on yield.
• maybe a sanity check in our context manager that catches some edge-cases, though this seems hard to guarantee

However, the footgun is going to be there regardless... this seems like something the python devs aren't really interested in fixing so not much we can do about it.