258 lines
8.8 KiB
Python
258 lines
8.8 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
优先级队列管理模块
|
||
实现动态阈值多级队列和平衡调度策略
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import heapq
|
||
from typing import Dict, List, Any, Optional, Tuple, Set
|
||
from datetime import datetime, timedelta
|
||
from collections import defaultdict
|
||
from utils.logger import get_logger
|
||
|
||
# 获取日志记录器
|
||
logger = get_logger("services.enhanced_scheduler.priority_queue_manager")
|
||
|
||
class PriorityQueueManager:
|
||
"""
|
||
优先级队列管理器
|
||
实现动态阈值多级队列和平衡调度策略
|
||
"""
|
||
|
||
def __init__(self,
|
||
queue_count: int = 3,
|
||
threshold_percentiles: List[float] = None,
|
||
worker_ratios: List[float] = None,
|
||
rebalance_interval: int = 300):
|
||
"""
|
||
初始化优先级队列管理器
|
||
|
||
Args:
|
||
queue_count: 队列数量
|
||
threshold_percentiles: 队列阈值百分比，如[0.1, 0.3, 1.0]表示前10%、接下来20%、剩余的70%
|
||
worker_ratios: 工作线程分配比例，如[0.6, 0.3, 0.1]
|
||
rebalance_interval: 重新平衡间隔（秒）
|
||
"""
|
||
self.queue_count = queue_count
|
||
self.threshold_percentiles = threshold_percentiles
|
||
self.worker_ratios = worker_ratios
|
||
self.rebalance_interval = rebalance_interval
|
||
|
||
# 验证配置
|
||
if len(self.threshold_percentiles) != queue_count:
|
||
raise ValueError(f"阈值百分比数量({len(self.threshold_percentiles)})必须等于队列数量({queue_count})")
|
||
if len(self.worker_ratios) != queue_count:
|
||
raise ValueError(f"工作线程比例数量({len(self.worker_ratios)})必须等于队列数量({queue_count})")
|
||
if sum(self.worker_ratios) != 1.0:
|
||
raise ValueError(f"工作线程比例总和必须等于1.0，当前为{sum(self.worker_ratios)}")
|
||
|
||
# 初始化多级队列
|
||
self.priority_queues = [asyncio.PriorityQueue() for _ in range(queue_count)]
|
||
|
||
# 优先级阈值，初始为默认值
|
||
self.priority_thresholds = [1, 5, 10]
|
||
|
||
# 优先级映射表 {task_id: priority}
|
||
self.priority_map = {}
|
||
|
||
# 最后一次重新平衡时间
|
||
self.last_rebalance_time = datetime.now()
|
||
|
||
# 优先级统计信息
|
||
self.priority_stats = []
|
||
|
||
logger.info(f"初始化优先级队列管理器: 队列数={queue_count}, 阈值百分比={self.threshold_percentiles}, "
|
||
f"工作线程比例={self.worker_ratios}")
|
||
|
||
async def enqueue(self, task_id: str, priority: int) -> int:
|
||
"""
|
||
将任务添加到适当的优先级队列
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
priority: 任务优先级
|
||
|
||
Returns:
|
||
int: 队列索引
|
||
"""
|
||
# 检查是否需要重新平衡
|
||
await self._check_rebalance()
|
||
|
||
# 确定任务应该进入哪个队列
|
||
queue_index = self._get_queue_index(priority)
|
||
|
||
# 保存优先级信息
|
||
self.priority_map[task_id] = priority
|
||
self.priority_stats.append(priority)
|
||
|
||
# 添加到队列（使用负优先级值）
|
||
await self.priority_queues[queue_index].put((-priority, task_id))
|
||
logger.info(f"任务 {task_id} (优先级 {priority}) 添加到队列 {queue_index}")
|
||
|
||
return queue_index
|
||
|
||
async def dequeue(self, worker_id: int, worker_count: int) -> Tuple[int, Any]:
|
||
"""
|
||
从适当的队列中获取任务
|
||
|
||
Args:
|
||
worker_id: 工作线程ID
|
||
worker_count: 工作线程总数
|
||
|
||
Returns:
|
||
Tuple[int, Any]: (队列索引, 任务数据)
|
||
"""
|
||
# 确定工作线程应该从哪个队列获取任务
|
||
queue_index = self._get_worker_queue(worker_id, worker_count)
|
||
# 检查指定队列是否有任务
|
||
if not self.priority_queues[queue_index].empty():
|
||
item = await self.priority_queues[queue_index].get()
|
||
return queue_index, item
|
||
|
||
# 如果指定队列为空，尝试从其他队列获取任务
|
||
for i in range(self.queue_count):
|
||
if not self.priority_queues[i].empty():
|
||
item = await self.priority_queues[i].get()
|
||
return i, item
|
||
|
||
# 所有队列都为空，返回None
|
||
return -1, None
|
||
|
||
def task_done(self, queue_index: int) -> None:
|
||
"""
|
||
标记任务完成
|
||
|
||
Args:
|
||
queue_index: 队列索引
|
||
"""
|
||
if 0 <= queue_index < self.queue_count:
|
||
self.priority_queues[queue_index].task_done()
|
||
|
||
def remove_task(self, task_id: str) -> bool:
|
||
"""
|
||
从优先级映射表中移除任务
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
bool: 是否成功移除
|
||
"""
|
||
if task_id in self.priority_map:
|
||
self.priority_map.pop(task_id)
|
||
return True
|
||
return False
|
||
|
||
async def _check_rebalance(self) -> None:
|
||
"""
|
||
检查是否需要重新平衡队列阈值
|
||
"""
|
||
now = datetime.now()
|
||
if (now - self.last_rebalance_time).total_seconds() > self.rebalance_interval:
|
||
await self._rebalance_thresholds()
|
||
self.last_rebalance_time = now
|
||
|
||
async def _rebalance_thresholds(self) -> None:
|
||
"""
|
||
重新平衡优先级阈值
|
||
基于历史优先级分布动态调整阈值
|
||
"""
|
||
if len(self.priority_stats) < 10:
|
||
# 数据不足，使用默认值
|
||
return
|
||
|
||
# 复制并排序优先级数据
|
||
sorted_priorities = sorted(self.priority_stats, reverse=True)
|
||
|
||
# 根据百分比计算新阈值
|
||
new_thresholds = []
|
||
for percentile in self.threshold_percentiles[:-1]: # 最后一个阈值不需要计算
|
||
index = min(int(len(sorted_priorities) * percentile), len(sorted_priorities) - 1)
|
||
new_thresholds.append(sorted_priorities[index])
|
||
|
||
# 添加最小阈值（1，确保所有任务都能分配到队列）
|
||
new_thresholds.append(1)
|
||
|
||
# 更新阈值
|
||
if new_thresholds != self.priority_thresholds:
|
||
old_thresholds = self.priority_thresholds.copy()
|
||
self.priority_thresholds = new_thresholds
|
||
logger.info(f"优先级阈值更新: {old_thresholds} -> {new_thresholds}")
|
||
|
||
# 限制历史数据大小，保留最近的数据
|
||
if len(self.priority_stats) > 1000:
|
||
self.priority_stats = self.priority_stats[-1000:]
|
||
|
||
def _get_queue_index(self, priority: int) -> int:
|
||
"""
|
||
根据优先级确定队列索引
|
||
|
||
Args:
|
||
priority: 任务优先级
|
||
|
||
Returns:
|
||
int: 队列索引
|
||
"""
|
||
for i, threshold in enumerate(self.priority_thresholds):
|
||
if priority >= threshold:
|
||
return i
|
||
|
||
# 默认放入最低优先级队列
|
||
return self.queue_count - 1
|
||
|
||
def _get_worker_queue(self, worker_id: int, worker_count: int) -> int:
|
||
"""
|
||
根据工作线程ID确定应该从哪个队列获取任务
|
||
|
||
Args:
|
||
worker_id: 工作线程ID
|
||
worker_count: 工作线程总数
|
||
|
||
Returns:
|
||
int: 队列索引
|
||
"""
|
||
# 计算每个队列的工作线程数量
|
||
worker_counts = [int(ratio * worker_count) for ratio in self.worker_ratios]
|
||
|
||
# 补充剩余的工作线程到最后一个队列
|
||
remaining = worker_count - sum(worker_counts)
|
||
worker_counts[-1] += remaining
|
||
|
||
# 确定工作线程应该处理哪个队列
|
||
current_count = 0
|
||
for i, count in enumerate(worker_counts):
|
||
current_count += count
|
||
if worker_id < current_count:
|
||
return i
|
||
|
||
# 默认返回最低优先级队列
|
||
return self.queue_count - 1
|
||
|
||
def get_queue_sizes(self) -> List[int]:
|
||
"""
|
||
获取各队列大小
|
||
|
||
Returns:
|
||
List[int]: 队列大小列表
|
||
"""
|
||
return [queue.qsize() for queue in self.priority_queues]
|
||
|
||
def get_queue_status(self) -> Dict[str, Any]:
|
||
"""
|
||
获取队列状态信息
|
||
|
||
Returns:
|
||
Dict[str, Any]: 队列状态
|
||
"""
|
||
return {
|
||
"queue_count": self.queue_count,
|
||
"thresholds": self.priority_thresholds,
|
||
"queue_sizes": self.get_queue_sizes(),
|
||
"worker_ratios": self.worker_ratios,
|
||
"task_count": len(self.priority_map),
|
||
"last_rebalance": self.last_rebalance_time.strftime("%Y-%m-%d %H:%M:%S")
|
||
} |