Asyncio中将同步函数改为异步调用

发表于 | 分类于 | | 阅读次数: |
字数统计: 1,881 字 | 阅读时长 ≈ 10 分钟

今天看小明大神的博客：深入理解asyncio(三) 里面有段将同步函数改为协程使用的代码。其中提到了run_in_executor，主要使用这个方法将同步变为异步。

我们先看下如何将一个同步函数变为异步的

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
In [35]: import time

In [36]: import asyncio

In [37]: def a():
    ...:     time.sleep(1)
    ...:     return 'A'
    ...:

In [38]: async def c():
    ...:     loop = asyncio.get_running_loop()
    ...:     return await loop.run_in_executor(None, a)
    ...:

In [39]: asyncio.run(c())
Out[39]: 'A'

上面使用run_in_executor可以将同步函数a以协程的方式执行，我们看下源码

1
2
3
4
5
6
7
8
9
10
11
def run_in_executor(self, executor, func, *args):
    self._check_closed()
    if self._debug:
        self._check_callback(func, 'run_in_executor')
    if executor is None:
        executor = self._default_executor
        if executor is None:
            executor = concurrent.futures.ThreadPoolExecutor()
            self._default_executor = executor
    return futures.wrap_future(
        executor.submit(func, *args), loop=self)

我们看到 当没有设置 executor的时候 会默认使用concurrent.futures.ThreadPoolExecutor() 那我们自己设置一下试试看。

1
2
3
4
5
6
7
8
9
10
11
12
In [40]: from concurrent.futures import ThreadPoolExecutor

In [41]: thread_executor = ThreadPoolExecutor(5)

In [44]: async def c():
    ...:     loop = asyncio.get_running_loop()
    ...:
    ...:     return await loop.run_in_executor(thread_executor, a)
    ...:

In [45]: asyncio.run(c())
Out[45]: 'A'

正确输出

还有其他方法实现吗？

我们看到 run_in_executor的源码进行了loop是否关闭的校验和是否是debug的判断以及executor验空和赋值。假设我们不去做这些操作的话，直接使用ThreadPoolExecutor是否可以呢？

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
In [1]: from concurrent.futures import ThreadPoolExecutor

In [2]: import time

In [3]: import asyncio

In [4]: def a():
   ...:     time.sleep(1)
   ...:     return 'A'
   ...:

In [5]: thread_executor = ThreadPoolExecutor(5)

In [6]: async def c():
   ...:     future = thread_executor.submit(a)
   ...:     return await asyncio.wrap_future(future)
   ...:

In [7]: asyncio.run(c())
Out[7]: 'A'

正确输出

上面的thread_executor.submit返回的是一个future对象，但是并不是一个符合asyncio模块的future是不可等待的，即无法调用await去等待该对象。

代码中的wrap_future是一个比较关键的函数，看下源码

1
2
3
4
5
6
7
8
9
10
11
def wrap_future(future, *, loop=None):
    """Wrap concurrent.futures.Future object."""
    if isfuture(future):
        return future
    assert isinstance(future, concurrent.futures.Future), \
        f'concurrent.futures.Future is expected, got {future!r}'
    if loop is None:
        loop = events.get_event_loop()
    new_future = loop.create_future()
    _chain_future(future, new_future)
    return new_future

结合着使用到的_chain_future源码一起看

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def _chain_future(source, destination):
    """Chain two futures so that when one completes, so does the other.

    The result (or exception) of source will be copied to destination.
    If destination is cancelled, source gets cancelled too.
    Compatible with both asyncio.Future and concurrent.futures.Future.
    """
    if not isfuture(source) and not isinstance(source,
                                               concurrent.futures.Future):
        raise TypeError('A future is required for source argument')
    if not isfuture(destination) and not isinstance(destination,
                                                    concurrent.futures.Future):
        raise TypeError('A future is required for destination argument')
    source_loop = _get_loop(source) if isfuture(source) else None
    dest_loop = _get_loop(destination) if isfuture(destination) else None

    def _set_state(future, other):
        if isfuture(future):
            _copy_future_state(other, future)
        else:
            _set_concurrent_future_state(future, other)

    def _call_check_cancel(destination):
        if destination.cancelled():
            if source_loop is None or source_loop is dest_loop:
                source.cancel()
            else:
                source_loop.call_soon_threadsafe(source.cancel)

    def _call_set_state(source):
        if (destination.cancelled() and
                dest_loop is not None and dest_loop.is_closed()):
            return
        if dest_loop is None or dest_loop is source_loop:
            _set_state(destination, source)
        else:
            dest_loop.call_soon_threadsafe(_set_state, destination, source)

    destination.add_done_callback(_call_check_cancel)
    source.add_done_callback(_call_set_state)

通过wrap_future函数可以将concurrent.futures.Future变成asyncio.Future实现可等待。

这样我们可以不用显示的获取当前的loop也可以直接去将同步函数变成协程去执行了。

装饰器模式 将同步函数变为异步方式

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import asyncio
import functools

from concurrent.futures import ThreadPoolExecutor



class ThreadPool():

    def __init__(self, max_workers):

        self._thread_pool = ThreadPoolExecutor(max_workers)

    async def run(self, _callable, *args, **kwargs):

        future = self._thread_pool.submit(_callable, *args, **kwargs)

        return await asyncio.wrap_future(future)


class ThreadWorker:

    def __init__(self, max_workers):

        self._thread_pool = ThreadPool(max_workers)

    def __call__(self, func):

        @functools.wraps(func)
        def _wrapper(*args, **kwargs):
            return self._thread_pool.run(func, *args, **kwargs)

        return _wrapper


thread_worker = ThreadWorker(32)


@thread_worker
def some_io_block():
    return 1


asyncio.run(some_io_block()) 
# 输出 1

其中的 max_workers 参数就是能够执行的最大线程数。

再深入看下源码

上面我们知道可直接使用ThreadPoolExecutor的submit方法获取到future对象。通过源码来分析下具体的流程。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class ThreadPoolExecutor(_base.Executor):

    # Used to assign unique thread names when thread_name_prefix is not supplied.
    _counter = itertools.count().__next__

    def __init__(self, max_workers=None, thread_name_prefix='',
                 initializer=None, initargs=()):
        if max_workers is None:
            # Use this number because ThreadPoolExecutor is often
            # used to overlap I/O instead of CPU work.
            max_workers = (os.cpu_count() or 1) * 5
        if max_workers <= 0:
            raise ValueError("max_workers must be greater than 0")

        if initializer is not None and not callable(initializer):
            raise TypeError("initializer must be a callable")

        self._max_workers = max_workers
        self._work_queue = queue.SimpleQueue()
        self._threads = set()
        self._broken = False
        self._shutdown = False
        self._shutdown_lock = threading.Lock()
        self._thread_name_prefix = (thread_name_prefix or
                                    ("ThreadPoolExecutor-%d" % self._counter()))
        self._initializer = initializer
        self._initargs = initargs

    def submit(self, fn, *args, **kwargs):
        with self._shutdown_lock:
            if self._broken:
                raise BrokenThreadPool(self._broken)

            if self._shutdown:
                raise RuntimeError('cannot schedule new futures after shutdown')
            if _shutdown:
                raise RuntimeError('cannot schedule new futures after '
                                   'interpreter shutdown')

            f = _base.Future() # 这个 Future 就是  concurrent.futures 里的 Future
            w = _WorkItem(f, fn, args, kwargs)

            self._work_queue.put(w)
            self._adjust_thread_count()
            return f
    submit.__doc__ = _base.Executor.submit.__doc__

我们看到在submit函数里面主要是生成了一个concurrent.futures 里的 Future对象。然后一个_WorkItem实例。接着将生成的_WorkItem实例放到了一个队列里面，然后执行self._adjust_thread_count() 函数。这个_WorkItem实例是什么呢？self._adjust_thread_count()函数里面是什么呢？我们看下源码。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class _WorkItem(object):
    def __init__(self, future, fn, args, kwargs):
        self.future = future
        self.fn = fn
        self.args = args
        self.kwargs = kwargs

    def run(self):
        if not self.future.set_running_or_notify_cancel():
            return

        try:
            result = self.fn(*self.args, **self.kwargs)
        except BaseException as exc:
            self.future.set_exception(exc)
            # Break a reference cycle with the exception 'exc'
            self = None
        else:
            self.future.set_result(result)

我们发现在_WorkItem类中的run方法执行了真正的同步函数 并将执行结果或者异常放到了之前生成的future对象中。

那这个run方法什么时候真正执行呢？我们返回看self._adjust_thread_count()的源码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def _adjust_thread_count(self):
    # When the executor gets lost, the weakref callback will wake up
    # the worker threads.
    def weakref_cb(_, q=self._work_queue):
        q.put(None)
    # TODO(bquinlan): Should avoid creating new threads if there are more
    # idle threads than items in the work queue.
    num_threads = len(self._threads)
    if num_threads < self._max_workers:
        thread_name = '%s_%d' % (self._thread_name_prefix or self,
                                 num_threads)
        t = threading.Thread(name=thread_name, target=_worker,
                             args=(weakref.ref(self, weakref_cb),
                                   self._work_queue,
                                   self._initializer,
                                   self._initargs))
        t.daemon = True
        t.start()
        self._threads.add(t)
        _threads_queues[t] = self._work_queue

我们看到在_adjust_thread_count方法中生成了一个线程并且去执行了，执行的函数是_worker。这个_worker是什么呢？

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def _worker(executor_reference, work_queue, initializer, initargs):
    if initializer is not None:
        try:
            initializer(*initargs)
        except BaseException:
            _base.LOGGER.critical('Exception in initializer:', exc_info=True)
            executor = executor_reference()
            if executor is not None:
                executor._initializer_failed()
            return
    try:
        while True:
            work_item = work_queue.get(block=True)
            if work_item is not None:
                work_item.run()
                # Delete references to object. See issue16284
                del work_item
                continue
            executor = executor_reference()
            # Exit if:
            #   - The interpreter is shutting down OR
            #   - The executor that owns the worker has been collected OR
            #   - The executor that owns the worker has been shutdown.
            if _shutdown or executor is None or executor._shutdown:
                # Flag the executor as shutting down as early as possible if it
                # is not gc-ed yet.
                if executor is not None:
                    executor._shutdown = True
                # Notice other workers
                work_queue.put(None)
                return
            del executor
    except BaseException:
        _base.LOGGER.critical('Exception in worker', exc_info=True)

我们看到在这个函数中会获取到之前放到队列里面的_WorkItem的实例。然后执行_WorkItem里面的run方法。

这样我们整个过程就完整了，成功将一个同步函数变成异步方式执行。

知识就是财富
如果您觉得文章对您有帮助, 欢迎请我喝杯水！