fix(scheduling): 负载均衡模式与无亲和键场景统一走随机排序

重构 CandidateSorter.shuffle_keys_by_internal_priority 中同优先级
Key 的排序逻辑,将三分支简化为两分支:
- 随机排序:TTL=0 / 负载均衡模式 / 无 affinity_key
- 哈希确定性排序:缓存亲和模式且有 affinity_key

移除了原先"无 affinity_key 时按 ID 排序"的冗余分支,新增
对应单元测试覆盖三种场景。
This commit is contained in:
fawney19
2026-03-02 22:36:24 +08:00
parent 8e98eed5c8
commit e26ed8481f
2 changed files with 104 additions and 9 deletions

View File

@@ -299,24 +299,26 @@ class CandidateSorter:
group_keys = priority_groups[priority] group_keys = priority_groups[priority]
if len(group_keys) > 1: if len(group_keys) > 1:
if use_random: should_randomize = (
# TTL=0 模式:使用随机排序实现 Key 轮换 use_random
or self._config.scheduling_mode == SchedulingConfig.SCHEDULING_MODE_LOAD_BALANCE
or not affinity_key
)
if should_randomize:
# 随机排序TTL=0 / 负载均衡模式 / 无 affinity_key
shuffled = list(group_keys) shuffled = list(group_keys)
random.shuffle(shuffled) random.shuffle(shuffled)
result.extend(shuffled) result.extend(shuffled)
elif affinity_key: else:
# 正常模式:使用哈希确定性打乱(保持缓存亲和性 # 缓存亲和模式:使用哈希确定性排序should_randomize=False 蕴含 affinity_key 非空
key_scores = [] key_scores = []
for key in group_keys: for key in group_keys:
hash_value = affinity_hash(affinity_key, key.id) hash_value = affinity_hash(affinity_key, key.id) # type: ignore[arg-type]
key_scores.append((hash_value, key)) key_scores.append((hash_value, key))
# 按哈希值排序
sorted_group = [key for _, key in sorted(key_scores, key=lambda x: x[0])] sorted_group = [key for _, key in sorted(key_scores, key=lambda x: x[0])]
result.extend(sorted_group) result.extend(sorted_group)
else:
# 没有 affinity_key 时按 ID 排序保持稳定性
result.extend(sorted(group_keys, key=lambda k: k.id))
else: else:
# 单个 Key 直接添加 # 单个 Key 直接添加
result.extend(group_keys) result.extend(group_keys)

View File

@@ -0,0 +1,93 @@
from __future__ import annotations
from types import SimpleNamespace
from typing import cast
from unittest.mock import patch
from src.models.database import ProviderAPIKey
from src.services.scheduling.candidate_sorter import CandidateSorter
from src.services.scheduling.scheduling_config import SchedulingConfig
from src.services.scheduling.utils import affinity_hash
def _make_key(key_id: str, internal_priority: int = 1) -> ProviderAPIKey:
return cast(
ProviderAPIKey,
SimpleNamespace(
id=key_id,
internal_priority=internal_priority,
),
)
def _reverse_in_place(items: list[object]) -> None:
items.reverse()
def test_shuffle_keys_random_in_load_balance_mode_even_with_affinity_key() -> None:
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_LOAD_BALANCE,
)
sorter = CandidateSorter(config)
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
with patch(
"src.services.scheduling.candidate_sorter.random.shuffle",
side_effect=_reverse_in_place,
) as shuffle_mock:
result = sorter.shuffle_keys_by_internal_priority(
keys,
affinity_key="affinity-1",
use_random=False,
)
assert [k.id for k in result] == ["k3", "k2", "k1"]
assert shuffle_mock.call_count == 1
def test_shuffle_keys_random_when_affinity_key_absent() -> None:
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_CACHE_AFFINITY,
)
sorter = CandidateSorter(config)
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
with patch(
"src.services.scheduling.candidate_sorter.random.shuffle",
side_effect=_reverse_in_place,
) as shuffle_mock:
result = sorter.shuffle_keys_by_internal_priority(
keys,
affinity_key=None,
use_random=False,
)
assert [k.id for k in result] == ["k3", "k2", "k1"]
assert shuffle_mock.call_count == 1
def test_shuffle_keys_still_hashes_with_affinity_in_non_load_balance_mode() -> None:
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_CACHE_AFFINITY,
)
sorter = CandidateSorter(config)
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
affinity_key = "affinity-1"
expected = sorted(keys, key=lambda k: affinity_hash(affinity_key, k.id))
with patch("src.services.scheduling.candidate_sorter.random.shuffle") as shuffle_mock:
result = sorter.shuffle_keys_by_internal_priority(
keys,
affinity_key=affinity_key,
use_random=False,
)
assert [k.id for k in result] == [k.id for k in expected]
shuffle_mock.assert_not_called()