VWED_server/services/execution/handlers/storage_queue_manager.py
2025-07-30 15:11:59 +08:00

432 lines
17 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
库位请求队列管理器
提供稳定的库位分配队列机制,避免高并发冲突
"""
import asyncio
import time
import json
import uuid
from typing import Dict, Any, List, Optional, Callable
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict
import heapq
from utils.logger import get_logger
from config.settings import settings
logger = get_logger("services.execution.handlers.storage_queue_manager")
class RequestPriority(Enum):
"""请求优先级"""
LOW = 1
NORMAL = 2
HIGH = 3
URGENT = 4
class RequestStatus(Enum):
"""请求状态"""
PENDING = "pending"
PROCESSING = "processing"
COMPLETED = "completed"
FAILED = "failed"
TIMEOUT = "timeout"
CANCELLED = "cancelled"
@dataclass
class StorageRequest:
"""库位请求"""
request_id: str
handler_type: str # "idle_site" 或 "idle_crowded_site"
input_params: Dict[str, Any]
context_data: Dict[str, Any]
map_id: str
task_record_id: str
priority: RequestPriority = RequestPriority.NORMAL
created_at: float = field(default_factory=time.time)
timeout: float = field(default_factory=lambda: settings.STORAGE_QUEUE_DEFAULT_TIMEOUT if settings.STORAGE_QUEUE_ENABLE_TIMEOUT else float('inf'))
retry_count: int = 0
max_retries: int = 3
status: RequestStatus = RequestStatus.PENDING
result: Optional[Dict[str, Any]] = None
error_message: Optional[str] = None
def __lt__(self, other):
"""用于优先级队列排序"""
if self.priority.value != other.priority.value:
return self.priority.value > other.priority.value # 高优先级优先
return self.created_at < other.created_at # 时间早的优先
class StorageQueueManager:
"""库位请求队列管理器"""
def __init__(self, max_workers: int = settings.STORAGE_QUEUE_MAX_WORKERS, max_queue_size: int = None):
self.max_workers = max_workers or settings.STORAGE_QUEUE_MAX_WORKERS
self.max_queue_size = max_queue_size or settings.STORAGE_QUEUE_MAX_SIZE
self.enable_timeout = settings.STORAGE_QUEUE_ENABLE_TIMEOUT
self.cleanup_interval = settings.STORAGE_QUEUE_CLEANUP_INTERVAL
self.completed_request_ttl = settings.STORAGE_QUEUE_COMPLETED_REQUEST_TTL
# 优先级队列
self.request_queue: List[StorageRequest] = []
self.queue_lock = asyncio.Lock()
# 请求跟踪
self.pending_requests: Dict[str, StorageRequest] = {}
self.processing_requests: Dict[str, StorageRequest] = {}
self.completed_requests: Dict[str, StorageRequest] = {}
# 工作者管理
self.workers: List[asyncio.Task] = []
self.worker_semaphore = asyncio.Semaphore(max_workers)
self.shutdown_event = asyncio.Event()
# 处理器注册
self.handlers: Dict[str, Callable] = {}
# 统计信息
self.stats = {
'requests_total': 0,
'requests_completed': 0,
'requests_failed': 0,
'requests_timeout': 0,
'avg_processing_time': 0,
'queue_size': 0,
'active_workers': 0
}
# 启动清理任务
self.cleanup_task = None
async def start(self):
"""启动队列管理器"""
logger.info("启动库位请求队列管理器")
# 启动工作者
for i in range(self.max_workers):
worker = asyncio.create_task(self._worker(f"worker-{i}"))
self.workers.append(worker)
# 启动清理任务
self.cleanup_task = asyncio.create_task(self._cleanup_completed_requests())
logger.info(f"队列管理器已启动,工作者数量: {self.max_workers}")
async def stop(self):
"""停止队列管理器"""
logger.info("停止库位请求队列管理器")
# 设置停止信号
self.shutdown_event.set()
# 等待所有工作者完成
if self.workers:
await asyncio.gather(*self.workers, return_exceptions=True)
# 停止清理任务
if self.cleanup_task:
self.cleanup_task.cancel()
try:
await self.cleanup_task
except asyncio.CancelledError:
pass
logger.info("队列管理器已停止")
def register_handler(self, handler_type: str, handler_func: Callable):
"""注册处理器"""
self.handlers[handler_type] = handler_func
logger.info(f"注册处理器: {handler_type}")
async def submit_request(self, request: StorageRequest) -> str:
"""提交请求到队列"""
if len(self.request_queue) >= self.max_queue_size:
raise Exception("队列已满,无法提交新请求")
async with self.queue_lock:
# 添加到优先级队列
heapq.heappush(self.request_queue, request)
self.pending_requests[request.request_id] = request
# 更新统计
self.stats['requests_total'] += 1
self.stats['queue_size'] = len(self.request_queue)
logger.info(f"提交请求到队列: {request.request_id}, 优先级: {request.priority.name}")
return request.request_id
async def get_request_status(self, request_id: str) -> Optional[Dict[str, Any]]:
"""获取请求状态"""
# 检查各个状态的请求
if request_id in self.pending_requests:
request = self.pending_requests[request_id]
return {
"request_id": request_id,
"status": request.status.value,
"created_at": request.created_at,
"queue_position": self._get_queue_position(request_id)
}
elif request_id in self.processing_requests:
request = self.processing_requests[request_id]
return {
"request_id": request_id,
"status": request.status.value,
"created_at": request.created_at,
"processing_time": time.time() - request.created_at
}
elif request_id in self.completed_requests:
request = self.completed_requests[request_id]
return {
"request_id": request_id,
"status": request.status.value,
"created_at": request.created_at,
"result": request.result,
"error_message": request.error_message,
"processing_time": time.time() - request.created_at
}
return None
async def wait_for_result(self, request_id: str, timeout: float = None) -> Dict[str, Any]:
"""等待请求结果"""
start_time = time.time()
# 根据配置决定是否使用超时
if not self.enable_timeout:
# 禁用超时,无限等待
check_timeout = float('inf')
logger.info(f"等待请求结果(无超时限制): {request_id}")
else:
# 启用超时,使用提供的超时时间或默认值
check_timeout = timeout or settings.STORAGE_QUEUE_DEFAULT_TIMEOUT
logger.info(f"等待请求结果(超时 {check_timeout}s: {request_id}")
while True:
# 检查是否完成
if request_id in self.completed_requests:
request = self.completed_requests[request_id]
if request.status == RequestStatus.COMPLETED:
return request.result
else:
raise Exception(f"请求失败: {request.error_message}")
# 只有在启用超时时才检查超时
if self.enable_timeout:
# 检查等待超时
if time.time() - start_time >= check_timeout:
await self._mark_request_timeout(request_id)
raise Exception("等待结果超时")
# 检查请求本身是否超时
if request_id in self.pending_requests:
request = self.pending_requests[request_id]
if request.timeout != float('inf') and time.time() - request.created_at > request.timeout:
await self._mark_request_timeout(request_id)
raise Exception("请求超时")
await asyncio.sleep(0.1) # 避免忙等待
async def cancel_request(self, request_id: str) -> bool:
"""取消请求"""
async with self.queue_lock:
if request_id in self.pending_requests:
request = self.pending_requests[request_id]
request.status = RequestStatus.CANCELLED
# 从队列中移除
self.request_queue = [r for r in self.request_queue if r.request_id != request_id]
heapq.heapify(self.request_queue)
# 移动到完成队列
del self.pending_requests[request_id]
self.completed_requests[request_id] = request
self.stats['queue_size'] = len(self.request_queue)
logger.info(f"取消请求: {request_id}")
return True
return False
def get_queue_stats(self) -> Dict[str, Any]:
"""获取队列统计信息"""
self.stats['queue_size'] = len(self.request_queue)
self.stats['active_workers'] = sum(1 for w in self.workers if not w.done())
return self.stats.copy()
async def _worker(self, worker_name: str):
"""工作者协程"""
logger.info(f"工作者 {worker_name} 启动")
while not self.shutdown_event.is_set():
try:
# 获取请求
request = await self._get_next_request()
if not request:
await asyncio.sleep(0.1)
continue
# 处理请求
async with self.worker_semaphore:
await self._process_request(request, worker_name)
except Exception as e:
logger.error(f"工作者 {worker_name} 处理请求异常: {str(e)}")
await asyncio.sleep(1)
logger.info(f"工作者 {worker_name} 停止")
async def _get_next_request(self) -> Optional[StorageRequest]:
"""获取下一个请求"""
async with self.queue_lock:
while self.request_queue:
request = heapq.heappop(self.request_queue)
# 检查请求是否仍然有效
if request.request_id in self.pending_requests:
# 只有在启用超时且请求设置了有限超时时才检查超时
if (self.enable_timeout and
request.timeout != float('inf') and
time.time() - request.created_at > request.timeout):
await self._mark_request_timeout(request.request_id)
continue
# 移动到处理队列
del self.pending_requests[request.request_id]
self.processing_requests[request.request_id] = request
request.status = RequestStatus.PROCESSING
self.stats['queue_size'] = len(self.request_queue)
return request
return None
async def _process_request(self, request: StorageRequest, worker_name: str):
"""处理单个请求"""
start_time = time.time()
logger.info(f"工作者 {worker_name} 开始处理请求: {request.request_id}")
try:
# 获取处理器
handler = self.handlers.get(request.handler_type)
if not handler:
raise Exception(f"未找到处理器: {request.handler_type}")
# 执行处理
result = await handler(request.input_params, request.context_data, request.map_id, request.task_record_id)
# 标记完成
await self._mark_request_completed(request.request_id, result)
processing_time = time.time() - start_time
logger.info(f"工作者 {worker_name} 完成请求: {request.request_id}, 用时: {processing_time:.3f}s")
except Exception as e:
# 处理失败
error_msg = str(e)
logger.error(f"工作者 {worker_name} 处理请求失败: {request.request_id}, 错误: {error_msg}")
# 检查是否需要重试
if request.retry_count < request.max_retries:
await self._retry_request(request)
else:
await self._mark_request_failed(request.request_id, error_msg)
finally:
# 更新平均处理时间
processing_time = time.time() - start_time
total_time = self.stats['avg_processing_time'] * self.stats['requests_completed']
self.stats['requests_completed'] += 1
self.stats['avg_processing_time'] = (total_time + processing_time) / self.stats['requests_completed']
async def _mark_request_completed(self, request_id: str, result: Dict[str, Any]):
"""标记请求完成"""
if request_id in self.processing_requests:
request = self.processing_requests[request_id]
request.status = RequestStatus.COMPLETED
request.result = result
del self.processing_requests[request_id]
self.completed_requests[request_id] = request
self.stats['requests_completed'] += 1
async def _mark_request_failed(self, request_id: str, error_message: str):
"""标记请求失败"""
if request_id in self.processing_requests:
request = self.processing_requests[request_id]
request.status = RequestStatus.FAILED
request.error_message = error_message
del self.processing_requests[request_id]
self.completed_requests[request_id] = request
self.stats['requests_failed'] += 1
async def _mark_request_timeout(self, request_id: str):
"""标记请求超时"""
request = None
if request_id in self.pending_requests:
request = self.pending_requests[request_id]
del self.pending_requests[request_id]
elif request_id in self.processing_requests:
request = self.processing_requests[request_id]
del self.processing_requests[request_id]
if request:
request.status = RequestStatus.TIMEOUT
request.error_message = "请求超时"
self.completed_requests[request_id] = request
self.stats['requests_timeout'] += 1
async def _retry_request(self, request: StorageRequest):
"""重试请求"""
request.retry_count += 1
request.status = RequestStatus.PENDING
# 重新加入队列
async with self.queue_lock:
heapq.heappush(self.request_queue, request)
del self.processing_requests[request.request_id]
self.pending_requests[request.request_id] = request
self.stats['queue_size'] = len(self.request_queue)
logger.info(f"重试请求: {request.request_id}, 第 {request.retry_count} 次重试")
def _get_queue_position(self, request_id: str) -> int:
"""获取请求在队列中的位置"""
for i, request in enumerate(self.request_queue):
if request.request_id == request_id:
return i + 1
return -1
async def _cleanup_completed_requests(self):
"""清理已完成的请求"""
while not self.shutdown_event.is_set():
try:
await asyncio.sleep(self.cleanup_interval) # 使用配置的清理间隔
current_time = time.time()
cleanup_threshold = self.completed_request_ttl # 使用配置的保留时间
to_remove = []
for request_id, request in self.completed_requests.items():
if current_time - request.created_at > cleanup_threshold:
to_remove.append(request_id)
for request_id in to_remove:
del self.completed_requests[request_id]
if to_remove:
logger.info(f"清理了 {len(to_remove)} 个已完成的请求")
except Exception as e:
logger.error(f"清理已完成请求异常: {str(e)}")
# 全局队列管理器实例
storage_queue_manager = StorageQueueManager()