refactor: isolate dispatch scheduling core

This commit is contained in:
fawney19
2026-05-12 13:15:41 +08:00
parent 81ee27cdea
commit fa73655134
50 changed files with 4467 additions and 3202 deletions

View File

@@ -0,0 +1,16 @@
[package]
name = "aether-dispatch-core"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Request-scoped dispatch sequence and pool cursor primitives for Aether"
[dependencies]
async-trait.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
[dev-dependencies]
tokio.workspace = true

View File

@@ -0,0 +1,68 @@
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ProviderEndpointRef {
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
pub selected_provider_model_name: String,
pub api_format: String,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct KeyRef {
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
pub model_id: String,
pub selected_provider_model_name: String,
pub api_format: String,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolRef {
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
pub selected_provider_model_name: String,
pub api_format: String,
pub pool_group_id: String,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DispatchRankFacts {
pub provider_priority: i32,
pub key_priority: Option<i32>,
pub ranking_reason: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum DispatchCandidateRef {
SingleKey {
key: KeyRef,
rank: DispatchRankFacts,
},
PoolRef {
pool: PoolRef,
rank: DispatchRankFacts,
},
}
impl DispatchCandidateRef {
pub fn provider_endpoint(&self) -> ProviderEndpointRef {
match self {
Self::SingleKey { key, .. } => ProviderEndpointRef {
provider_id: key.provider_id.clone(),
endpoint_id: key.endpoint_id.clone(),
model_id: key.model_id.clone(),
selected_provider_model_name: key.selected_provider_model_name.clone(),
api_format: key.api_format.clone(),
},
Self::PoolRef { pool, .. } => ProviderEndpointRef {
provider_id: pool.provider_id.clone(),
endpoint_id: pool.endpoint_id.clone(),
model_id: pool.model_id.clone(),
selected_provider_model_name: pool.selected_provider_model_name.clone(),
api_format: pool.api_format.clone(),
},
}
}
}

View File

@@ -0,0 +1,16 @@
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum DispatchEffectKind {
CandidateFailed,
RateLimited,
Succeeded,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DispatchEffect {
pub kind: DispatchEffectKind,
pub provider_id: String,
pub endpoint_id: String,
pub key_id: Option<String>,
pub candidate_index: u32,
pub reason: Option<String>,
}

View File

@@ -0,0 +1,15 @@
pub mod candidate;
pub mod effects;
pub mod pool;
pub mod sequence;
pub use candidate::{
DispatchCandidateRef, DispatchRankFacts, KeyRef, PoolRef, ProviderEndpointRef,
};
pub use effects::{DispatchEffect, DispatchEffectKind};
pub use pool::{
run_pool_dispatch_cursor, PoolDispatchCursorOutcome, PoolDispatchError, PoolDispatchPort,
PoolDispatchWindow, PoolWindowConfig, DEFAULT_POOL_MAX_SCAN, DEFAULT_POOL_PAGE_SIZE,
DEFAULT_POOL_WINDOW_SIZE,
};
pub use sequence::{DispatchSequence, DispatchSequenceItem, DispatchSequenceMark};

View File

@@ -0,0 +1,216 @@
use async_trait::async_trait;
pub const DEFAULT_POOL_WINDOW_SIZE: u32 = 16;
pub const DEFAULT_POOL_PAGE_SIZE: u32 = 64;
pub const DEFAULT_POOL_MAX_SCAN: u32 = 512;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolWindowConfig {
pub window_size: u32,
pub page_size: u32,
pub max_scan: u32,
}
impl Default for PoolWindowConfig {
fn default() -> Self {
Self {
window_size: DEFAULT_POOL_WINDOW_SIZE,
page_size: DEFAULT_POOL_PAGE_SIZE,
max_scan: DEFAULT_POOL_MAX_SCAN,
}
}
}
impl PoolWindowConfig {
pub fn normalized(self) -> Self {
let page_size = self.page_size.max(1);
let window_size = self.window_size.max(1).min(page_size);
let max_scan = self.max_scan.max(window_size);
Self {
window_size,
page_size,
max_scan,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PoolDispatchWindow<Candidate> {
pub candidates: Vec<Candidate>,
pub scanned_count: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PoolDispatchCursorOutcome<Candidate> {
pub candidates: Vec<Candidate>,
pub scanned_count: u32,
pub exhausted: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum PoolDispatchError<Error> {
#[error("pool dispatch port failed")]
Port(Error),
}
#[async_trait]
pub trait PoolDispatchPort {
type Candidate: Send;
type Error: Send;
async fn read_page(
&mut self,
offset: u32,
limit: u32,
) -> Result<Vec<Self::Candidate>, Self::Error>;
async fn rank_and_filter_window(
&mut self,
candidates: Vec<Self::Candidate>,
window_size: u32,
) -> Result<PoolDispatchWindow<Self::Candidate>, Self::Error>;
}
pub async fn run_pool_dispatch_cursor<Port>(
port: &mut Port,
config: PoolWindowConfig,
) -> Result<PoolDispatchCursorOutcome<Port::Candidate>, PoolDispatchError<Port::Error>>
where
Port: PoolDispatchPort + Send,
{
let config = config.normalized();
let mut offset = 0_u32;
let mut scanned_count = 0_u32;
while scanned_count < config.max_scan {
let limit = config.page_size.min(config.max_scan - scanned_count);
let page = port
.read_page(offset, limit)
.await
.map_err(PoolDispatchError::Port)?;
if page.is_empty() {
return Ok(PoolDispatchCursorOutcome {
candidates: Vec::new(),
scanned_count,
exhausted: true,
});
}
let page_len = u32::try_from(page.len()).unwrap_or(u32::MAX);
offset = offset.saturating_add(page_len);
scanned_count = scanned_count.saturating_add(page_len);
let window = port
.rank_and_filter_window(page, config.window_size)
.await
.map_err(PoolDispatchError::Port)?;
if !window.candidates.is_empty() {
return Ok(PoolDispatchCursorOutcome {
candidates: window.candidates,
scanned_count,
exhausted: false,
});
}
}
Ok(PoolDispatchCursorOutcome {
candidates: Vec::new(),
scanned_count,
exhausted: true,
})
}
#[cfg(test)]
mod tests {
use std::collections::VecDeque;
use async_trait::async_trait;
use super::{run_pool_dispatch_cursor, PoolDispatchPort, PoolDispatchWindow, PoolWindowConfig};
#[derive(Default)]
struct TestPort {
pages: VecDeque<Vec<u32>>,
read_limits: Vec<u32>,
}
#[async_trait]
impl PoolDispatchPort for TestPort {
type Candidate = u32;
type Error = std::convert::Infallible;
async fn read_page(
&mut self,
_offset: u32,
limit: u32,
) -> Result<Vec<Self::Candidate>, Self::Error> {
self.read_limits.push(limit);
Ok(self.pages.pop_front().unwrap_or_default())
}
async fn rank_and_filter_window(
&mut self,
mut candidates: Vec<Self::Candidate>,
window_size: u32,
) -> Result<PoolDispatchWindow<Self::Candidate>, Self::Error> {
candidates.sort();
candidates.truncate(window_size as usize);
let scanned_count = u32::try_from(candidates.len()).unwrap_or(u32::MAX);
Ok(PoolDispatchWindow {
candidates,
scanned_count,
})
}
}
#[tokio::test]
async fn small_pool_is_returned_in_one_frozen_window() {
let mut port = TestPort {
pages: VecDeque::from([vec![3, 1, 2]]),
read_limits: Vec::new(),
};
let outcome = run_pool_dispatch_cursor(&mut port, PoolWindowConfig::default())
.await
.unwrap();
assert_eq!(outcome.candidates, [1, 2, 3]);
assert_eq!(outcome.scanned_count, 3);
assert_eq!(port.read_limits, [64]);
}
#[tokio::test]
async fn large_pool_returns_bounded_window() {
let mut port = TestPort {
pages: VecDeque::from([(0..100).rev().collect::<Vec<_>>()]),
read_limits: Vec::new(),
};
let outcome = run_pool_dispatch_cursor(&mut port, PoolWindowConfig::default())
.await
.unwrap();
assert_eq!(outcome.candidates.len(), 16);
assert_eq!(outcome.candidates[0], 0);
assert_eq!(outcome.candidates[15], 15);
assert_eq!(port.read_limits, [64]);
}
#[tokio::test]
async fn max_scan_caps_page_reads() {
let mut port = TestPort {
pages: VecDeque::from([Vec::new()]),
read_limits: Vec::new(),
};
let config = PoolWindowConfig {
window_size: 16,
page_size: 64,
max_scan: 32,
};
let outcome = run_pool_dispatch_cursor(&mut port, config).await.unwrap();
assert!(outcome.exhausted);
assert_eq!(port.read_limits, [32]);
}
}

View File

@@ -0,0 +1,110 @@
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DispatchSequenceItem<Candidate> {
pub candidate_index: u32,
pub retry_index: u32,
pub candidate: Candidate,
pub mark: DispatchSequenceMark,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum DispatchSequenceMark {
Pending,
Failed,
Succeeded,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DispatchSequence<Candidate> {
items: Vec<DispatchSequenceItem<Candidate>>,
cursor: usize,
}
impl<Candidate> DispatchSequence<Candidate> {
pub fn new(items: Vec<DispatchSequenceItem<Candidate>>) -> Self {
Self { items, cursor: 0 }
}
pub fn from_candidates(candidates: Vec<Candidate>) -> Self {
let items = candidates
.into_iter()
.enumerate()
.map(|(index, candidate)| DispatchSequenceItem {
candidate_index: u32::try_from(index).unwrap_or(u32::MAX),
retry_index: 0,
candidate,
mark: DispatchSequenceMark::Pending,
})
.collect();
Self::new(items)
}
pub fn peek_current(&self) -> Option<&DispatchSequenceItem<Candidate>> {
self.items.get(self.cursor)
}
#[allow(clippy::should_implement_trait)]
pub fn next(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
while self
.items
.get(self.cursor)
.is_some_and(|item| item.mark != DispatchSequenceMark::Pending)
{
self.cursor = self.cursor.saturating_add(1);
}
self.items.get(self.cursor)
}
pub fn mark_failed(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
self.mark_current(DispatchSequenceMark::Failed)
}
pub fn mark_succeeded(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
self.mark_current(DispatchSequenceMark::Succeeded)
}
pub fn len(&self) -> usize {
self.items.len()
}
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
pub fn cursor(&self) -> usize {
self.cursor
}
pub fn items(&self) -> &[DispatchSequenceItem<Candidate>] {
&self.items
}
fn mark_current(
&mut self,
mark: DispatchSequenceMark,
) -> Option<&DispatchSequenceItem<Candidate>> {
let item = self.items.get_mut(self.cursor)?;
item.mark = mark;
self.cursor = self.cursor.saturating_add(1);
self.items.get(self.cursor)
}
}
#[cfg(test)]
mod tests {
use super::{DispatchSequence, DispatchSequenceMark};
#[test]
fn mark_failed_advances_without_reordering() {
let mut sequence = DispatchSequence::from_candidates(vec!["a", "b", "c"]);
assert_eq!(sequence.next().map(|item| item.candidate), Some("a"));
assert_eq!(sequence.mark_failed().map(|item| item.candidate), Some("b"));
assert_eq!(sequence.next().map(|item| item.candidate), Some("b"));
assert_eq!(sequence.mark_failed().map(|item| item.candidate), Some("c"));
assert_eq!(sequence.next().map(|item| item.candidate), Some("c"));
assert_eq!(sequence.items()[0].mark, DispatchSequenceMark::Failed);
assert_eq!(sequence.items()[1].mark, DispatchSequenceMark::Failed);
assert_eq!(sequence.items()[2].mark, DispatchSequenceMark::Pending);
}
}