Python `ThreadPoolExecutor`线程池

Python ThreadPoolExecutor线程池

1. 简介

在Python中，多线程常用于处理一些IO密集型的任务，以便在等待IO操作完成的同时，能够执行其他任务以提高效率。Python标准库中提供了ThreadPoolExecutor类，它可用于创建一个线程池，用来管理和调度线程的执行。本文将详细介绍ThreadPoolExecutor的使用。

2. `ThreadPoolExecutor`类的基本用法

ThreadPoolExecutor是concurrent.futures模块中的一个类，它提供了一个高级的接口来管理线程池和线程的执行。

首先，我们需要导入相关模块和类：

from concurrent.futures import ThreadPoolExecutor
import time

然后，我们可以创建一个ThreadPoolExecutor对象，指定线程池的大小：

with ThreadPoolExecutor(max_workers=3) as executor:
    ...

上述代码中，max_workers参数指定了线程池的大小。这里我们创建了一个大小为3的线程池。

接下来，我们可以通过submit()方法提交任务给线程池，并获得一个Future对象。Future对象代表将来完成的操作，它提供了一种异步编程的方式。我们可以使用Future.done()方法来检查任务是否完成，使用Future.result()方法来获取任务的结果。

下面是一个示例，我们在线程池中执行一些简单的任务，并获取它们的结果：

def task(name):
    print(f'Task {name} is running')
    time.sleep(2)  # 模拟任务执行的耗时
    return f'Task {name} is done'

with ThreadPoolExecutor(max_workers=3) as executor:
    future1 = executor.submit(task, 'A')
    future2 = executor.submit(task, 'B')
    future3 = executor.submit(task, 'C')

    print(f'Task A done? {future1.done()}')
    print(f'Task B done? {future2.done()}')
    print(f'Task C done? {future3.done()}')

    print(f'Task A result: {future1.result()}')
    print(f'Task B result: {future2.result()}')
    print(f'Task C result: {future3.result()}')

运行以上示例代码，输出结果如下：

Task A is running
Task B is running
Task C is running
Task A done? True
Task B done? True
Task C done? True
Task A result: Task A is done
Task B result: Task B is done
Task C result: Task C is done

可以看到，任务”A”、”B”和”C”都在不同的线程中执行，并分别返回了执行结果。

3. `map`方法的使用

如果需要批量提交任务，并且按照顺序获取任务的结果，可以使用map方法。map方法将函数和参数列表作为输入，依次执行函数，并返回执行结果的迭代器。

下面是一个示例，我们定义一个求平方的函数，并使用map方法来对一组数字进行平方计算：

def square(x):
    return x ** 2

numbers = [1, 2, 3, 4, 5]

with ThreadPoolExecutor(max_workers=3) as executor:
    results = executor.map(square, numbers)

    for result in results:
        print(result)

运行以上示例代码，输出结果如下：

可以看到，任务按照顺序依次执行，并返回了对应的平方结果。

4. 线程池的异常处理

ThreadPoolExecutor可以自动捕获任务中发生的异常，并将异常抛出到主线程。我们可以通过add_done_callback()方法注册回调函数来处理异常。

下面是一个示例，我们在任务中故意引发一个异常，并通过回调函数来处理异常：

def task(num):
    print(f'Task {num} is running')
    time.sleep(2)
    if num == 3:
        raise ValueError(f'Task {num} has an error')
    return f'Task {num} is done'

def handle_exception(future):
    try:
        future.result()
    except Exception as e:
        print(f'Exception: {e}')

with ThreadPoolExecutor(max_workers=3) as executor:
    future1 = executor.submit(task, 1)
    future2 = executor.submit(task, 2)
    future3 = executor.submit(task, 3)

    future1.add_done_callback(handle_exception)
    future2.add_done_callback(handle_exception)
    future3.add_done_callback(handle_exception)

运行以上示例代码，输出结果如下：

Task 1 is running
Task 2 is running
Task 3 is running
Exception: Task 3 has an error
Task 1 is done
Task 2 is done

可以看到，任务”1″和”2″都完成了，而任务”3″发生了异常，并被异常处理函数捕获和打印。

5. 线程池的关闭

当不再需要使用线程池时，可以通过调用shutdown()方法来关闭线程池，这将导致线程池中的所有任务被终止，而且无法再向线程池中提交新的任务。

下面是一个示例，我们使用shutdown()方法关闭线程池，并在关闭前等待所有任务完成：

def task(name):
    print(f'Task {name} is running')
    time.sleep(2)  # 模拟任务执行的耗时
    return f'Task {name} is done'

with ThreadPoolExecutor(max_workers=3) as executor:
    future1 = executor.submit(task, 'A')
    future2 = executor.submit(task, 'B')
    future3 = executor.submit(task, 'C')

    executor.shutdown(wait=True)

    print(f'Task A done? {future1.done()}')
    print(f'Task B done? {future2.done()}')
    print(f'Task C done? {future3.done()}')

    print(f'Task A result: {future1.result()}')
    print(f'Task B result: {future2.result()}')
    print(f'Task C result: {future3.result()}')

运行以上示例代码，输出结果如下：

Task A is running
Task B is running
Task C is running
Task A done? True
Task B done? True
Task C done? True
Task A result: Task A is done
Task B result: Task B is done
Task C result: Task C is done

可以看到，线程池被关闭后，任务”A”、”B”和”C”仍然能够正常完成。

6. 控制线程池的大小

在创建ThreadPoolExecutor对象时，我们可以通过max_workers参数来指定线程池的大小。但有时，我们可能需要根据任务的类型和数量动态地调整线程池的大小。

ThreadPoolExecutor类提供了两个方法来实现线程池的动态大小调整：allow_thread_shutdown()和adjust_thread_count()。

allow_thread_shutdown()方法允许正在运行的线程在完成后退出，但不会接受新的任务。
adjust_thread_count()方法可以动态地调整线程池的大小。

下面是一个示例，我们先创建一个大小为3的线程池，然后向线程池中提交6个任务，并在任务执行完毕后动态调整线程池的大小：

def task(name):
    print(f'Task {name} is running')
    time.sleep(2)  # 模拟任务执行的耗时
    return f'Task {name} is done'

with ThreadPoolExecutor(max_workers=3) as executor:
    futures = [executor.submit(task, i) for i in range(6)]

    for future in futures:
        future.result()

    executor.allow_thread_shutdown()
    executor._adjust_thread_count(1)

    print(f'Task A done? {futures[0].done()}')
    print(f'Task B done? {futures[1].done()}')
    print(f'Task C done? {futures[2].done()}')
    print(f'Task D done? {futures[3].done()}')
    print(f'Task E done? {futures[4].done()}')
    print(f'Task F done? {futures[5].done()}')

    print(f'Task A result: {futures[0].result()}')
    print(f'Task B result: {futures[1].result()}')
    print(f'Task C result: {futures[2].result()}')
    print(f'Task D result: {futures[3].result()}')
    print(f'Task E result: {futures[4].result()}')
    print(f'Task F result: {futures[5].result()}')

运行以上示例代码，输出结果如下：

Task 0 is running
Task 1 is running
Task 2 is running
Task 3 is running
Task 4 is running
Task 5 is running
Task 0 is done
Task 1 is done
Task 2 is done
Task 3 is done
Task 4 is done
Task 5 is done
Task A done? True
Task B done? True
Task C done? True
Task D done? True
Task E done? True
Task F done? True
Task A result: Task 0 is done
Task B result: Task 1 is done
Task C result: Task 2 is done
Task D result: Task 3 is done
Task E result: Task 4 is done
Task F result: Task 5 is done