Move dg_stack into a ContextVar to support with blocks in separate threads

[eric-skydio]

Describe your changes

Previously, using a DeltaGenerator as a context manager appended to a dg_stack variable associated with the current ScriptRunContext. Unfortunately, if that ScriptRunContext is shared between multiple threads or async tasks, they would end up modifying each other's stacks, leading to broken behavior.

This PR instead uses the contextvars module to make dg_stack a context-local variable, which means it automatically does the "right thing" in those cases.

Note that this also requires making it immutable, so this PR switches from a list to a tuple. This does mean entering/exiting
contexts is now linear compute cost in the depth of the dg_stack, but it's really weird to have super deeply nested context managers, so I'm not too concerned.

Tagging @AnOctopus as the most recent person to touch this code.

GitHub Issue Link (if applicable)

#7668

Testing Plan

• Explanation of why no additional tests are needed
• Unit Tests (JS and/or Python)
• E2E Tests
• Any manual testing needed?

---

It would be reasonable to have a unit test that verifies the correct behavior here, but I'm not sure where to put it.

There's some reproducing code in the linked issue, and I also have an internal test case that could be adapted. I'm definitely open to suggestions/guidance here.

Contribution License Agreement

By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.

[eric-skydio]

It looks like I accidentally introduced a circular import cycle, the solution to which is probably moving dg_stack back into the script_run_context.py file. I can put up a revision that does that or you can do so.

[AnOctopus]

@eric-skydio This is interesting, I hadn't heard of contextvars before. I'll need to familiarize myself with it more to fully address the proposed change.

The current code mostly doesn't try to support scripts doing Streamlit Things in other threads or tasks, because many parts of the script runner still have implicit assumptions from when it was even more single-threaded. That said, I am open to changes that improve things in those situations, when the code complexity cost is low and we have confidence that it won't break other stuff.

If you have a small script that demonstrates the problem this is fixing, that would be helpful for both understanding the fix and as a test case.

[eric-skydio]

#7668 has a small demo script in it

[eric-skydio]

This came up because we're trying to improve the performance of a fairly large Streamlit app, and it quickly became clear from profiling that overall performance was limited by the single-threaded implementation of our app. Since Python has nice support for asyncio now and encourages using those patterns in other contexts, we're in the process of gradually rewriting things to use those patterns. That revealed the #7668 bug, and this PR is intended to fix that. Because our streamlit app is mostly spending its time loading and waiting for other data, making it concurrent shows about a 10x performance improvement, which gets our initial page load from 30s -> 3s and is a huge improvement in UX.

[AnOctopus]

Now that I have a better understanding of the problem and fix, I'm on board with this change.

Only thing the PR is missing is tests. I think it would suffice to turn the repro script from #7668 into a python test checking the element order is correct

[eric-skydio]

Done. I added tests for both asyncio and threading, and also used Events to avoid needing to put sleeps in (which are both nondeterministic and slow)

[AnOctopus]

Excellent use of diagram explanation

[AnOctopus]

It might be cleaner to implement the tests using AppTest and assertions on the resulting elements, but this is still plenty easy to follow, so I'd be fine with merging it as is. Just let me know which you'd prefer.

[AnOctopus]

And as a small note to myself later, there are probably other places where we use thread locals with context managers that would benefit from this kind of change.

[eric-skydio]

I decided to go with this test pattern to keep it next to other tests of similar functionality, and because I don't actually have much familiarity with AppTest yet since we just upgraded to a streamlit version that has it.

Notably, this was a global variable not even thread-local. The only explicit thread-local variables I've seen are in the caching decorators, which currently don't support async for other reasons. I'm mildly interested in fixing that, but it'll require more design work, and would possibly even benefit from a live design brainstorming session or something.