Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 21 additions & 13 deletions zstash/extract.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import logging
import multiprocessing
import os.path
import queue
import re
import sqlite3
import sys
Expand Down Expand Up @@ -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())

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolves #402 (comment)

except queue.Empty:
pass

manager.shutdown()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolves #402 (comment)


# Sort the failures, since they can come in at any order.
failures.sort(key=lambda t: (t.name, t.tar, t.offset))
Expand Down Expand Up @@ -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
)
Comment on lines -552 to -554

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolves #402 (comment)

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))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolves #402 (comment)


tries: int = args.retries + 1
# Set to True to test the `--retries` option with a forced failure.
Expand Down Expand Up @@ -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)
Expand Down
19 changes: 11 additions & 8 deletions zstash/parallel.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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:
Expand Down