mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix(scheduling): 负载均衡模式与无亲和键场景统一走随机排序
重构 CandidateSorter.shuffle_keys_by_internal_priority 中同优先级 Key 的排序逻辑,将三分支简化为两分支: - 随机排序:TTL=0 / 负载均衡模式 / 无 affinity_key - 哈希确定性排序:缓存亲和模式且有 affinity_key 移除了原先"无 affinity_key 时按 ID 排序"的冗余分支,新增 对应单元测试覆盖三种场景。
This commit is contained in:
@@ -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)
|
||||||
|
|||||||
93
tests/unit/test_candidate_sorter_key_shuffle.py
Normal file
93
tests/unit/test_candidate_sorter_key_shuffle.py
Normal 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()
|
||||||
Reference in New Issue
Block a user