We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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).

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/python/pants/engine/streaming_workunit_handler.py`
Content:
```
1 # Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
2 # Licensed under the Apache License, Version 2.0 (see LICENSE).
3 
4 from __future__ import annotations
5 
6 import logging
7 import threading
8 from abc import ABC, abstractmethod
9 from dataclasses import dataclass
10 from typing import Any, Callable, Iterable, Sequence, Tuple
11 
12 from pants.base.specs import Specs
13 from pants.engine.addresses import Addresses
14 from pants.engine.fs import Digest, DigestContents, FileDigest, Snapshot
15 from pants.engine.internals import native_engine
16 from pants.engine.internals.scheduler import SchedulerSession, Workunit
17 from pants.engine.internals.selectors import Params
18 from pants.engine.rules import Get, MultiGet, QueryRule, collect_rules, rule
19 from pants.engine.target import Targets
20 from pants.engine.unions import UnionMembership, union
21 from pants.goal.run_tracker import RunTracker
22 from pants.option.options_bootstrapper import OptionsBootstrapper
23 from pants.util.logging import LogLevel
24 
25 logger = logging.getLogger(__name__)
26 
27 
28 # -----------------------------------------------------------------------------------------------
29 # Streaming workunits plugin API
30 # -----------------------------------------------------------------------------------------------
31 
32 
33 @dataclass(frozen=True)
34 class TargetInfo:
35     filename: str
36 
37 
38 @dataclass(frozen=True)
39 class ExpandedSpecs:
40     targets: dict[str, list[TargetInfo]]
41 
42 
43 @dataclass(frozen=True)
44 class StreamingWorkunitContext:
45     _scheduler: SchedulerSession
46     _run_tracker: RunTracker
47     _specs: Specs
48     _options_bootstrapper: OptionsBootstrapper
49 
50     @property
51     def run_tracker(self) -> RunTracker:
52         """Returns the RunTracker for the current run of Pants."""
53         return self._run_tracker
54 
55     def single_file_digests_to_bytes(self, digests: Sequence[FileDigest]) -> list[bytes]:
56         """Return `bytes` for each `FileDigest`."""
57         return self._scheduler.single_file_digests_to_bytes(digests)
58 
59     def snapshots_to_file_contents(
60         self, snapshots: Sequence[Snapshot]
61     ) -> tuple[DigestContents, ...]:
62         """Given a sequence of Snapshot objects, return a tuple of DigestContents representing the
63         files contained in those `Snapshot`s in sequence."""
64         return self._scheduler.snapshots_to_file_contents(snapshots)
65 
66     def ensure_remote_has_recursive(self, digests: Sequence[Digest | FileDigest]) -> None:
67         """Invoke the internal ensure_remote_has_recursive function, which ensures that a remote
68         ByteStore, if it exists, has a copy of the files fingerprinted by each Digest."""
69         return self._scheduler.ensure_remote_has_recursive(digests)
70 
71     def get_metrics(self) -> dict[str, int]:
72         """Invoke the internal get_metrics function, which returns metrics for the Session."""
73         return self._scheduler.get_metrics()
74 
75     def get_observation_histograms(self) -> dict[str, Any]:
76         """Invoke the internal get_observation_histograms function, which serializes histograms
77         generated from Pants-internal observation metrics observed during the current run of Pants.
78 
79         These metrics are useful for debugging Pants internals.
80         """
81         return self._scheduler.get_observation_histograms()
82 
83     def get_expanded_specs(self) -> ExpandedSpecs:
84         """Return a dict containing the canonicalized addresses of the specs for this run, and what
85         files they expand to."""
86 
87         (unexpanded_addresses,) = self._scheduler.product_request(
88             Addresses, [Params(self._specs, self._options_bootstrapper)]
89         )
90 
91         expanded_targets = self._scheduler.product_request(
92             Targets, [Params(Addresses([addr])) for addr in unexpanded_addresses]
93         )
94         targets_dict: dict[str, list[TargetInfo]] = {}
95         for addr, targets in zip(unexpanded_addresses, expanded_targets):
96             targets_dict[addr.spec] = [
97                 TargetInfo(
98                     filename=(
99                         tgt.address.filename if tgt.address.is_file_target else str(tgt.address)
100                     )
101                 )
102                 for tgt in targets
103             ]
104         return ExpandedSpecs(targets=targets_dict)
105 
106 
107 class WorkunitsCallback(ABC):
108     @abstractmethod
109     def __call__(
110         self,
111         *,
112         started_workunits: tuple[Workunit, ...],
113         completed_workunits: tuple[Workunit, ...],
114         finished: bool,
115         context: StreamingWorkunitContext,
116     ) -> None:
117         """
118         :started_workunits: Workunits that have started but not completed.
119         :completed_workunits: Workunits that have completed.
120         :finished: True when the last chunk of workunit data is reported to the callback.
121         :context: A context providing access to functionality relevant to the run.
122         """
123 
124     @property
125     @abstractmethod
126     def can_finish_async(self) -> bool:
127         """Can this callback finish its work in the background after the Pants run has already
128         completed?
129 
130         The main reason to `return False` is if your callback logs in its final call, when
131         `finished=True`, as it may end up logging to `.pantsd.d/pants.log` instead of the console,
132         which is harder for users to find. Otherwise, most callbacks should return `True` to avoid
133         slowing down Pants from finishing the run.
134         """
135 
136 
137 @dataclass(frozen=True)
138 class WorkunitsCallbackFactory:
139     """A wrapper around a callable that constructs WorkunitsCallbacks.
140 
141     NB: This extra wrapping is because subtyping is not supported in the return position of a
142     rule. See #11354 for discussion of that limitation.
143     """
144 
145     callback_factory: Callable[[], WorkunitsCallback | None]
146 
147 
148 class WorkunitsCallbackFactories(Tuple[WorkunitsCallbackFactory, ...]):
149     """A list of registered factories for WorkunitsCallback instances."""
150 
151 
152 @union
153 class WorkunitsCallbackFactoryRequest:
154     """A request for a particular WorkunitsCallbackFactory."""
155 
156 
157 @rule
158 async def construct_workunits_callback_factories(
159     union_membership: UnionMembership,
160 ) -> WorkunitsCallbackFactories:
161     request_types = union_membership.get(WorkunitsCallbackFactoryRequest)
162     workunit_callback_factories = await MultiGet(
163         Get(WorkunitsCallbackFactory, WorkunitsCallbackFactoryRequest, request_type())
164         for request_type in request_types
165     )
166     return WorkunitsCallbackFactories(workunit_callback_factories)
167 
168 
169 # -----------------------------------------------------------------------------------------------
170 # Streaming workunits handler
171 # -----------------------------------------------------------------------------------------------
172 
173 
174 class StreamingWorkunitHandler:
175     """Periodically calls each registered WorkunitsCallback in a dedicated thread.
176 
177     This class should be used as a context manager.
178     """
179 
180     def __init__(
181         self,
182         scheduler: SchedulerSession,
183         run_tracker: RunTracker,
184         callbacks: Iterable[WorkunitsCallback],
185         options_bootstrapper: OptionsBootstrapper,
186         specs: Specs,
187         report_interval_seconds: float,
188         allow_async_completion: bool,
189         max_workunit_verbosity: LogLevel,
190     ) -> None:
191         scheduler = scheduler.isolated_shallow_clone("streaming_workunit_handler_session")
192         self.callbacks = callbacks
193         self.context = StreamingWorkunitContext(
194             _scheduler=scheduler,
195             _run_tracker=run_tracker,
196             _specs=specs,
197             _options_bootstrapper=options_bootstrapper,
198         )
199         self.thread_runner = (
200             _InnerHandler(
201                 scheduler=scheduler,
202                 context=self.context,
203                 callbacks=self.callbacks,
204                 report_interval=report_interval_seconds,
205                 # TODO(10092) The max verbosity should be a per-client setting, rather than a global
206                 #  setting.
207                 max_workunit_verbosity=max_workunit_verbosity,
208                 allow_async_completion=allow_async_completion,
209             )
210             if callbacks
211             else None
212         )
213 
214     def __enter__(self) -> None:
215         if not self.thread_runner:
216             return
217         self.thread_runner.start()
218 
219     def __exit__(self, exc_type, exc_value, traceback) -> None:
220         if not self.thread_runner:
221             return
222         self.thread_runner.end()
223         if exc_type is not None:
224             self.thread_runner.join()
225 
226 
227 class _InnerHandler(threading.Thread):
228     def __init__(
229         self,
230         scheduler: Any,
231         context: StreamingWorkunitContext,
232         callbacks: Iterable[WorkunitsCallback],
233         report_interval: float,
234         max_workunit_verbosity: LogLevel,
235         allow_async_completion: bool,
236     ) -> None:
237         super().__init__(daemon=True)
238         self.scheduler = scheduler
239         self.context = context
240         self.stop_request = threading.Event()
241         self.report_interval = report_interval
242         self.callbacks = callbacks
243         self.max_workunit_verbosity = max_workunit_verbosity
244         # TODO: Have a thread per callback so that some callbacks can always finish async even
245         #  if others must be finished synchronously.
246         self.block_until_complete = not allow_async_completion or any(
247             callback.can_finish_async is False for callback in self.callbacks
248         )
249         # Get the parent thread's logging destination. Note that this thread has not yet started
250         # as we are only in the constructor.
251         self.logging_destination = native_engine.stdio_thread_get_destination()
252 
253     def poll_workunits(self, *, finished: bool) -> None:
254         workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)
255         for callback in self.callbacks:
256             callback(
257                 started_workunits=workunits["started"],
258                 completed_workunits=workunits["completed"],
259                 finished=finished,
260                 context=self.context,
261             )
262 
263     def run(self) -> None:
264         # First, set the thread's logging destination to the parent thread's, meaning the console.
265         native_engine.stdio_thread_set_destination(self.logging_destination)
266         while not self.stop_request.isSet():  # type: ignore[attr-defined]
267             self.poll_workunits(finished=False)
268             self.stop_request.wait(timeout=self.report_interval)
269         else:
270             # Make one final call. Note that this may run after the Pants run has already
271             # completed, depending on whether the thread was joined or not.
272             self.poll_workunits(finished=True)
273 
274     def end(self) -> None:
275         self.stop_request.set()
276         if self.block_until_complete:
277             logger.debug(
278                 "Async completion is disabled: waiting for workunit callbacks to complete..."
279             )
280             super().join()
281         else:
282             logger.debug(
283                 "Async completion is enabled: workunit callbacks will complete in the background."
284             )
285 
286 
287 def rules():
288     return [
289         QueryRule(WorkunitsCallbackFactories, (UnionMembership,)),
290         QueryRule(Targets, (Addresses,)),
291         QueryRule(Addresses, (Specs, OptionsBootstrapper)),
292         *collect_rules(),
293     ]
294
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 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():  # type: ignore[attr-defined]
             self.poll_workunits(finished=False)
             self.stop_request.wait(timeout=self.report_interval)

{"golden_diff": "diff --git a/src/python/pants/engine/streaming_workunit_handler.py b/src/python/pants/engine/streaming_workunit_handler.py\n--- a/src/python/pants/engine/streaming_workunit_handler.py\n+++ b/src/python/pants/engine/streaming_workunit_handler.py\n@@ -12,7 +12,7 @@\n from pants.base.specs import Specs\n from pants.engine.addresses import Addresses\n from pants.engine.fs import Digest, DigestContents, FileDigest, Snapshot\n-from pants.engine.internals import native_engine\n+from pants.engine.internals.native_engine import PyThreadLocals\n from pants.engine.internals.scheduler import SchedulerSession, Workunit\n from pants.engine.internals.selectors import Params\n from pants.engine.rules import Get, MultiGet, QueryRule, collect_rules, rule\n@@ -30,6 +30,24 @@\n # -----------------------------------------------------------------------------------------------\n \n \n+def thread_locals_get_for_current_thread() -> PyThreadLocals:\n+    \"\"\"Gets the engine's thread local state for the current thread.\n+\n+    In order to safely use StreamingWorkunitContext methods from additional threads,\n+    StreamingWorkunit plugins should propagate thread local state from the threads that they are\n+    initialized on to any additional threads that they spawn.\n+    \"\"\"\n+    return PyThreadLocals.get_for_current_thread()\n+\n+\n+def thread_locals_set_for_current_thread(thread_locals: PyThreadLocals) -> None:\n+    \"\"\"Sets the engine's thread local state for the current thread.\n+\n+    See `thread_locals_get`.\n+    \"\"\"\n+    thread_locals.set_for_current_thread()\n+\n+\n @dataclass(frozen=True)\n class TargetInfo:\n     filename: str\n@@ -246,9 +264,9 @@\n         self.block_until_complete = not allow_async_completion or any(\n             callback.can_finish_async is False for callback in self.callbacks\n         )\n-        # Get the parent thread's logging destination. Note that this thread has not yet started\n+        # Get the parent thread's thread locals. Note that this thread has not yet started\n         # as we are only in the constructor.\n-        self.logging_destination = native_engine.stdio_thread_get_destination()\n+        self.thread_locals = PyThreadLocals.get_for_current_thread()\n \n     def poll_workunits(self, *, finished: bool) -> None:\n         workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)\n@@ -261,8 +279,9 @@\n             )\n \n     def run(self) -> None:\n-        # First, set the thread's logging destination to the parent thread's, meaning the console.\n-        native_engine.stdio_thread_set_destination(self.logging_destination)\n+        # First, set the thread's thread locals to the parent thread's in order to propagate the\n+        # console, workunit stores, etc.\n+        self.thread_locals.set_for_current_thread()\n         while not self.stop_request.isSet():  # type: ignore[attr-defined]\n             self.poll_workunits(finished=False)\n             self.stop_request.wait(timeout=self.report_interval)\n", "issue": "Adjust plugin thread-local APIs to account for `WorkunitStore` state\nWhen `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:\r\n> A WorkunitStore has not been set for this thread.\r\n\r\n...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.\r\n\r\n----\r\n\r\nTo 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).\n", "before_files": [{"content": "# Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).\n# Licensed under the Apache License, Version 2.0 (see LICENSE).\n\nfrom __future__ import annotations\n\nimport logging\nimport threading\nfrom abc import ABC, abstractmethod\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Iterable, Sequence, Tuple\n\nfrom pants.base.specs import Specs\nfrom pants.engine.addresses import Addresses\nfrom pants.engine.fs import Digest, DigestContents, FileDigest, Snapshot\nfrom pants.engine.internals import native_engine\nfrom pants.engine.internals.scheduler import SchedulerSession, Workunit\nfrom pants.engine.internals.selectors import Params\nfrom pants.engine.rules import Get, MultiGet, QueryRule, collect_rules, rule\nfrom pants.engine.target import Targets\nfrom pants.engine.unions import UnionMembership, union\nfrom pants.goal.run_tracker import RunTracker\nfrom pants.option.options_bootstrapper import OptionsBootstrapper\nfrom pants.util.logging import LogLevel\n\nlogger = logging.getLogger(__name__)\n\n\n# -----------------------------------------------------------------------------------------------\n# Streaming workunits plugin API\n# -----------------------------------------------------------------------------------------------\n\n\n@dataclass(frozen=True)\nclass TargetInfo:\n    filename: str\n\n\n@dataclass(frozen=True)\nclass ExpandedSpecs:\n    targets: dict[str, list[TargetInfo]]\n\n\n@dataclass(frozen=True)\nclass StreamingWorkunitContext:\n    _scheduler: SchedulerSession\n    _run_tracker: RunTracker\n    _specs: Specs\n    _options_bootstrapper: OptionsBootstrapper\n\n    @property\n    def run_tracker(self) -> RunTracker:\n        \"\"\"Returns the RunTracker for the current run of Pants.\"\"\"\n        return self._run_tracker\n\n    def single_file_digests_to_bytes(self, digests: Sequence[FileDigest]) -> list[bytes]:\n        \"\"\"Return `bytes` for each `FileDigest`.\"\"\"\n        return self._scheduler.single_file_digests_to_bytes(digests)\n\n    def snapshots_to_file_contents(\n        self, snapshots: Sequence[Snapshot]\n    ) -> tuple[DigestContents, ...]:\n        \"\"\"Given a sequence of Snapshot objects, return a tuple of DigestContents representing the\n        files contained in those `Snapshot`s in sequence.\"\"\"\n        return self._scheduler.snapshots_to_file_contents(snapshots)\n\n    def ensure_remote_has_recursive(self, digests: Sequence[Digest | FileDigest]) -> None:\n        \"\"\"Invoke the internal ensure_remote_has_recursive function, which ensures that a remote\n        ByteStore, if it exists, has a copy of the files fingerprinted by each Digest.\"\"\"\n        return self._scheduler.ensure_remote_has_recursive(digests)\n\n    def get_metrics(self) -> dict[str, int]:\n        \"\"\"Invoke the internal get_metrics function, which returns metrics for the Session.\"\"\"\n        return self._scheduler.get_metrics()\n\n    def get_observation_histograms(self) -> dict[str, Any]:\n        \"\"\"Invoke the internal get_observation_histograms function, which serializes histograms\n        generated from Pants-internal observation metrics observed during the current run of Pants.\n\n        These metrics are useful for debugging Pants internals.\n        \"\"\"\n        return self._scheduler.get_observation_histograms()\n\n    def get_expanded_specs(self) -> ExpandedSpecs:\n        \"\"\"Return a dict containing the canonicalized addresses of the specs for this run, and what\n        files they expand to.\"\"\"\n\n        (unexpanded_addresses,) = self._scheduler.product_request(\n            Addresses, [Params(self._specs, self._options_bootstrapper)]\n        )\n\n        expanded_targets = self._scheduler.product_request(\n            Targets, [Params(Addresses([addr])) for addr in unexpanded_addresses]\n        )\n        targets_dict: dict[str, list[TargetInfo]] = {}\n        for addr, targets in zip(unexpanded_addresses, expanded_targets):\n            targets_dict[addr.spec] = [\n                TargetInfo(\n                    filename=(\n                        tgt.address.filename if tgt.address.is_file_target else str(tgt.address)\n                    )\n                )\n                for tgt in targets\n            ]\n        return ExpandedSpecs(targets=targets_dict)\n\n\nclass WorkunitsCallback(ABC):\n    @abstractmethod\n    def __call__(\n        self,\n        *,\n        started_workunits: tuple[Workunit, ...],\n        completed_workunits: tuple[Workunit, ...],\n        finished: bool,\n        context: StreamingWorkunitContext,\n    ) -> None:\n        \"\"\"\n        :started_workunits: Workunits that have started but not completed.\n        :completed_workunits: Workunits that have completed.\n        :finished: True when the last chunk of workunit data is reported to the callback.\n        :context: A context providing access to functionality relevant to the run.\n        \"\"\"\n\n    @property\n    @abstractmethod\n    def can_finish_async(self) -> bool:\n        \"\"\"Can this callback finish its work in the background after the Pants run has already\n        completed?\n\n        The main reason to `return False` is if your callback logs in its final call, when\n        `finished=True`, as it may end up logging to `.pantsd.d/pants.log` instead of the console,\n        which is harder for users to find. Otherwise, most callbacks should return `True` to avoid\n        slowing down Pants from finishing the run.\n        \"\"\"\n\n\n@dataclass(frozen=True)\nclass WorkunitsCallbackFactory:\n    \"\"\"A wrapper around a callable that constructs WorkunitsCallbacks.\n\n    NB: This extra wrapping is because subtyping is not supported in the return position of a\n    rule. See #11354 for discussion of that limitation.\n    \"\"\"\n\n    callback_factory: Callable[[], WorkunitsCallback | None]\n\n\nclass WorkunitsCallbackFactories(Tuple[WorkunitsCallbackFactory, ...]):\n    \"\"\"A list of registered factories for WorkunitsCallback instances.\"\"\"\n\n\n@union\nclass WorkunitsCallbackFactoryRequest:\n    \"\"\"A request for a particular WorkunitsCallbackFactory.\"\"\"\n\n\n@rule\nasync def construct_workunits_callback_factories(\n    union_membership: UnionMembership,\n) -> WorkunitsCallbackFactories:\n    request_types = union_membership.get(WorkunitsCallbackFactoryRequest)\n    workunit_callback_factories = await MultiGet(\n        Get(WorkunitsCallbackFactory, WorkunitsCallbackFactoryRequest, request_type())\n        for request_type in request_types\n    )\n    return WorkunitsCallbackFactories(workunit_callback_factories)\n\n\n# -----------------------------------------------------------------------------------------------\n# Streaming workunits handler\n# -----------------------------------------------------------------------------------------------\n\n\nclass StreamingWorkunitHandler:\n    \"\"\"Periodically calls each registered WorkunitsCallback in a dedicated thread.\n\n    This class should be used as a context manager.\n    \"\"\"\n\n    def __init__(\n        self,\n        scheduler: SchedulerSession,\n        run_tracker: RunTracker,\n        callbacks: Iterable[WorkunitsCallback],\n        options_bootstrapper: OptionsBootstrapper,\n        specs: Specs,\n        report_interval_seconds: float,\n        allow_async_completion: bool,\n        max_workunit_verbosity: LogLevel,\n    ) -> None:\n        scheduler = scheduler.isolated_shallow_clone(\"streaming_workunit_handler_session\")\n        self.callbacks = callbacks\n        self.context = StreamingWorkunitContext(\n            _scheduler=scheduler,\n            _run_tracker=run_tracker,\n            _specs=specs,\n            _options_bootstrapper=options_bootstrapper,\n        )\n        self.thread_runner = (\n            _InnerHandler(\n                scheduler=scheduler,\n                context=self.context,\n                callbacks=self.callbacks,\n                report_interval=report_interval_seconds,\n                # TODO(10092) The max verbosity should be a per-client setting, rather than a global\n                #  setting.\n                max_workunit_verbosity=max_workunit_verbosity,\n                allow_async_completion=allow_async_completion,\n            )\n            if callbacks\n            else None\n        )\n\n    def __enter__(self) -> None:\n        if not self.thread_runner:\n            return\n        self.thread_runner.start()\n\n    def __exit__(self, exc_type, exc_value, traceback) -> None:\n        if not self.thread_runner:\n            return\n        self.thread_runner.end()\n        if exc_type is not None:\n            self.thread_runner.join()\n\n\nclass _InnerHandler(threading.Thread):\n    def __init__(\n        self,\n        scheduler: Any,\n        context: StreamingWorkunitContext,\n        callbacks: Iterable[WorkunitsCallback],\n        report_interval: float,\n        max_workunit_verbosity: LogLevel,\n        allow_async_completion: bool,\n    ) -> None:\n        super().__init__(daemon=True)\n        self.scheduler = scheduler\n        self.context = context\n        self.stop_request = threading.Event()\n        self.report_interval = report_interval\n        self.callbacks = callbacks\n        self.max_workunit_verbosity = max_workunit_verbosity\n        # TODO: Have a thread per callback so that some callbacks can always finish async even\n        #  if others must be finished synchronously.\n        self.block_until_complete = not allow_async_completion or any(\n            callback.can_finish_async is False for callback in self.callbacks\n        )\n        # Get the parent thread's logging destination. Note that this thread has not yet started\n        # as we are only in the constructor.\n        self.logging_destination = native_engine.stdio_thread_get_destination()\n\n    def poll_workunits(self, *, finished: bool) -> None:\n        workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)\n        for callback in self.callbacks:\n            callback(\n                started_workunits=workunits[\"started\"],\n                completed_workunits=workunits[\"completed\"],\n                finished=finished,\n                context=self.context,\n            )\n\n    def run(self) -> None:\n        # First, set the thread's logging destination to the parent thread's, meaning the console.\n        native_engine.stdio_thread_set_destination(self.logging_destination)\n        while not self.stop_request.isSet():  # type: ignore[attr-defined]\n            self.poll_workunits(finished=False)\n            self.stop_request.wait(timeout=self.report_interval)\n        else:\n            # Make one final call. Note that this may run after the Pants run has already\n            # completed, depending on whether the thread was joined or not.\n            self.poll_workunits(finished=True)\n\n    def end(self) -> None:\n        self.stop_request.set()\n        if self.block_until_complete:\n            logger.debug(\n                \"Async completion is disabled: waiting for workunit callbacks to complete...\"\n            )\n            super().join()\n        else:\n            logger.debug(\n                \"Async completion is enabled: workunit callbacks will complete in the background.\"\n            )\n\n\ndef rules():\n    return [\n        QueryRule(WorkunitsCallbackFactories, (UnionMembership,)),\n        QueryRule(Targets, (Addresses,)),\n        QueryRule(Addresses, (Specs, OptionsBootstrapper)),\n        *collect_rules(),\n    ]\n", "path": "src/python/pants/engine/streaming_workunit_handler.py"}], "after_files": [{"content": "# Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).\n# Licensed under the Apache License, Version 2.0 (see LICENSE).\n\nfrom __future__ import annotations\n\nimport logging\nimport threading\nfrom abc import ABC, abstractmethod\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Iterable, Sequence, Tuple\n\nfrom pants.base.specs import Specs\nfrom pants.engine.addresses import Addresses\nfrom pants.engine.fs import Digest, DigestContents, FileDigest, Snapshot\nfrom pants.engine.internals.native_engine import PyThreadLocals\nfrom pants.engine.internals.scheduler import SchedulerSession, Workunit\nfrom pants.engine.internals.selectors import Params\nfrom pants.engine.rules import Get, MultiGet, QueryRule, collect_rules, rule\nfrom pants.engine.target import Targets\nfrom pants.engine.unions import UnionMembership, union\nfrom pants.goal.run_tracker import RunTracker\nfrom pants.option.options_bootstrapper import OptionsBootstrapper\nfrom pants.util.logging import LogLevel\n\nlogger = logging.getLogger(__name__)\n\n\n# -----------------------------------------------------------------------------------------------\n# Streaming workunits plugin API\n# -----------------------------------------------------------------------------------------------\n\n\ndef thread_locals_get_for_current_thread() -> PyThreadLocals:\n    \"\"\"Gets the engine's thread local state for the current thread.\n\n    In order to safely use StreamingWorkunitContext methods from additional threads,\n    StreamingWorkunit plugins should propagate thread local state from the threads that they are\n    initialized on to any additional threads that they spawn.\n    \"\"\"\n    return PyThreadLocals.get_for_current_thread()\n\n\ndef thread_locals_set_for_current_thread(thread_locals: PyThreadLocals) -> None:\n    \"\"\"Sets the engine's thread local state for the current thread.\n\n    See `thread_locals_get`.\n    \"\"\"\n    thread_locals.set_for_current_thread()\n\n\n@dataclass(frozen=True)\nclass TargetInfo:\n    filename: str\n\n\n@dataclass(frozen=True)\nclass ExpandedSpecs:\n    targets: dict[str, list[TargetInfo]]\n\n\n@dataclass(frozen=True)\nclass StreamingWorkunitContext:\n    _scheduler: SchedulerSession\n    _run_tracker: RunTracker\n    _specs: Specs\n    _options_bootstrapper: OptionsBootstrapper\n\n    @property\n    def run_tracker(self) -> RunTracker:\n        \"\"\"Returns the RunTracker for the current run of Pants.\"\"\"\n        return self._run_tracker\n\n    def single_file_digests_to_bytes(self, digests: Sequence[FileDigest]) -> list[bytes]:\n        \"\"\"Return `bytes` for each `FileDigest`.\"\"\"\n        return self._scheduler.single_file_digests_to_bytes(digests)\n\n    def snapshots_to_file_contents(\n        self, snapshots: Sequence[Snapshot]\n    ) -> tuple[DigestContents, ...]:\n        \"\"\"Given a sequence of Snapshot objects, return a tuple of DigestContents representing the\n        files contained in those `Snapshot`s in sequence.\"\"\"\n        return self._scheduler.snapshots_to_file_contents(snapshots)\n\n    def ensure_remote_has_recursive(self, digests: Sequence[Digest | FileDigest]) -> None:\n        \"\"\"Invoke the internal ensure_remote_has_recursive function, which ensures that a remote\n        ByteStore, if it exists, has a copy of the files fingerprinted by each Digest.\"\"\"\n        return self._scheduler.ensure_remote_has_recursive(digests)\n\n    def get_metrics(self) -> dict[str, int]:\n        \"\"\"Invoke the internal get_metrics function, which returns metrics for the Session.\"\"\"\n        return self._scheduler.get_metrics()\n\n    def get_observation_histograms(self) -> dict[str, Any]:\n        \"\"\"Invoke the internal get_observation_histograms function, which serializes histograms\n        generated from Pants-internal observation metrics observed during the current run of Pants.\n\n        These metrics are useful for debugging Pants internals.\n        \"\"\"\n        return self._scheduler.get_observation_histograms()\n\n    def get_expanded_specs(self) -> ExpandedSpecs:\n        \"\"\"Return a dict containing the canonicalized addresses of the specs for this run, and what\n        files they expand to.\"\"\"\n\n        (unexpanded_addresses,) = self._scheduler.product_request(\n            Addresses, [Params(self._specs, self._options_bootstrapper)]\n        )\n\n        expanded_targets = self._scheduler.product_request(\n            Targets, [Params(Addresses([addr])) for addr in unexpanded_addresses]\n        )\n        targets_dict: dict[str, list[TargetInfo]] = {}\n        for addr, targets in zip(unexpanded_addresses, expanded_targets):\n            targets_dict[addr.spec] = [\n                TargetInfo(\n                    filename=(\n                        tgt.address.filename if tgt.address.is_file_target else str(tgt.address)\n                    )\n                )\n                for tgt in targets\n            ]\n        return ExpandedSpecs(targets=targets_dict)\n\n\nclass WorkunitsCallback(ABC):\n    @abstractmethod\n    def __call__(\n        self,\n        *,\n        started_workunits: tuple[Workunit, ...],\n        completed_workunits: tuple[Workunit, ...],\n        finished: bool,\n        context: StreamingWorkunitContext,\n    ) -> None:\n        \"\"\"\n        :started_workunits: Workunits that have started but not completed.\n        :completed_workunits: Workunits that have completed.\n        :finished: True when the last chunk of workunit data is reported to the callback.\n        :context: A context providing access to functionality relevant to the run.\n        \"\"\"\n\n    @property\n    @abstractmethod\n    def can_finish_async(self) -> bool:\n        \"\"\"Can this callback finish its work in the background after the Pants run has already\n        completed?\n\n        The main reason to `return False` is if your callback logs in its final call, when\n        `finished=True`, as it may end up logging to `.pantsd.d/pants.log` instead of the console,\n        which is harder for users to find. Otherwise, most callbacks should return `True` to avoid\n        slowing down Pants from finishing the run.\n        \"\"\"\n\n\n@dataclass(frozen=True)\nclass WorkunitsCallbackFactory:\n    \"\"\"A wrapper around a callable that constructs WorkunitsCallbacks.\n\n    NB: This extra wrapping is because subtyping is not supported in the return position of a\n    rule. See #11354 for discussion of that limitation.\n    \"\"\"\n\n    callback_factory: Callable[[], WorkunitsCallback | None]\n\n\nclass WorkunitsCallbackFactories(Tuple[WorkunitsCallbackFactory, ...]):\n    \"\"\"A list of registered factories for WorkunitsCallback instances.\"\"\"\n\n\n@union\nclass WorkunitsCallbackFactoryRequest:\n    \"\"\"A request for a particular WorkunitsCallbackFactory.\"\"\"\n\n\n@rule\nasync def construct_workunits_callback_factories(\n    union_membership: UnionMembership,\n) -> WorkunitsCallbackFactories:\n    request_types = union_membership.get(WorkunitsCallbackFactoryRequest)\n    workunit_callback_factories = await MultiGet(\n        Get(WorkunitsCallbackFactory, WorkunitsCallbackFactoryRequest, request_type())\n        for request_type in request_types\n    )\n    return WorkunitsCallbackFactories(workunit_callback_factories)\n\n\n# -----------------------------------------------------------------------------------------------\n# Streaming workunits handler\n# -----------------------------------------------------------------------------------------------\n\n\nclass StreamingWorkunitHandler:\n    \"\"\"Periodically calls each registered WorkunitsCallback in a dedicated thread.\n\n    This class should be used as a context manager.\n    \"\"\"\n\n    def __init__(\n        self,\n        scheduler: SchedulerSession,\n        run_tracker: RunTracker,\n        callbacks: Iterable[WorkunitsCallback],\n        options_bootstrapper: OptionsBootstrapper,\n        specs: Specs,\n        report_interval_seconds: float,\n        allow_async_completion: bool,\n        max_workunit_verbosity: LogLevel,\n    ) -> None:\n        scheduler = scheduler.isolated_shallow_clone(\"streaming_workunit_handler_session\")\n        self.callbacks = callbacks\n        self.context = StreamingWorkunitContext(\n            _scheduler=scheduler,\n            _run_tracker=run_tracker,\n            _specs=specs,\n            _options_bootstrapper=options_bootstrapper,\n        )\n        self.thread_runner = (\n            _InnerHandler(\n                scheduler=scheduler,\n                context=self.context,\n                callbacks=self.callbacks,\n                report_interval=report_interval_seconds,\n                # TODO(10092) The max verbosity should be a per-client setting, rather than a global\n                #  setting.\n                max_workunit_verbosity=max_workunit_verbosity,\n                allow_async_completion=allow_async_completion,\n            )\n            if callbacks\n            else None\n        )\n\n    def __enter__(self) -> None:\n        if not self.thread_runner:\n            return\n        self.thread_runner.start()\n\n    def __exit__(self, exc_type, exc_value, traceback) -> None:\n        if not self.thread_runner:\n            return\n        self.thread_runner.end()\n        if exc_type is not None:\n            self.thread_runner.join()\n\n\nclass _InnerHandler(threading.Thread):\n    def __init__(\n        self,\n        scheduler: Any,\n        context: StreamingWorkunitContext,\n        callbacks: Iterable[WorkunitsCallback],\n        report_interval: float,\n        max_workunit_verbosity: LogLevel,\n        allow_async_completion: bool,\n    ) -> None:\n        super().__init__(daemon=True)\n        self.scheduler = scheduler\n        self.context = context\n        self.stop_request = threading.Event()\n        self.report_interval = report_interval\n        self.callbacks = callbacks\n        self.max_workunit_verbosity = max_workunit_verbosity\n        # TODO: Have a thread per callback so that some callbacks can always finish async even\n        #  if others must be finished synchronously.\n        self.block_until_complete = not allow_async_completion or any(\n            callback.can_finish_async is False for callback in self.callbacks\n        )\n        # Get the parent thread's thread locals. Note that this thread has not yet started\n        # as we are only in the constructor.\n        self.thread_locals = PyThreadLocals.get_for_current_thread()\n\n    def poll_workunits(self, *, finished: bool) -> None:\n        workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)\n        for callback in self.callbacks:\n            callback(\n                started_workunits=workunits[\"started\"],\n                completed_workunits=workunits[\"completed\"],\n                finished=finished,\n                context=self.context,\n            )\n\n    def run(self) -> None:\n        # First, set the thread's thread locals to the parent thread's in order to propagate the\n        # console, workunit stores, etc.\n        self.thread_locals.set_for_current_thread()\n        while not self.stop_request.isSet():  # type: ignore[attr-defined]\n            self.poll_workunits(finished=False)\n            self.stop_request.wait(timeout=self.report_interval)\n        else:\n            # Make one final call. Note that this may run after the Pants run has already\n            # completed, depending on whether the thread was joined or not.\n            self.poll_workunits(finished=True)\n\n    def end(self) -> None:\n        self.stop_request.set()\n        if self.block_until_complete:\n            logger.debug(\n                \"Async completion is disabled: waiting for workunit callbacks to complete...\"\n            )\n            super().join()\n        else:\n            logger.debug(\n                \"Async completion is enabled: workunit callbacks will complete in the background.\"\n            )\n\n\ndef rules():\n    return [\n        QueryRule(WorkunitsCallbackFactories, (UnionMembership,)),\n        QueryRule(Targets, (Addresses,)),\n        QueryRule(Addresses, (Specs, OptionsBootstrapper)),\n        *collect_rules(),\n    ]\n", "path": "src/python/pants/engine/streaming_workunit_handler.py"}]}