zephyr_mirror: Port sharding wrapper to asyncio.

andersk · andersk · commit 02586f1d348d · 2023-01-25T16:53:09.000-08:00
Signed-off-by: Anders Kaseorg &lt;anders@zulip.com&gt;
diff --git a/zulip/integrations/zephyr/zephyr_mirror.py b/zulip/integrations/zephyr/zephyr_mirror.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import asyncio
 import os
 import signal
 import subprocess
@@ -32,21 +33,22 @@ def die(signal: int, frame: FrameType) -> None:
     sys.exit(0)
 
 if options.forward_class_messages and not options.noshard:
-    # Needed to get access to zephyr.lib.parallel
-    sys.path.append("/home/zulip/zulip")
     if options.on_startup_command is not None:
         subprocess.call([options.on_startup_command])
-    from zerver.lib.parallel import run_parallel
 
     print("Starting parallel zephyr class mirroring bot")
-    jobs = list("0123456789abcdef")
+    shards = list("0123456789abcdef")
 
-    def run_job(shard: str) -> int:
-        subprocess.call(args + [f"--shard={shard}"])
-        return 0
+    async def run_shard(shard: str) -> int:
+        process = await asyncio.create_subprocess_exec(*args, f"--shard={shard}")
+        return await process.wait()
 
-    for (status, job) in run_parallel(run_job, jobs, threads=16):
-        print("A mirroring shard died!")
+    async def run_shards():
+        for coro in asyncio.as_completed(map(run_shard, shards)):
+            await coro
+            print("A mirroring shard died!")
+
+    asyncio.run(run_shards())
     sys.exit(0)
 
 backoff = RandomExponentialBackoff(timeout_success_equivalent=300)
