blob: e92d2c841ea56637b48b6f23a60180977195785b [file] [log] [blame]
#!/usr/bin/env python3
# Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
# for details. All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.
import sys
import threading
from threading import Thread
import traceback
# A thread that is given a list of jobs. The thread will repeatedly take a job
# from the list of jobs (which is shared with other threads) and execute it
# until there is no more jobs.
#
# If stop_on_first_failure is True, then the thread will exit upon the first
# failing job. The thread will then clear the jobs to ensure that all other
# workers will also terminate after completing there current job.
#
# Each job is a lambda that takes the worker_id as an argument. To guarantee
# termination each job must itself terminate (i.e., each job is responsible for
# setting an appropriate timeout).
class WorkerThread(Thread):
# The initialization of a WorkerThread is never run concurrently with the
# initialization of other WorkerThreads.
def __init__(self, jobs, jobs_lock, stop_on_first_failure, worker_id):
Thread.__init__(self)
self.completed = False
self.jobs = jobs
self.jobs_lock = jobs_lock
self.number_of_jobs = len(jobs)
self.stop_on_first_failure = stop_on_first_failure
self.success = True
self.worker_id = worker_id
def run(self):
print_thread("Starting worker", self.worker_id)
while True:
(job, job_id) = self.take_job(self.jobs, self.jobs_lock)
if job is None:
break
try:
print_thread(
"Starting job %s/%s" % (job_id, self.number_of_jobs),
self.worker_id)
exit_code = job(self.worker_id)
print_thread(
"Job %s finished with exit code %s" % (job_id, exit_code),
self.worker_id)
if exit_code:
self.success = False
if self.stop_on_first_failure:
self.clear_jobs(self.jobs, self.jobs_lock)
break
except:
print_thread("Job %s crashed" % job_id, self.worker_id)
print_thread(traceback.format_exc(), self.worker_id)
self.success = False
if self.stop_on_first_failure:
self.clear_jobs(self.jobs, self.jobs_lock)
break
print_thread("Exiting", self.worker_id)
self.completed = True
def take_job(self, jobs, jobs_lock):
jobs_lock.acquire()
job_id = self.number_of_jobs - len(jobs) + 1
job = jobs.pop(0) if jobs else None
jobs_lock.release()
return (job, job_id)
def clear_jobs(self, jobs, jobs_lock):
jobs_lock.acquire()
jobs.clear()
jobs_lock.release()
def run_in_parallel(jobs, number_of_workers, stop_on_first_failure):
assert number_of_workers > 0
if number_of_workers > len(jobs):
number_of_workers = len(jobs)
if number_of_workers == 1:
return run_in_sequence(jobs, stop_on_first_failure)
jobs_lock = threading.Lock()
threads = []
for worker_id in range(1, number_of_workers + 1):
threads.append(
WorkerThread(jobs, jobs_lock, stop_on_first_failure, worker_id))
for thread in threads:
thread.start()
for thread in threads:
thread.join()
for thread in threads:
if not thread.completed or not thread.success:
return 1
return 0
def run_in_sequence(jobs, stop_on_first_failure):
combined_exit_code = 0
worker_id = None
for job in jobs:
try:
exit_code = job(worker_id)
if exit_code:
combined_exit_code = exit_code
if stop_on_first_failure:
break
except:
print(traceback.format_exc())
combined_exit_code = 1
if stop_on_first_failure:
break
return combined_exit_code
def print_thread(msg, worker_id):
if worker_id is None:
print(msg)
else:
print('WORKER %s: %s' % (worker_id, msg))