Repository: pantsbuild/pants

Issue:
Adjust plugin thread-local APIs to account for `WorkunitStore` state
When `eager_fetch=False`, it's possible that a workunit's "artifacts" contain `Digest`s which haven't actually been fetched. When that's the case for a `Digest`, and a `StreamingWorkunit` plugin is using any of the [context methods which fetch files](https://github.com/pantsbuild/pants/blob/1d8205538a2932badcc1738fb1288600908b01e1/src/python/pants/engine/streaming_workunit_handler.py#L55-L69) from a background thread, they will encounter a:
> A WorkunitStore has not been set for this thread.

...error. That's because our existing `native_engine.stdio_thread_set_destination` statics only set the thread local `stdio` state, and not also our workunit state.

----

To fix this, we should adjust the existing method to additionally set the workunit store. But we should also deprecate the existing method and add a new one with a more accurate name (replacing #12295).


Discussion:

Resolving patch:
diff --git a/src/python/pants/bsp/testutil.py b/src/python/pants/bsp/testutil.py
--- a/src/python/pants/bsp/testutil.py
+++ b/src/python/pants/bsp/testutil.py
@@ -16,7 +16,7 @@
 from pants.bsp.context import BSPContext
 from pants.bsp.protocol import BSPConnection
 from pants.bsp.rules import rules as bsp_rules
-from pants.engine.internals import native_engine
+from pants.engine.internals.native_engine import PyThreadLocals
 from pants.testutil.rule_runner import RuleRunner
 
 
@@ -93,7 +93,7 @@ def setup_bsp_server(
 ):
     rule_runner = rule_runner or RuleRunner(rules=bsp_rules())
     notification_names = notification_names or set()
-    stdio_destination = native_engine.stdio_thread_get_destination()
+    thread_locals = PyThreadLocals.get_for_current_thread()
 
     with setup_pipes() as pipes:
         context = BSPContext()
@@ -107,7 +107,7 @@ def setup_bsp_server(
         )
 
         def run_bsp_server():
-            native_engine.stdio_thread_set_destination(stdio_destination)
+            thread_locals.set_for_current_thread()
             conn.run()
 
         bsp_thread = Thread(target=run_bsp_server)
diff --git a/src/python/pants/engine/streaming_workunit_handler.py b/src/python/pants/engine/streaming_workunit_handler.py
--- a/src/python/pants/engine/streaming_workunit_handler.py
+++ b/src/python/pants/engine/streaming_workunit_handler.py
@@ -12,7 +12,7 @@
 from pants.base.specs import Specs
 from pants.engine.addresses import Addresses
 from pants.engine.fs import Digest, DigestContents, FileDigest, Snapshot
-from pants.engine.internals import native_engine
+from pants.engine.internals.native_engine import PyThreadLocals
 from pants.engine.internals.scheduler import SchedulerSession, Workunit
 from pants.engine.internals.selectors import Params
 from pants.engine.rules import Get, MultiGet, QueryRule, collect_rules, rule
@@ -30,6 +30,24 @@
 # -----------------------------------------------------------------------------------------------
 
 
+def thread_locals_get_for_current_thread() -> PyThreadLocals:
+    """Gets the engine's thread local state for the current thread.
+
+    In order to safely use StreamingWorkunitContext methods from additional threads,
+    StreamingWorkunit plugins should propagate thread local state from the threads that they are
+    initialized on to any additional threads that they spawn.
+    """
+    return PyThreadLocals.get_for_current_thread()
+
+
+def thread_locals_set_for_current_thread(thread_locals: PyThreadLocals) -> None:
+    """Sets the engine's thread local state for the current thread.
+
+    See `thread_locals_get`.
+    """
+    thread_locals.set_for_current_thread()
+
+
 @dataclass(frozen=True)
 class TargetInfo:
     filename: str
@@ -246,9 +264,9 @@ def __init__(
         self.block_until_complete = not allow_async_completion or any(
             callback.can_finish_async is False for callback in self.callbacks
         )
-        # Get the parent thread's logging destination. Note that this thread has not yet started
+        # Get the parent thread's thread locals. Note that this thread has not yet started
         # as we are only in the constructor.
-        self.logging_destination = native_engine.stdio_thread_get_destination()
+        self.thread_locals = PyThreadLocals.get_for_current_thread()
 
     def poll_workunits(self, *, finished: bool) -> None:
         workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)
@@ -261,8 +279,9 @@ def poll_workunits(self, *, finished: bool) -> None:
             )
 
     def run(self) -> None:
-        # First, set the thread's logging destination to the parent thread's, meaning the console.
-        native_engine.stdio_thread_set_destination(self.logging_destination)
+        # First, set the thread's thread locals to the parent thread's in order to propagate the
+        # console, workunit stores, etc.
+        self.thread_locals.set_for_current_thread()
         while not self.stop_request.isSet():
             self.poll_workunits(finished=False)
             self.stop_request.wait(timeout=self.report_interval)


Test patch: