Test async completion.

# Building wheels and fs_util will be skipped. Delete if not intended.
[ci skip-build-wheels]

src/python/pants/engine/BUILD

@@ -17,4 +17,5 @@ python_tests(
       # Loaded reflectively as a backend in `streaming_workunit_handler_integration_test.py`.
       "testprojects/pants-plugins/src/python/workunit_logger",
   ],
+  timeout=120,
 )

src/python/pants/engine/streaming_workunit_handler_integration_test.py

@@ -35,14 +35,19 @@ def run(
         dest = os.path.join(tmpdir, "dest.log")
         normalized_args = [arg.format(tmpdir=tmpdir) for arg in args]
         pants_run = run_pants(normalized_args, config=workunit_logger_config(dest))
-        log_content = maybe_read_file(dest)
         if success:
             pants_run.assert_success()
-            assert log_content
-            assert FINISHED_SUCCESSFULLY in log_content
+            confirm_eventual_success(dest)
         else:
             pants_run.assert_failure()
-        return pants_run, log_content
+        return pants_run, maybe_read_file(dest)
+
+
+def confirm_eventual_success(log_dest: str) -> None:
+    for _ in attempts("The log should eventually show that the SWH shut down."):
+        content = maybe_read_file(log_dest)
+        if content and FINISHED_SUCCESSFULLY in content:
+            break
 
 
 def test_list() -> None:
@@ -64,7 +69,9 @@ def test_ctrl_c() -> None:
         os.kill(client_pid, signal.SIGINT)
 
         # Confirm that finish is still called (even though it may be backgrounded in the server).
-        for _ in attempts("The log should eventually show that the SWH shut down."):
-            content = maybe_read_file(dest)
-            if content and FINISHED_SUCCESSFULLY in content:
-                break
+        confirm_eventual_success(dest)
+
+
+def test_restart() -> None:
+    # Will trigger a restart
+    run(["--pantsd-max-memory-usage=1", "roots"])

testprojects/pants-plugins/src/python/workunit_logger/register.py

@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import logging
+import time
 from dataclasses import dataclass
 from typing import Tuple
 
@@ -41,8 +42,7 @@ class WorkunitsLogger(WorkunitsCallback):
 
     @property
     def can_finish_async(self) -> bool:
-        # We'd like to synchronously fail the run on the final call if need be.
-        return False
+        return True
 
     def __call__(
         self,
@@ -55,6 +55,8 @@ def __call__(
         with open(self.dest, "a") as dest:
             print(str(completed_workunits), file=dest)
             if finished and context.run_tracker.has_ended():
+                # Sleep a little while to ensure that we're finishing asynchronously.
+                time.sleep(2)
                 print(FINISHED_SUCCESSFULLY, file=dest)