asyncio 协程动态添加任务、支持阻塞任务

待我称王封你为后i 2022-12-29 05:00 354阅读 0赞

## 介绍 ##

最后一版，支持阻塞的任务，但需要线程池进行支持，阻塞任务会被放置在线程池中执行。

1.  支持动态添加任务
2.  支持停止事件循环
3.  支持最大协程数：maxsize
4.  支持进度条
5.  实时获取剩余协程数
6.  支持阻塞协程，需要线程进行支持，注意设置线程池：pool\_maxsize

## asyncPool.py ##

# -*- coding:utf-8 -*-
    import asyncio
    import aiohttp
    import time
    import queue
    from concurrent.futures import ThreadPoolExecutor
    
    
    class AsyncPool(object):
        """
        1. 支持动态添加任务
        2. 支持停止事件循环
        3. 支持最大协程数：maxsize
        4. 支持进度条
        5. 实时获取剩余协程数
        6. 支持阻塞协程，需要线程进行支持，注意设置线程池：pool_maxsize
        """
    
        def __init__(self, maxsize=1, pool_maxsize=None, loop=None):
            """
            初始化
            :param loop: 事件循环对象
            :param maxsize: 默认为1
            :param pool_maxsize: 默认为系统内核数
            """
            # 在jupyter需要这个，不然asyncio运行出错
            # import nest_asyncio
            # nest_asyncio.apply()
    
            # 任务计数器
            self._task, self._task_done_count, self._task_progress_list = self.task_create()
    
            # 协程池
            self._loop, self._loop_thread, self._thread_pool = self.start_loop(
                loop,
                pool_maxsize=pool_maxsize)
    
            # 限制并发量为500
            self.semaphore = asyncio.Semaphore(maxsize, loop=self._loop)
    
        @staticmethod
        def task_create():
            """
            创建任务对象
    
            :param
            :return:
            """
            # 队列，先进先出，根据队列是否为空判断，退出协程
            # self.task_done_count: 任务完成后的计数器
            # 任务组，用来阻塞任务用的
            task = queue.Queue()
    
            # 任务完成后的计数器
            task_done_count = queue.Queue()
    
            # 任务进度值存储
            # 确保任务值是唯一的
            task_progress_list = []
    
            return task, task_done_count, task_progress_list
    
        def task_add(self, item=1):
            """
            添加任务
            :param item:
            :return:
            """
            self._task.put(item)
    
        def task_done(self, fn):
            """
            任务完成
            回调函数
            已完成任务计数
    
            :param fn:
            :return:
            """
            if fn:
                pass
            self._task.get()
            self._task.task_done()
    
            # 已完成任务计数
            self._task_done_count.put(1)
    
        def task_progress(self, total):
            """
            任务进度条
            适用于已知任务总数的情况
    
            :param total: 任务总数
            :return:
            """
            # 任务完成计数
            # 需要剔除任务停止的任务数
            count = self._task_done_count.qsize() - 1
            if count < 0:
                count = 0
    
            item = int(count / total * 100)
            if count == total:
                # 任务完成
                self._task_progress_list.append(100)
            elif item not in self._task_progress_list:
                # 过程值
                self._task_progress_list.append(item)
            else:
                pass
    
            self._task_progress_list = list(set(self._task_progress_list))
            self._task_progress_list.sort()
            return self._task_progress_list[-1]
    
        def get_task(self):
            """
            获取事件循环任务列表
            :return:
            """
            task_list = asyncio.Task.all_tasks(loop=self._loop)
            # task_list = asyncio.all_tasks(loop=self._loop)
            return task_list
    
        def wait(self):
            """
            任务阻塞
            等待所有任务执行完毕
    
            :return:
            """
            self._task.join()
            # self._thread_pool.shutdown()
    
        @property
        def running(self):
            """
            获取剩余协程数
            :return:
            """
            # 剩余任务数
            # 需要剔除任务停止的任务数
            count = self._task.qsize() - 1
            if count < 0:
                count = 0
            return count
    
        @staticmethod
        def _start_thread_loop(loop):
            """
            运行事件循环
            :param loop: loop以参数的形式传递进来运行
            :return:
            """
            # 将当前上下文的事件循环设置为循环。
            asyncio.set_event_loop(loop)
            # 开始事件循环
            loop.run_forever()
    
        async def _stop_thread_loop(self, loop_time=1):
            """
            停止协程
            关闭线程
            :return:
            """
            while True:
                if self._task.empty():
                    # 停止协程
                    self._loop.stop()
                    break
                await asyncio.sleep(loop_time)
    
        def start_loop(self, loop, pool_maxsize=None):
            """
            运行事件循环
            开启新线程
            :param loop: 协程
            :param pool_maxsize: 线程池大小，默认为系统内核数
            :return:
            """
            # 获取一个事件循环
            if not loop:
                loop = asyncio.new_event_loop()
    
            # 线程池
            thread_pool = ThreadPoolExecutor(pool_maxsize)
    
            # 设置线程池
            loop.set_default_executor(thread_pool)
    
            # from threading import Thread
            # thread_pool = Thread(target=self._start_thread_loop, args=(loop,))
            # 设置守护进程
            # thread_pool.setDaemon(True)
            # 运行线程，同时协程事件循环也会运行
            # thread_pool.start()
    
            # 启动子线程
            # 协程开始启动
            thread_pool.submit(self._start_thread_loop, loop)
    
            return loop, thread_pool, thread_pool
    
        def stop_loop(self, loop_time=1):
            """
            队列为空，则关闭线程
            :param loop_time:
            :return:
            """
            # 关闭线程任务
            asyncio.run_coroutine_threadsafe(self._stop_thread_loop(loop_time), self._loop)
    
            # 取消单个任务
            # task.cancel()
    
        def _close(self):
            """
            关闭事件循环，不然会重启
            会导致错误
            # RuntimeError: Cannot close a running event loop
            :return:
            """
            self._loop.close()
    
        def release(self, loop_time=1):
            """
            释放线程
            :param loop_time:
            :return:
            """
            self.stop_loop(loop_time)
    
        async def async_semaphore_func(self, func):
            """
            信号代理
            :param func:
            :return:
            """
            async with self.semaphore:
                return await func
    
        async def async_thread_pool_func(self, block_func, *args, thread_pool=True):
            """
            信号代理
            线程池代理
    
            loop: asyncio.AbstractEventLoop
    
            :param block_func: 阻塞任务
            :param args: 参数
            :param thread_pool: 是否使用自定义线程池
            :return:
            """
            async with self.semaphore:
                if not thread_pool:
                    # 默认线程池
                    return await self._loop.run_in_executor(None, block_func, *args)
    
                # 使用自定义线程池
                future = await self._loop.run_in_executor(self._thread_pool, block_func, *args)
                # gather = await asyncio.gather(future)  # 阻塞
                # gather = await asyncio.wait(future)  # 阻塞
                return future
    
        def _submit(self, func, callback=None):
            """
            非阻塞模式
            提交任务到事件循环
    
            :param func: 异步函数对象
            :param callback: 回调函数
            :return:
            """
            self.task_add()
    
            # 将协程注册一个到运行在线程中的循环，thread_loop 会获得一个环任务
            # 注意：run_coroutine_threadsafe 这个方法只能用在运行在线程中的循环事件使用
            # future = asyncio.run_coroutine_threadsafe(func, self.loop)
            future = asyncio.run_coroutine_threadsafe(self.async_semaphore_func(func), self._loop)
    
            # 添加回调函数,添加顺序调用
            if callback:
                future.add_done_callback(callback)
            future.add_done_callback(self.task_done)
    
        def submit(self, func, *args, callback=None):
            """
            非阻塞模式
            提交任务到事件循环
    
            :param func: 异步函数对象
            :param args: 参数
            :param callback: 回调函数
            :return:
            """
            self.task_add()
    
            # 将协程注册一个到运行在线程中的循环，thread_loop 会获得一个环任务
            # 注意：run_coroutine_threadsafe 这个方法只能用在运行在线程中的循环事件使用
            # future = asyncio.run_coroutine_threadsafe(func, self.loop)
            future = asyncio.run_coroutine_threadsafe(self.async_semaphore_func(func(*args)), self._loop)
    
            # 添加回调函数,添加顺序调用
            if callback:
                future.add_done_callback(callback)
            future.add_done_callback(self.task_done)
    
        def submit2(self, func, *args, callback=None):
            """
            阻塞模式
            提交任务到事件循环
    
            :param func: 异步函数对象
            :param args: 入参
            :param callback: 回调函数
            :return:
            """
            self.task_add()
    
            # 将协程注册一个到运行在线程中的循环，thread_loop 会获得一个环任务
            # 注意：run_coroutine_threadsafe 这个方法只能用在运行在线程中的循环事件使用
    
            # future = self._loop.run_in_executor(None, func)
            # future = asyncio.ensure_future(self._loop.run_in_executor(None, func))  # 非阻塞
            future = asyncio.run_coroutine_threadsafe(
                self.async_thread_pool_func(func, *args),
                self._loop)
    
            # 添加回调函数,添加顺序调用
            if callback:
                future.add_done_callback(callback)
            future.add_done_callback(self.task_done)

## demo.py ##

import asyncio
    import time
    from asyncPool import AsyncPool
    
    
    async def thread_example2(i):
        await asyncio.sleep(1)
        return i
    
    
    async def thread_example3(i):
        time.sleep(1)
        await asyncio.sleep(0.1)
        return i
    
    
    def thread_example4(i):
        time.sleep(4)
        return i
    
    
    def my_callback(future):
        result = future.result()
        print('返回值: ', result)
    
    
    def demo():
        # 任务组， 最大协程数
        pool = AsyncPool(maxsize=1000, pool_maxsize=900)
    
        # 插入任务任务
        for i in range(3000 + 1):
            # 非阻塞协程
            # pool.submit(thread_example2, i, callback=my_callback)
            # 阻塞协程
            pool.submit2(thread_example4, i, callback=my_callback)
    
        print("等待子线程结束1...")
    
        # 停止子线程
        pool.release()
    
        print("等待子线程结束2...")
    
        # 等待
        # pool.wait()
    
        print("等待子线程结束3...")
    
        # 进度条使用
        while True:
            value = pool.task_progress(3000)
            print("value: ", value)
    
            # 获取线程数
            print("running: ", pool.running)
    
            if value == 100:
                break
    
            # 获取任务
            # print(pool.get_task())
            time.sleep(0.1)
    
        print("等待子线程结束4...")
    
    
    if __name__ == '__main__':
        start_time = time.time()
        demo()
        end_time = time.time()
        print("run time: ", end_time - start_time)