From 91863175d4181d085f49180b2bb3c216cfc1e051 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 25 Feb 2026 18:44:18 +0000
Subject: [PATCH 1/2] Initial plan


From db9b68c87e9eda79acb0cdb69d538042c5c3800f Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 25 Feb 2026 18:52:44 +0000
Subject: [PATCH 2/2] Address code review: remove wait_turn from extraction,
 fix HPSS auth, manager cleanup, queue drain, precomputed tar index map

Co-authored-by: forsyth2 <30700190+forsyth2@users.noreply.github.com>
---
 zstash/extract.py  | 34 +++++++++++++++++++++-------------
 zstash/parallel.py | 19 +++++++++++--------
 2 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/zstash/extract.py b/zstash/extract.py
index 4ee91ee3..c72432fa 100644
--- a/zstash/extract.py
+++ b/zstash/extract.py
@@ -7,6 +7,7 @@
 import logging
 import multiprocessing
 import os.path
+import queue
 import re
 import sqlite3
 import sys
@@ -389,12 +390,21 @@ def multiprocess_extract(
     # because we'll be in this loop until completion.
     failures: List[FilesRow] = []
     while any(p.is_alive() for p in processes):
-        while not failure_queue.empty():
-            failures.append(failure_queue.get())
+        try:
+            while True:
+                failures.append(failure_queue.get_nowait())
+        except queue.Empty:
+            pass
         time.sleep(0.01)
 
-    while not failure_queue.empty():
-        failures.append(failure_queue.get())
+    # Drain any remaining failures after all processes have exited.
+    try:
+        while True:
+            failures.append(failure_queue.get_nowait())
+    except queue.Empty:
+        pass
+
+    manager.shutdown()
 
     # Sort the failures, since they can come in at any order.
     failures.sort(key=lambda t: (t.name, t.tar, t.offset))
@@ -547,18 +557,16 @@ def extractFiles(  # noqa: C901
             # let the process know.
             # This is to synchronize the print statements.
 
-            # Wait for turn before processing this tar
             if multiprocess_worker:
-                multiprocess_worker.print_monitor.wait_turn(
-                    multiprocess_worker, files_row.tar
-                )
                 multiprocess_worker.set_curr_tar(files_row.tar)
 
-            # Use args.hpss directly
+            # Use args.hpss, falling back to config.hpss when not provided
             if args.hpss is not None:
                 hpss: str = args.hpss
+            elif config.hpss is not None:
+                hpss = config.hpss
             else:
-                raise TypeError("Invalid args.hpss={}".format(args.hpss))
+                raise TypeError("Invalid config.hpss={}".format(config.hpss))
 
             tries: int = args.retries + 1
             # Set to True to test the `--retries` option with a forced failure.
@@ -735,9 +743,9 @@ def extractFiles(  # noqa: C901
         cur.close()
         con.close()
 
-        # Add the failures to the queue.
-        # When running with multiprocessing, the function multiprocess_extract()
-        # that calls this extractFiles() function will return the failures as a list.
+    # Add the failures to the queue.
+    # When running with multiprocessing, the function multiprocess_extract()
+    # that calls this extractFiles() function will return the failures as a list.
     if multiprocess_worker:
         for f in failures:
             multiprocess_worker.failure_queue.put(f)
diff --git a/zstash/parallel.py b/zstash/parallel.py
index 4470a399..4bcda598 100644
--- a/zstash/parallel.py
+++ b/zstash/parallel.py
@@ -38,6 +38,11 @@ def __init__(self, tars_to_print: List[str], manager=None, *args, **kwargs):
         # Store the ordered list of tars
         self._tars_list: List[str] = tars_to_print
 
+        # Precomputed mapping from tar name to its position in the ordered list.
+        self._tar_to_index: Dict[str, int] = {
+            tar: i for i, tar in enumerate(tars_to_print)
+        }
+
         # Use a simple counter to track which tar we're on
         self._current_tar_index: multiprocessing.managers.ValueProxy = manager.Value(
             "i", 0
@@ -52,10 +57,9 @@ def wait_turn(
         """
         Wait until it's this worker's turn to process workers_curr_tar.
         """
-        try:
-            tar_index = self._tars_list.index(workers_curr_tar)
-        except ValueError:
-            return
+        if workers_curr_tar not in self._tar_to_index:
+            raise RuntimeError("Tar {} not in ordered list".format(workers_curr_tar))
+        tar_index = self._tar_to_index[workers_curr_tar]
 
         attempted = False
         while True:
@@ -75,10 +79,9 @@ def done_enqueuing_output_for_tar(
         A worker has finished printing output for workers_curr_tar.
         Advance to the next tar in the sequence.
         """
-        try:
-            tar_index = self._tars_list.index(workers_curr_tar)
-        except ValueError:
-            return
+        if workers_curr_tar not in self._tar_to_index:
+            raise RuntimeError("Tar {} not in ordered list".format(workers_curr_tar))
+        tar_index = self._tar_to_index[workers_curr_tar]
 
         with self._lock:
             if self._current_tar_index.value == tar_index: