dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4146869333
##########
task-sdk/tests/task_sdk/execution_time/test_context.py:
##########
@@ -554,6 +559,65 @@ def test_nested_context(self):
# End of with statement
ctx_list[i].__exit__(None, None, None)
+ @pytest.mark.parametrize("start", ["thread_pool", "thread"])
+ def test_thread_started_by_the_task_sees_its_context(self, start):
+ """A thread the task starts has empty ContextVars, and must still find
the task's context."""
+ task_context = {"Hello": "World"}
+
+ def read():
+ return get_current_context()
+
+ with set_current_context(task_context):
+ if start == "thread_pool":
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ seen = pool.submit(read).result()
+ else:
+ results = []
+ thread = threading.Thread(target=lambda:
results.append(read()))
+ thread.start()
+ thread.join()
+ (seen,) = results
+
+ assert seen is task_context
+
+ def test_indexed_context_covers_the_task_context_within_its_block(self):
+ task_context = {"ContextId": "task"}
+ indexed_context = {"ContextId": "iteration"}
+
+ with set_current_context(task_context):
+ with set_indexed_context(indexed_context):
+ assert get_current_context() is indexed_context
+ assert get_current_context() is task_context
+
+ def test_indexed_context_is_not_seen_by_other_threads(self):
+ """A thread started inside an iteration sees the task's context, not
the iteration's."""
+ task_context = {"ContextId": "task"}
+
+ with set_current_context(task_context):
+ with set_indexed_context({"ContextId": "iteration"}):
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ seen = pool.submit(get_current_context).result()
+
+ assert seen is task_context
+
+ @pytest.mark.asyncio
+ async def test_concurrent_iterations_each_see_their_own_context(self):
+ """Iterations interleaving on one event loop never see each other's
context."""
+ entered: list[int] = []
+ both_entered = asyncio.Event()
+
+ async def iteration(index):
+ with set_indexed_context({"ContextId": index}):
+ entered.append(index)
+ if len(entered) == 2:
+ both_entered.set()
+ await both_entered.wait()
Review Comment:
Right, it passed with a thread-local stack swapped in; I checked. 7925b470be
adds the `sleep(0)` and also a second barrier so neither iteration leaves
before both have read. `sleep(0)` alone still let the reads happen in nesting
order. With a thread-local stack it now fails with iteration 0 seeing 1.
Drafted-by: Claude Opus 5.5; reviewed by @dabla before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
