Waitgroups in Python

To wait for multiple threads to finish, we can use a ThreadPoolExecutor with a Future list in Python.

import time
from concurrent.futures import ThreadPoolExecutor, as_completed

# This is the function we'll run in every thread.
def worker(id):
    print(f"Worker {id} starting")
    
    # Sleep to simulate an expensive task.
    time.sleep(1)
    print(f"Worker {id} done")

def main():
    # This ThreadPoolExecutor is used to manage the thread pool
    with ThreadPoolExecutor(max_workers=5) as executor:
        # Launch several threads and store their futures
        futures = []
        for i in range(1, 6):
            # Submit the worker function to the executor
            future = executor.submit(worker, i)
            futures.append(future)
        
        # Wait for all futures to complete
        for future in as_completed(futures):
            future.result()

    # Note that this approach allows for straightforward error propagation.
    # If a worker raises an exception, it will be re-raised here when calling future.result()

if __name__ == "__main__":
    main()

To run the program:

$ python threadpool.py
Worker 1 starting
Worker 2 starting
Worker 3 starting
Worker 4 starting
Worker 5 starting
Worker 1 done
Worker 2 done
Worker 3 done
Worker 4 done
Worker 5 done

The order of workers starting up and finishing is likely to be different for each invocation.

In this Python version, we use the concurrent.futures module, which provides a high-level interface for asynchronously executing callables. The ThreadPoolExecutor manages a pool of worker threads, similar to how Go manages goroutines.

Instead of a WaitGroup, we use the Future objects returned by executor.submit(). These Future objects represent the asynchronous execution of our worker function. We collect all these futures and then use as_completed() to iterate over them as they complete.

This approach provides a way to wait for all threads to complete, similar to the WaitGroup.Wait() in the original example. It also allows for easy error propagation: if a worker raises an exception, it will be re-raised when we call future.result().

The with statement ensures that the ThreadPoolExecutor is properly shut down after all tasks are complete, similar to how Go manages the lifecycle of goroutines.