FastAPI异步应用中ProcessPoolExecutor的高效管理策略

@app.post("/addContent")
async def add_content(content:dict):
  # ...
  with concurrent.futures.ProcessPoolExecutor(max_workers=nworkers) as executor:
    # ...
    await asyncio.gather(*async_tasks)

import asyncio
import concurrent.futures
import functools
import re
from contextlib import asynccontextmanager
from fastapi import FastAPI

# 全局进程池变量
process_pool: concurrent.futures.ProcessPoolExecutor = None

@asynccontextmanager
async def executor_pool_lifespan(app: FastAPI):
    """
    FastAPI应用生命周期管理器，用于初始化和关闭ProcessPoolExecutor。
    """
    global process_pool
    # 根据CPU核心数和任务特性调整工作进程数
    # 通常，对于CPU密集型任务，max_workers不应超过CPU核心数。
    # 但考虑到I/O等待或网络延迟，适当增加可能带来收益，需监控生产环境。
    nworkers = 18 
    process_pool = concurrent.futures.ProcessPoolExecutor(max_workers=nworkers)
    try:
        yield # 在此点，FastAPI将启动服务器并运行应用程序
    finally:
        # 当服务器停止时，安全关闭进程池
        process_pool.shutdown()
        print("ProcessPoolExecutor has been shut down.")

# 初始化FastAPI应用时，传入lifespan参数
app = FastAPI(lifespan=executor_pool_lifespan)

# 辅助函数：将阻塞任务提交到executor
async def run_in_process_executor(fn, executor=None):
    """
    将一个可调用对象提交到指定的executor中运行，并等待其结果。
    """
    if executor is None:
        # 如果未指定executor，则使用全局进程池
        executor = process_pool
    if executor is None:
        raise RuntimeError("ProcessPoolExecutor is not initialized.")
    event_loop = asyncio.get_event_loop()
    return await event_loop.run_in_executor(executor, fn)

# 示例：CPU密集型任务函数
def run_regex_on_content_chunk(content: str):
    """
    在一个内容块上运行正则表达式，提取域名。
    这是一个CPU密集型任务的示例。
    """
    domains = []
    domain_patt = re.compile(r'([a-zA-Z0-9\-_]+\.){1,}[a-zA-Z0-9\-_]+')
    for match in domain_patt.finditer(content):
        domains.append(match.group(0))
    return domains

# 示例：内容分块函数
def split_on_whitespace(content: str, count: int = 6):
    """
    根据空白字符将内容分割成大致相等的部分。
    这是一个简化的示例，实际分割逻辑可能更复杂。
    """
    if not content:
        return ['' for _ in range(count)]

    chunks = []
    current_start = 0
    total_len = len(content)
    # 尝试按比例分割，并确保在空白处断开
    for i in range(count):
        target_end = (i + 1) * total_len // count
        if target_end >= total_len:
            chunks.append(content[current_start:])
            break

        # 寻找最近的空白字符作为分割点
        split_point = content.rfind(' ', current_start, target_end)
        if split_point == -1 or split_point <= current_start: # 找不到空白或空白在当前块开头
            split_point = target_end # 强制分割

        chunks.append(content[current_start:split_point].strip())
        current_start = split_point + 1

    # 填充剩余块或处理最后一个块
    while len(chunks) < count:
        chunks.append('')

    return chunks[:count]

@app.post("/addContent")
async def add_content(content: dict):
    all_content = content['data']
    nworkers = 6 # 这里的nworkers现在仅用于决定分块数量，而非进程池大小

    # 确保进程池已初始化
    if process_pool is None:
        raise RuntimeError("ProcessPoolExecutor is not initialized. Server might not be running correctly.")

    content_chunks = split_on_whitespace(all_content, nworkers) # 根据分块数调整
    async_tasks = []

    for chunk in content_chunks:
        # 使用functools.partial封装任务函数及其参数
        regex_fn = functools.partial(run_regex_on_content_chunk, chunk)
        # 将任务提交到全局进程池
        async_tasks.append(run_in_process_executor(regex_fn, process_pool))

    # 并发等待所有任务完成
    results = await asyncio.gather(*async_tasks)

    # 汇总结果（示例）
    all_domains = [domain for sublist in results for domain in sublist]
    return {"status": "success", "domains_found": all_domains, "count": len(all_domains)}