Merge remote-tracking branch 'origin/codex/pool-member-scores' into aether-rust-pioneer

This commit is contained in:
fawney19
2026-05-12 09:17:25 +08:00
81 changed files with 7206 additions and 77 deletions

View File

@@ -9,6 +9,7 @@ description = "AI serving application contracts and ports for Aether"
[dependencies]
aether-ai-formats.workspace = true
aether-contracts.workspace = true
aether-data-contracts.workspace = true
aether-scheduler-core.workspace = true
async-trait.workspace = true
http.workspace = true

View File

@@ -16,6 +16,7 @@ pub mod execution_path;
pub mod failure_diagnostic;
pub mod plan_payload;
pub mod pool_scheduler;
pub mod pool_scores;
pub mod ports;
pub mod ranking_metadata;
pub mod report_context;
@@ -99,6 +100,13 @@ pub use pool_scheduler::{
AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, AI_POOL_COOLDOWN_SKIP_REASON,
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
pub use pool_scores::{
probe_freshness_score, probe_freshness_score_with_ttl, score_pool_member,
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreOutput,
PoolMemberScoreRules, PoolMemberScoreWeights, POOL_SCORE_VERSION,
PROBE_FAILURE_COOLDOWN_THRESHOLD, PROBE_FAILURE_PENALTY, PROBE_FRESHNESS_TTL_SECONDS,
REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP,
};
pub use ranking_metadata::append_ai_ranking_metadata_to_object;
pub use report_context::{
build_ai_execution_report_context, build_ai_report_context_original_request_echo,

View File

@@ -0,0 +1,540 @@
use aether_data_contracts::repository::pool_scores::{
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope,
};
use serde_json::{json, Value};
pub const POOL_SCORE_VERSION: u64 = 1;
pub const PROBE_FRESHNESS_TTL_SECONDS: u64 = 30 * 60;
pub const UNSCHEDULABLE_SCORE_CAP: f64 = 0.05;
pub const PROBE_FAILURE_PENALTY: f64 = 0.05;
pub const REQUEST_FAILURE_PENALTY: f64 = 0.005;
pub const PROBE_FAILURE_COOLDOWN_THRESHOLD: u64 = 3;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PoolMemberScoreWeights {
pub manual_priority: f64,
pub health: f64,
pub probe_freshness: f64,
pub quota_remaining: f64,
pub latency: f64,
pub cost_lru: f64,
}
impl Default for PoolMemberScoreWeights {
fn default() -> Self {
Self {
manual_priority: 0.30,
health: 0.20,
probe_freshness: 0.15,
quota_remaining: 0.15,
latency: 0.10,
cost_lru: 0.10,
}
}
}
impl PoolMemberScoreWeights {
pub fn normalized(self) -> Self {
let sanitized = Self {
manual_priority: finite_non_negative(self.manual_priority),
health: finite_non_negative(self.health),
probe_freshness: finite_non_negative(self.probe_freshness),
quota_remaining: finite_non_negative(self.quota_remaining),
latency: finite_non_negative(self.latency),
cost_lru: finite_non_negative(self.cost_lru),
};
let total = sanitized.manual_priority
+ sanitized.health
+ sanitized.probe_freshness
+ sanitized.quota_remaining
+ sanitized.latency
+ sanitized.cost_lru;
if total <= f64::EPSILON {
return sanitized;
}
Self {
manual_priority: sanitized.manual_priority / total,
health: sanitized.health / total,
probe_freshness: sanitized.probe_freshness / total,
quota_remaining: sanitized.quota_remaining / total,
latency: sanitized.latency / total,
cost_lru: sanitized.cost_lru / total,
}
}
fn as_reason_json(self) -> Value {
json!({
"manual_priority": self.manual_priority,
"health": self.health,
"probe_freshness": self.probe_freshness,
"quota_remaining": self.quota_remaining,
"latency": self.latency,
"cost_lru": self.cost_lru
})
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PoolMemberScoreRules {
pub weights: PoolMemberScoreWeights,
pub probe_freshness_ttl_seconds: u64,
pub unschedulable_score_cap: f64,
pub probe_failure_penalty: f64,
pub request_failure_penalty: f64,
pub probe_failure_cooldown_threshold: u64,
}
impl Default for PoolMemberScoreRules {
fn default() -> Self {
Self {
weights: PoolMemberScoreWeights::default(),
probe_freshness_ttl_seconds: PROBE_FRESHNESS_TTL_SECONDS,
unschedulable_score_cap: UNSCHEDULABLE_SCORE_CAP,
probe_failure_penalty: PROBE_FAILURE_PENALTY,
request_failure_penalty: REQUEST_FAILURE_PENALTY,
probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
}
}
}
impl PoolMemberScoreRules {
pub fn effective(self) -> Self {
let defaults = Self::default();
Self {
weights: self.weights.normalized(),
probe_freshness_ttl_seconds: if self.probe_freshness_ttl_seconds == 0 {
defaults.probe_freshness_ttl_seconds
} else {
self.probe_freshness_ttl_seconds
},
unschedulable_score_cap: if self.unschedulable_score_cap.is_finite() {
self.unschedulable_score_cap.clamp(0.0, 1.0)
} else {
defaults.unschedulable_score_cap
},
probe_failure_penalty: if self.probe_failure_penalty.is_finite() {
self.probe_failure_penalty.clamp(0.0, 1.0)
} else {
defaults.probe_failure_penalty
},
request_failure_penalty: if self.request_failure_penalty.is_finite() {
self.request_failure_penalty.clamp(0.0, 1.0)
} else {
defaults.request_failure_penalty
},
probe_failure_cooldown_threshold: self.probe_failure_cooldown_threshold,
}
}
}
#[derive(Debug, Clone)]
pub struct PoolMemberScoreInput {
pub identity: PoolMemberIdentity,
pub scope: PoolScoreScope,
pub internal_priority: i32,
pub is_active: bool,
pub health_score: Option<f64>,
pub quota_usage_ratio: Option<f64>,
pub quota_exhausted: bool,
pub account_blocked: bool,
pub oauth_invalid_reason: Option<String>,
pub circuit_open: bool,
pub success_count: u64,
pub error_count: u64,
pub total_response_time_ms: u64,
pub total_tokens: u64,
pub total_cost_usd: f64,
pub last_used_at: Option<u64>,
pub last_probe_success_at: Option<u64>,
pub probe_failure_count: u64,
pub probe_status: PoolMemberProbeStatus,
pub now_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PoolMemberScoreOutput {
pub score: f64,
pub hard_state: PoolMemberHardState,
pub score_reason: Value,
}
pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput {
score_pool_member_with_rules(input, PoolMemberScoreRules::default())
}
pub fn score_pool_member_with_rules(
input: &PoolMemberScoreInput,
rules: PoolMemberScoreRules,
) -> PoolMemberScoreOutput {
let rules = rules.effective();
let weights = rules.weights;
let hard_state = derive_hard_state(input, &rules);
let manual_priority = manual_priority_score(input.internal_priority);
let health = input.health_score.unwrap_or(0.5).clamp(0.0, 1.0);
let probe_freshness = probe_freshness_score_with_ttl(
input.last_probe_success_at,
input.probe_status,
input.now_unix_secs,
rules.probe_freshness_ttl_seconds,
);
let quota_remaining = input
.quota_usage_ratio
.map(|ratio| 1.0 - ratio.clamp(0.0, 1.0))
.unwrap_or(0.5);
let latency = latency_score(input.success_count, input.total_response_time_ms);
let cost_lru = cost_lru_score(input.total_cost_usd, input.total_tokens, input.last_used_at);
let weighted_score = manual_priority * weights.manual_priority
+ health * weights.health
+ probe_freshness * weights.probe_freshness
+ quota_remaining * weights.quota_remaining
+ latency * weights.latency
+ cost_lru * weights.cost_lru;
let probe_failure_penalty =
(input.probe_failure_count.min(10) as f64 * rules.probe_failure_penalty).min(0.5);
let request_failure_penalty =
(input.error_count.min(20) as f64 * rules.request_failure_penalty).min(0.5);
let total_penalty = (probe_failure_penalty + request_failure_penalty).min(1.0);
let mut score = weighted_score - total_penalty;
if !hard_state.schedulable() {
score = score.min(rules.unschedulable_score_cap);
}
score = score.clamp(0.0, 1.0);
PoolMemberScoreOutput {
score,
hard_state,
score_reason: json!({
"weights": weights.as_reason_json(),
"factors": {
"manual_priority": manual_priority,
"health": health,
"probe_freshness": probe_freshness,
"quota_remaining": quota_remaining,
"latency": latency,
"cost_lru": cost_lru
},
"rules": {
"probe_freshness_ttl_seconds": rules.probe_freshness_ttl_seconds,
"unschedulable_score_cap": rules.unschedulable_score_cap,
"probe_failure_penalty": rules.probe_failure_penalty,
"request_failure_penalty": rules.request_failure_penalty,
"probe_failure_cooldown_threshold": rules.probe_failure_cooldown_threshold
},
"penalties": {
"probe_failure": probe_failure_penalty,
"request_failure": request_failure_penalty,
"total": total_penalty
},
"hard_state": hard_state.as_database(),
"score_version": POOL_SCORE_VERSION
}),
}
}
fn finite_non_negative(value: f64) -> f64 {
if value.is_finite() {
value.max(0.0)
} else {
0.0
}
}
fn derive_hard_state(
input: &PoolMemberScoreInput,
rules: &PoolMemberScoreRules,
) -> PoolMemberHardState {
if !input.is_active {
return PoolMemberHardState::Inactive;
}
if let Some(reason) = input.oauth_invalid_reason.as_deref() {
let reason = reason.to_ascii_lowercase();
if reason.contains("ban") || reason.contains("blocked") || reason.contains("suspended") {
return PoolMemberHardState::Banned;
}
return PoolMemberHardState::AuthInvalid;
}
if input.account_blocked {
return PoolMemberHardState::Banned;
}
if input.quota_exhausted {
return PoolMemberHardState::QuotaExhausted;
}
if input.circuit_open {
return PoolMemberHardState::Cooldown;
}
if input.probe_status == PoolMemberProbeStatus::Failed
&& rules.probe_failure_cooldown_threshold > 0
&& input.probe_failure_count >= rules.probe_failure_cooldown_threshold
{
return PoolMemberHardState::Cooldown;
}
if input.health_score.is_some() || input.probe_status == PoolMemberProbeStatus::Ok {
PoolMemberHardState::Available
} else {
PoolMemberHardState::Unknown
}
}
fn manual_priority_score(internal_priority: i32) -> f64 {
(1.0 - (f64::from(internal_priority).clamp(0.0, 100.0) / 100.0)).clamp(0.0, 1.0)
}
pub fn probe_freshness_score(
last_probe_success_at: Option<u64>,
probe_status: PoolMemberProbeStatus,
now_unix_secs: u64,
) -> f64 {
probe_freshness_score_with_ttl(
last_probe_success_at,
probe_status,
now_unix_secs,
PROBE_FRESHNESS_TTL_SECONDS,
)
}
pub fn probe_freshness_score_with_ttl(
last_probe_success_at: Option<u64>,
probe_status: PoolMemberProbeStatus,
now_unix_secs: u64,
ttl_seconds: u64,
) -> f64 {
if probe_status != PoolMemberProbeStatus::Ok {
return 0.0;
}
let Some(success_at) = last_probe_success_at else {
return 0.0;
};
let ttl_seconds = ttl_seconds.max(1);
let age = now_unix_secs.saturating_sub(success_at);
if age >= ttl_seconds {
0.0
} else {
1.0 - (age as f64 / ttl_seconds as f64)
}
}
fn latency_score(success_count: u64, total_response_time_ms: u64) -> f64 {
if success_count == 0 || total_response_time_ms == 0 {
return 0.5;
}
let avg = total_response_time_ms as f64 / success_count as f64;
if avg <= 500.0 {
1.0
} else if avg >= 60_000.0 {
0.0
} else {
1.0 - ((avg - 500.0) / 59_500.0)
}
}
fn cost_lru_score(total_cost_usd: f64, total_tokens: u64, last_used_at: Option<u64>) -> f64 {
let cost_penalty = if total_cost_usd.is_finite() {
(total_cost_usd.max(0.0) / 100.0).min(0.5)
} else {
0.0
};
let token_penalty = (total_tokens as f64 / 10_000_000.0).min(0.25);
let lru_bonus = if last_used_at.unwrap_or(0) == 0 {
0.25
} else {
0.0
};
(0.75 - cost_penalty - token_penalty + lru_bonus).clamp(0.0, 1.0)
}
#[cfg(test)]
mod tests {
use super::*;
use aether_data_contracts::repository::pool_scores::{
POOL_KIND_PROVIDER_KEY_POOL, POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_SCOPE_KIND_MODEL,
};
fn input() -> PoolMemberScoreInput {
PoolMemberScoreInput {
identity: PoolMemberIdentity {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: "key-1".to_string(),
},
scope: PoolScoreScope {
capability: "openai:responses".to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
},
internal_priority: 10,
is_active: true,
health_score: Some(1.0),
quota_usage_ratio: Some(0.1),
quota_exhausted: false,
account_blocked: false,
oauth_invalid_reason: None,
circuit_open: false,
success_count: 10,
error_count: 0,
total_response_time_ms: 2_000,
total_tokens: 10,
total_cost_usd: 0.01,
last_used_at: None,
last_probe_success_at: Some(1_000),
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Ok,
now_unix_secs: 1_000,
}
}
#[test]
fn hard_state_caps_unavailable_member_score() {
let mut input = input();
input.oauth_invalid_reason = Some("token invalid".to_string());
let output = score_pool_member(&input);
assert_eq!(output.hard_state, PoolMemberHardState::AuthInvalid);
assert!(output.score <= 0.05);
}
#[test]
fn probe_freshness_has_ttl() {
assert_eq!(
probe_freshness_score(Some(1_000), PoolMemberProbeStatus::Ok, 1_000),
1.0
);
assert_eq!(
probe_freshness_score(
Some(1_000),
PoolMemberProbeStatus::Ok,
1_000 + PROBE_FRESHNESS_TTL_SECONDS
),
0.0
);
}
#[test]
fn custom_rules_change_weights_and_probe_ttl() {
let mut input = input();
input.last_probe_success_at = Some(1_000);
input.now_unix_secs = 1_900;
let rules = PoolMemberScoreRules {
weights: PoolMemberScoreWeights {
manual_priority: 0.0,
health: 0.0,
probe_freshness: 1.0,
quota_remaining: 0.0,
latency: 0.0,
cost_lru: 0.0,
},
probe_freshness_ttl_seconds: 1_000,
unschedulable_score_cap: 0.05,
probe_failure_penalty: 0.0,
request_failure_penalty: 0.0,
probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
};
let output = score_pool_member_with_rules(&input, rules);
assert!((output.score - 0.1).abs() < 0.000_001);
assert_eq!(
output.score_reason["rules"]["probe_freshness_ttl_seconds"],
1_000
);
assert_eq!(output.score_reason["weights"]["probe_freshness"], 1.0);
}
#[test]
fn custom_rules_normalize_weights() {
let rules = PoolMemberScoreRules {
weights: PoolMemberScoreWeights {
manual_priority: 2.0,
health: 2.0,
probe_freshness: 0.0,
quota_remaining: 0.0,
latency: -1.0,
cost_lru: f64::NAN,
},
probe_freshness_ttl_seconds: 0,
unschedulable_score_cap: f64::INFINITY,
probe_failure_penalty: f64::NAN,
request_failure_penalty: f64::INFINITY,
probe_failure_cooldown_threshold: 2,
}
.effective();
assert_eq!(rules.weights.manual_priority, 0.5);
assert_eq!(rules.weights.health, 0.5);
assert_eq!(
rules.probe_freshness_ttl_seconds,
PROBE_FRESHNESS_TTL_SECONDS
);
assert_eq!(rules.unschedulable_score_cap, UNSCHEDULABLE_SCORE_CAP);
assert_eq!(rules.probe_failure_penalty, PROBE_FAILURE_PENALTY);
assert_eq!(rules.request_failure_penalty, REQUEST_FAILURE_PENALTY);
assert_eq!(rules.probe_failure_cooldown_threshold, 2);
}
#[test]
fn custom_rules_preserve_zero_weight_total() {
let rules = PoolMemberScoreRules {
weights: PoolMemberScoreWeights {
manual_priority: 0.0,
health: 0.0,
probe_freshness: 0.0,
quota_remaining: 0.0,
latency: 0.0,
cost_lru: 0.0,
},
..PoolMemberScoreRules::default()
}
.effective();
assert_eq!(
rules.weights,
PoolMemberScoreWeights {
manual_priority: 0.0,
health: 0.0,
probe_freshness: 0.0,
quota_remaining: 0.0,
latency: 0.0,
cost_lru: 0.0,
}
);
}
#[test]
fn repeated_probe_failures_move_member_to_cooldown() {
let mut input = input();
input.probe_status = PoolMemberProbeStatus::Failed;
input.probe_failure_count = 3;
let output = score_pool_member(&input);
assert_eq!(output.hard_state, PoolMemberHardState::Cooldown);
assert!(output.score <= UNSCHEDULABLE_SCORE_CAP);
assert_eq!(
output.score_reason["rules"]["probe_failure_cooldown_threshold"],
PROBE_FAILURE_COOLDOWN_THRESHOLD
);
}
#[test]
fn probe_and_request_failures_penalize_schedulable_score() {
let mut input = input();
input.last_probe_success_at = None;
input.probe_status = PoolMemberProbeStatus::Never;
input.probe_failure_count = 1;
input.error_count = 2;
let rules = PoolMemberScoreRules {
probe_failure_penalty: 0.1,
request_failure_penalty: 0.01,
probe_failure_cooldown_threshold: 3,
..PoolMemberScoreRules::default()
};
let output = score_pool_member_with_rules(&input, rules);
assert_eq!(output.hard_state, PoolMemberHardState::Available);
assert_eq!(output.score_reason["penalties"]["probe_failure"], 0.1);
assert_eq!(output.score_reason["penalties"]["request_failure"], 0.02);
}
}

View File

@@ -3,6 +3,6 @@ mod types;
pub use types::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
};

View File

@@ -67,6 +67,16 @@ pub struct StoredPoolKeyCandidateRowsQuery {
pub limit: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredPoolKeyCandidateRowsByKeyIdsQuery {
pub api_format: String,
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
pub selected_provider_model_name: String,
pub key_ids: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredRequestedModelCandidateRowsQuery {
pub api_format: String,
@@ -124,6 +134,11 @@ pub trait MinimalCandidateSelectionReadRepository: Send + Sync {
&self,
query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
}
pub trait MinimalCandidateSelectionRepository:

View File

@@ -3,6 +3,7 @@ pub mod billing;
pub mod candidate_selection;
pub mod candidates;
pub mod global_models;
pub mod pool_scores;
pub mod provider_catalog;
pub mod quota;
pub mod settlement;

View File

@@ -0,0 +1,11 @@
mod types;
pub use types::{
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_CAPABILITY_API_FORMAT,
POOL_SCORE_SCOPE_KIND_MODEL,
};

View File

@@ -0,0 +1,359 @@
use async_trait::async_trait;
pub const POOL_KIND_PROVIDER_KEY_POOL: &str = "provider_key_pool";
pub const POOL_MEMBER_KIND_PROVIDER_API_KEY: &str = "provider_api_key";
pub const POOL_SCORE_CAPABILITY_API_FORMAT: &str = "api_format";
pub const POOL_SCORE_SCOPE_KIND_MODEL: &str = "model";
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(rename_all = "snake_case")]
pub enum PoolMemberHardState {
Available,
Unknown,
Cooldown,
QuotaExhausted,
AuthInvalid,
Banned,
Inactive,
}
impl PoolMemberHardState {
pub fn as_database(self) -> &'static str {
match self {
Self::Available => "available",
Self::Unknown => "unknown",
Self::Cooldown => "cooldown",
Self::QuotaExhausted => "quota_exhausted",
Self::AuthInvalid => "auth_invalid",
Self::Banned => "banned",
Self::Inactive => "inactive",
}
}
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"available" => Ok(Self::Available),
"unknown" => Ok(Self::Unknown),
"cooldown" => Ok(Self::Cooldown),
"quota_exhausted" => Ok(Self::QuotaExhausted),
"auth_invalid" => Ok(Self::AuthInvalid),
"banned" => Ok(Self::Banned),
"inactive" => Ok(Self::Inactive),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unknown pool member hard_state: {other}"
))),
}
}
pub fn schedulable(self) -> bool {
matches!(self, Self::Available | Self::Unknown)
}
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(rename_all = "snake_case")]
pub enum PoolMemberProbeStatus {
Never,
Ok,
Failed,
Stale,
InProgress,
}
impl PoolMemberProbeStatus {
pub fn as_database(self) -> &'static str {
match self {
Self::Never => "never",
Self::Ok => "ok",
Self::Failed => "failed",
Self::Stale => "stale",
Self::InProgress => "in_progress",
}
}
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"never" => Ok(Self::Never),
"ok" => Ok(Self::Ok),
"failed" => Ok(Self::Failed),
"stale" => Ok(Self::Stale),
"in_progress" => Ok(Self::InProgress),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unknown pool member probe_status: {other}"
))),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolScoreScope {
pub capability: String,
pub scope_kind: String,
pub scope_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberIdentity {
pub pool_kind: String,
pub pool_id: String,
pub member_kind: String,
pub member_id: String,
}
impl PoolMemberIdentity {
pub fn provider_api_key(provider_id: impl Into<String>, key_id: impl Into<String>) -> Self {
Self {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: provider_id.into(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: key_id.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredPoolMemberScore {
pub id: String,
pub pool_kind: String,
pub pool_id: String,
pub member_kind: String,
pub member_id: String,
pub capability: String,
pub scope_kind: String,
pub scope_id: Option<String>,
pub score: f64,
pub hard_state: PoolMemberHardState,
pub score_version: u64,
pub score_reason: serde_json::Value,
pub last_ranked_at: Option<u64>,
pub last_scheduled_at: Option<u64>,
pub last_success_at: Option<u64>,
pub last_failure_at: Option<u64>,
pub failure_count: u64,
pub last_probe_attempt_at: Option<u64>,
pub last_probe_success_at: Option<u64>,
pub last_probe_failure_at: Option<u64>,
pub probe_failure_count: u64,
pub probe_status: PoolMemberProbeStatus,
pub updated_at: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpsertPoolMemberScore {
pub id: String,
pub identity: PoolMemberIdentity,
pub scope: PoolScoreScope,
pub score: f64,
pub hard_state: PoolMemberHardState,
pub score_version: u64,
pub score_reason: serde_json::Value,
pub last_ranked_at: Option<u64>,
pub last_scheduled_at: Option<u64>,
pub last_success_at: Option<u64>,
pub last_failure_at: Option<u64>,
pub failure_count: u64,
pub last_probe_attempt_at: Option<u64>,
pub last_probe_success_at: Option<u64>,
pub last_probe_failure_at: Option<u64>,
pub probe_failure_count: u64,
pub probe_status: PoolMemberProbeStatus,
pub updated_at: u64,
}
impl UpsertPoolMemberScore {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
validate_non_empty(&self.id, "pool_member_scores.id")?;
validate_non_empty(&self.identity.pool_kind, "pool_member_scores.pool_kind")?;
validate_non_empty(&self.identity.pool_id, "pool_member_scores.pool_id")?;
validate_non_empty(&self.identity.member_kind, "pool_member_scores.member_kind")?;
validate_non_empty(&self.identity.member_id, "pool_member_scores.member_id")?;
validate_non_empty(&self.scope.capability, "pool_member_scores.capability")?;
validate_non_empty(&self.scope.scope_kind, "pool_member_scores.scope_kind")?;
if !self.score.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"pool_member_scores.score must be finite".to_string(),
));
}
Ok(())
}
pub fn into_stored(self) -> StoredPoolMemberScore {
StoredPoolMemberScore {
id: self.id,
pool_kind: self.identity.pool_kind,
pool_id: self.identity.pool_id,
member_kind: self.identity.member_kind,
member_id: self.identity.member_id,
capability: self.scope.capability,
scope_kind: self.scope.scope_kind,
scope_id: self.scope.scope_id,
score: self.score,
hard_state: self.hard_state,
score_version: self.score_version,
score_reason: self.score_reason,
last_ranked_at: self.last_ranked_at,
last_scheduled_at: self.last_scheduled_at,
last_success_at: self.last_success_at,
last_failure_at: self.last_failure_at,
failure_count: self.failure_count,
last_probe_attempt_at: self.last_probe_attempt_at,
last_probe_success_at: self.last_probe_success_at,
last_probe_failure_at: self.last_probe_failure_at,
probe_failure_count: self.probe_failure_count,
probe_status: self.probe_status,
updated_at: self.updated_at,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ListRankedPoolMembersQuery {
pub pool_kind: String,
pub pool_id: String,
pub capability: String,
pub scope_kind: String,
pub scope_id: Option<String>,
pub hard_states: Vec<PoolMemberHardState>,
pub probe_statuses: Option<Vec<PoolMemberProbeStatus>>,
pub offset: usize,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ListPoolMemberScoresQuery {
pub pool_kind: String,
pub pool_id: String,
pub capability: Option<String>,
pub scope_kind: Option<String>,
pub scope_id: Option<String>,
pub hard_states: Vec<PoolMemberHardState>,
pub probe_statuses: Option<Vec<PoolMemberProbeStatus>>,
pub offset: usize,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ListPoolMemberProbeCandidatesQuery {
pub pool_kind: String,
pub pool_id: String,
pub capability: Option<String>,
pub stale_before_unix_secs: u64,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct GetPoolMemberScoresByIdsQuery {
pub ids: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberProbeResult {
pub identity: PoolMemberIdentity,
pub scope: Option<PoolScoreScope>,
pub attempted_at: u64,
pub succeeded: bool,
pub hard_state: Option<PoolMemberHardState>,
pub probe_status: PoolMemberProbeStatus,
pub score_reason_patch: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberProbeAttempt {
pub identity: PoolMemberIdentity,
pub scope: Option<PoolScoreScope>,
pub attempted_at: u64,
pub score_reason_patch: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberScheduleFeedback {
pub identity: PoolMemberIdentity,
pub scope: Option<PoolScoreScope>,
pub scheduled_at: u64,
pub succeeded: Option<bool>,
pub hard_state: Option<PoolMemberHardState>,
pub score_delta: Option<i32>,
pub score_reason_patch: Option<serde_json::Value>,
}
#[async_trait]
pub trait PoolScoreReadRepository: Send + Sync {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
}
#[async_trait]
pub trait PoolMemberScoreWriteRepository: Send + Sync {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, crate::DataLayerError>;
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, crate::DataLayerError>;
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, crate::DataLayerError>;
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, crate::DataLayerError>;
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, crate::DataLayerError>;
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, crate::DataLayerError>;
}
pub trait PoolMemberScoreRepository:
PoolScoreReadRepository + PoolMemberScoreWriteRepository + Send + Sync
{
}
impl<T> PoolMemberScoreRepository for T where
T: PoolScoreReadRepository + PoolMemberScoreWriteRepository + Send + Sync
{
}
fn validate_non_empty(value: &str, field: &str) -> Result<(), crate::DataLayerError> {
if value.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(format!(
"{field} is empty"
)));
}
Ok(())
}

View File

@@ -294,6 +294,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (provider_id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id VARCHAR(64) PRIMARY KEY,
file_name VARCHAR(512) NOT NULL,

View File

@@ -0,0 +1,29 @@
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);

View File

@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--
@@ -1659,6 +1691,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
@@ -3401,6 +3448,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

View File

@@ -0,0 +1,40 @@
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) PRIMARY KEY,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision NOT NULL DEFAULT 0,
hard_state character varying(64) NOT NULL DEFAULT 'unknown',
score_version bigint NOT NULL DEFAULT 1,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint NOT NULL DEFAULT 0,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint NOT NULL DEFAULT 0,
probe_status character varying(64) NOT NULL DEFAULT 'never',
updated_at bigint NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
ON public.pool_member_scores USING btree (updated_at);

View File

@@ -301,6 +301,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id TEXT PRIMARY KEY,
file_name TEXT NOT NULL UNIQUE,

View File

@@ -0,0 +1,37 @@
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
ON pool_member_scores (updated_at);

View File

@@ -537,6 +537,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--

View File

@@ -252,6 +252,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--

View File

@@ -877,6 +877,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

View File

@@ -127,6 +127,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (provider_id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id VARCHAR(64) PRIMARY KEY,
file_name VARCHAR(512) NOT NULL,

View File

@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--

View File

@@ -252,6 +252,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--

View File

@@ -869,6 +869,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

View File

@@ -134,6 +134,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id TEXT PRIMARY KEY,
file_name TEXT NOT NULL UNIQUE,

View File

@@ -130,6 +130,37 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (`provider_id`)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
`id` VARCHAR(192) NOT NULL,
`pool_kind` VARCHAR(64) NOT NULL,
`pool_id` VARCHAR(64) NOT NULL,
`member_kind` VARCHAR(64) NOT NULL,
`member_id` VARCHAR(64) NOT NULL,
`capability` VARCHAR(64) NOT NULL,
`scope_kind` VARCHAR(64) NOT NULL,
`scope_id` VARCHAR(128),
`score` DOUBLE NOT NULL DEFAULT 0,
`hard_state` VARCHAR(64) NOT NULL DEFAULT 'unknown',
`score_version` BIGINT NOT NULL DEFAULT 1,
`score_reason` JSON NOT NULL,
`last_ranked_at` BIGINT,
`last_scheduled_at` BIGINT,
`last_success_at` BIGINT,
`last_failure_at` BIGINT,
`failure_count` BIGINT NOT NULL DEFAULT 0,
`last_probe_attempt_at` BIGINT,
`last_probe_success_at` BIGINT,
`last_probe_failure_at` BIGINT,
`probe_failure_count` BIGINT NOT NULL DEFAULT 0,
`probe_status` VARCHAR(64) NOT NULL DEFAULT 'never',
`updated_at` BIGINT NOT NULL,
PRIMARY KEY (`id`),
KEY pool_member_scores_rank_idx (`pool_kind`, `pool_id`, `capability`, `scope_kind`, `scope_id`, `hard_state`, `score`),
KEY pool_member_scores_member_idx (`pool_kind`, `pool_id`, `member_kind`, `member_id`),
KEY pool_member_scores_probe_idx (`pool_kind`, `pool_id`, `probe_status`, `last_probe_success_at`),
KEY pool_member_scores_updated_at_idx (`updated_at`)
);
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
`id` VARCHAR(64) NOT NULL,
`api_key_id` VARCHAR(64) NOT NULL,

View File

@@ -134,6 +134,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
ALTER TABLE ONLY public.provider_api_keys ADD CONSTRAINT provider_api_keys_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON public.provider_api_keys USING btree (provider_id);
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown' NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never' NOT NULL,
updated_at bigint NOT NULL
);
ALTER TABLE ONLY public.pool_member_scores ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
CREATE TABLE IF NOT EXISTS public.api_key_provider_mappings (
id character varying(64) NOT NULL,
api_key_id character varying(64) NOT NULL,

View File

@@ -126,6 +126,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY NOT NULL,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
id TEXT PRIMARY KEY NOT NULL,
api_key_id TEXT NOT NULL,

View File

@@ -579,6 +579,143 @@ type = "unix_seconds"
name = "provider_api_keys_provider_id_idx"
columns = ["provider_id"]
[table.pool_member_scores]
domain = "provider_catalog"
order = 44
primary_key = ["id"]
[[table.pool_member_scores.columns]]
name = "id"
type = "text_id"
length = 192
[[table.pool_member_scores.columns]]
name = "pool_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "pool_id"
type = "text_id"
length = 64
[[table.pool_member_scores.columns]]
name = "member_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "member_id"
type = "text_id"
length = 64
[[table.pool_member_scores.columns]]
name = "capability"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "scope_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "scope_id"
type = "text_id"
length = 128
nullable = true
[[table.pool_member_scores.columns]]
name = "score"
type = "float64"
default = 0
[[table.pool_member_scores.columns]]
name = "hard_state"
type = "text"
length = 64
default = "unknown"
[[table.pool_member_scores.columns]]
name = "score_version"
type = "int64"
default = 1
[[table.pool_member_scores.columns]]
name = "score_reason"
type = "json"
[[table.pool_member_scores.columns]]
name = "last_ranked_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_scheduled_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_success_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_failure_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "failure_count"
type = "int64"
default = 0
[[table.pool_member_scores.columns]]
name = "last_probe_attempt_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_probe_success_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_probe_failure_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "probe_failure_count"
type = "int64"
default = 0
[[table.pool_member_scores.columns]]
name = "probe_status"
type = "text"
length = 64
default = "never"
[[table.pool_member_scores.columns]]
name = "updated_at"
type = "unix_seconds"
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_rank_idx"
columns = ["pool_kind", "pool_id", "capability", "scope_kind", "scope_id", "hard_state", "score"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_member_idx"
columns = ["pool_kind", "pool_id", "member_kind", "member_id"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_probe_idx"
columns = ["pool_kind", "pool_id", "probe_status", "last_probe_success_at"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_updated_at_idx"
columns = ["updated_at"]
[table.api_key_provider_mappings]
domain = "provider_catalog"
order = 45

View File

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository,
};
use crate::repository::pool_scores::{
MysqlPoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
};
use crate::repository::provider_catalog::{
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository,
@@ -184,6 +187,14 @@ impl MysqlBackend {
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
}

View File

@@ -39,6 +39,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PostgresPoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqlxProviderCatalogReadRepository,
@@ -195,6 +198,14 @@ impl PostgresBackend {
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
}

View File

@@ -14,6 +14,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingReadRepository;
use crate::repository::global_models::GlobalModelReadRepository;
use crate::repository::management_tokens::ManagementTokenReadRepository;
use crate::repository::oauth_providers::OAuthProviderReadRepository;
use crate::repository::pool_scores::PoolScoreReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
use crate::repository::quota::ProviderQuotaReadRepository;
@@ -34,6 +35,7 @@ pub struct DataReadRepositories {
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
pool_scores: Option<Arc<dyn PoolScoreReadRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
@@ -61,6 +63,7 @@ impl fmt::Debug for DataReadRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field(
"has_minimal_candidate_selection",
@@ -124,6 +127,10 @@ impl DataReadRepositories {
.map(PostgresBackend::oauth_provider_read_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_read_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_read_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_read_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_read_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_read_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_read_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
@@ -208,6 +215,10 @@ impl DataReadRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolScoreReadRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
self.proxy_nodes.clone()
}
@@ -257,6 +268,7 @@ impl DataReadRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.minimal_candidate_selection.is_some()
|| self.request_candidates.is_some()

View File

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, SqlitePoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqliteProviderCatalogReadRepository,
@@ -197,6 +200,14 @@ impl SqliteBackend {
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
}

View File

@@ -11,6 +11,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
use crate::repository::global_models::GlobalModelWriteRepository;
use crate::repository::management_tokens::ManagementTokenWriteRepository;
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
use crate::repository::quota::ProviderQuotaWriteRepository;
@@ -30,6 +31,7 @@ pub struct DataWriteRepositories {
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
pool_scores: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
@@ -54,6 +56,7 @@ impl fmt::Debug for DataWriteRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_provider_quotas", &self.provider_quotas.is_some())
@@ -108,6 +111,10 @@ impl DataWriteRepositories {
.map(PostgresBackend::oauth_provider_write_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_write_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_write_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_write_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_write_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_write_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_write_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
@@ -184,6 +191,10 @@ impl DataWriteRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolMemberScoreWriteRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
self.proxy_nodes.clone()
}
@@ -218,6 +229,7 @@ impl DataWriteRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.provider_catalog.is_some()
|| self.provider_quotas.is_some()

View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260511120000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260511130000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

View File

@@ -300,6 +300,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260510120000,
20260511000000,
20260511120000,
20260511130000,
]
);
}
@@ -564,7 +565,8 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260509000000,
20260509120000,
20260510120000,
20260511120000
20260511120000,
20260511130000,
]
);
assert_eq!(
@@ -576,7 +578,8 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260509000000,
20260509120000,
20260510120000,
20260511120000
20260511120000,
20260511130000,
]
);
}
@@ -1088,6 +1091,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260510120000,
20260511000000,
20260511120000,
20260511130000,
]
);
}

View File

@@ -4,8 +4,8 @@ use async_trait::async_trait;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
};
use crate::DataLayerError;
@@ -138,6 +138,39 @@ impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelecti
.take(query.limit as usize)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<std::collections::BTreeMap<_, _>>();
let mut rows = self
.list_for_exact_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.provider_id == query.provider_id
&& row.endpoint_id == query.endpoint_id
&& row.model_id == query.model_id
&& key_order.contains_key(row.key_id.as_str())
})
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(rows)
}
}
fn sort_pool_key_rows(

View File

@@ -7,8 +7,8 @@ mod sqlite;
pub(crate) use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
};
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
pub use mysql::MysqlMinimalCandidateSelectionReadRepository;

View File

@@ -5,7 +5,8 @@ use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::mysql::MysqlPool;
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for MysqlMinimalCandidateSelectionR
.map(|item| item.row)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<BTreeMap<_, _>>();
let mut rows = self
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
&& key_order.contains_key(row.row.key_id.as_str())
})
.map(|item| item.row)
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(dedupe_candidate_selection_rows(rows))
}
}
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {

View File

@@ -5,7 +5,8 @@ use std::collections::BTreeSet;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::{error::SqlxResultExt, DataLayerError};
@@ -535,6 +536,13 @@ fn pool_key_candidate_selection_sql(order: &StoredPoolKeyCandidateOrder) -> Stri
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, &replacement)
}
fn pool_key_candidate_selection_by_key_ids_sql() -> String {
let default_order =
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC\nLIMIT $7\nOFFSET $8\n";
let replacement = "AND pak.id = ANY($7::text[])\nORDER BY\n array_position($7::text[], pak.id) ASC,\n pak.id ASC\n";
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, replacement)
}
#[derive(Debug, Clone)]
pub struct SqlxMinimalCandidateSelectionReadRepository {
pool: PgPool,
@@ -704,6 +712,51 @@ impl SqlxMinimalCandidateSelectionReadRepository {
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql = pool_key_candidate_selection_by_key_ids_sql();
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(sql.as_str())
.bind(api_format)
.bind(query.provider_id.as_str())
.bind(query.endpoint_id.as_str())
.bind(query.model_id.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(query.key_ids.clone())
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<std::collections::BTreeMap<_, _>>();
let mut rows = dedupe_candidate_selection_rows(rows);
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(rows)
}
}
fn requested_model_selection_sql() -> String {
@@ -917,6 +970,13 @@ impl MinimalCandidateSelectionReadRepository for SqlxMinimalCandidateSelectionRe
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group(self, query).await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group_key_ids(self, query).await
}
}
fn map_candidate_selection_row(

View File

@@ -5,7 +5,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::sqlite::SqlitePool;
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
.map(|item| item.row)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<BTreeMap<_, _>>();
let mut rows = self
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
&& key_order.contains_key(row.row.key_id.as_str())
})
.map(|item| item.row)
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(dedupe_candidate_selection_rows(rows))
}
}
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {

View File

@@ -16,6 +16,7 @@ pub mod gemini_file_mappings;
pub mod global_models;
pub mod management_tokens;
pub mod oauth_providers;
pub mod pool_scores;
pub mod provider_catalog;
pub mod provider_oauth;
pub mod proxy_nodes;

View File

@@ -0,0 +1,513 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::{
score_with_delta, GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery,
ListPoolMemberScoresQuery, ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
PoolMemberScheduleFeedback, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore,
};
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryPoolMemberScoreRepository {
scores: RwLock<BTreeMap<String, StoredPoolMemberScore>>,
}
impl InMemoryPoolMemberScoreRepository {
pub fn seed<I>(scores: I) -> Self
where
I: IntoIterator<Item = StoredPoolMemberScore>,
{
Self {
scores: RwLock::new(
scores
.into_iter()
.map(|score| (score.id.clone(), score))
.collect(),
),
}
}
fn matches_identity(score: &StoredPoolMemberScore, identity: &PoolMemberIdentity) -> bool {
score.pool_kind == identity.pool_kind
&& score.pool_id == identity.pool_id
&& score.member_kind == identity.member_kind
&& score.member_id == identity.member_id
}
fn matches_scope(score: &StoredPoolMemberScore, scope: Option<&PoolScoreScope>) -> bool {
let Some(scope) = scope else {
return true;
};
score.capability == scope.capability
&& score.scope_kind == scope.scope_kind
&& score.scope_id == scope.scope_id
}
fn sort_ranked(scores: &mut [StoredPoolMemberScore]) {
scores.sort_by(|left, right| {
right
.score
.partial_cmp(&left.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| {
right
.last_ranked_at
.unwrap_or(0)
.cmp(&left.last_ranked_at.unwrap_or(0))
})
.then_with(|| left.member_id.cmp(&right.member_id))
.then_with(|| left.id.cmp(&right.id))
});
}
fn sort_probe(scores: &mut [StoredPoolMemberScore]) {
scores.sort_by(|left, right| {
probe_priority(left)
.cmp(&probe_priority(right))
.then_with(|| right.probe_failure_count.cmp(&left.probe_failure_count))
.then_with(|| {
left.last_probe_success_at
.unwrap_or(0)
.cmp(&right.last_probe_success_at.unwrap_or(0))
})
.then_with(|| {
left.last_scheduled_at
.unwrap_or(0)
.cmp(&right.last_scheduled_at.unwrap_or(0))
.reverse()
})
.then_with(|| left.member_id.cmp(&right.member_id))
});
}
}
#[async_trait]
impl PoolScoreReadRepository for InMemoryPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let hard_states = query
.hard_states
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>();
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
items
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>()
});
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& score.capability == query.capability
&& score.scope_kind == query.scope_kind
&& score.scope_id == query.scope_id
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
&& probe_statuses
.as_ref()
.is_none_or(|statuses| statuses.contains(&score.probe_status))
})
.cloned()
.collect::<Vec<_>>();
Self::sort_ranked(&mut scores);
Ok(scores
.into_iter()
.skip(query.offset)
.take(query.limit.max(1))
.collect())
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let hard_states = query
.hard_states
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>();
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
items
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>()
});
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& query
.capability
.as_ref()
.is_none_or(|capability| score.capability == *capability)
&& query
.scope_kind
.as_ref()
.is_none_or(|scope_kind| score.scope_kind == *scope_kind)
&& query
.scope_id
.as_ref()
.is_none_or(|scope_id| score.scope_id.as_ref() == Some(scope_id))
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
&& probe_statuses
.as_ref()
.is_none_or(|statuses| statuses.contains(&score.probe_status))
})
.cloned()
.collect::<Vec<_>>();
Self::sort_ranked(&mut scores);
Ok(scores
.into_iter()
.skip(query.offset)
.take(query.limit.max(1))
.collect())
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& query
.capability
.as_ref()
.is_none_or(|capability| score.capability == *capability)
&& matches!(
score.hard_state,
PoolMemberHardState::Available
| PoolMemberHardState::Unknown
| PoolMemberHardState::Cooldown
| PoolMemberHardState::QuotaExhausted
)
&& match score.probe_status {
PoolMemberProbeStatus::Never
| PoolMemberProbeStatus::Failed
| PoolMemberProbeStatus::Stale => true,
PoolMemberProbeStatus::Ok => score
.last_probe_success_at
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
PoolMemberProbeStatus::InProgress => score
.last_probe_attempt_at
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
}
})
.cloned()
.collect::<Vec<_>>();
Self::sort_probe(&mut scores);
Ok(scores.into_iter().take(query.limit.max(1)).collect())
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let ids = query
.ids
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| ids.contains(&score.id))
.cloned()
.collect::<Vec<_>>();
Ok(scores)
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for InMemoryPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
self.scores
.write()
.expect("pool member score repository lock")
.insert(stored.id.clone(), stored.clone());
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &attempt.identity)
|| !Self::matches_scope(score, attempt.scope.as_ref())
{
continue;
}
score.last_probe_attempt_at = Some(attempt.attempted_at);
score.probe_status = PoolMemberProbeStatus::InProgress;
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
attempt.score_reason_patch.clone(),
);
score.updated_at = attempt.attempted_at;
updated += 1;
}
Ok(updated)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &result.identity)
|| !Self::matches_scope(score, result.scope.as_ref())
{
continue;
}
score.last_probe_attempt_at = Some(result.attempted_at);
score.probe_status = result.probe_status;
if result.succeeded {
score.last_probe_success_at = Some(result.attempted_at);
score.probe_failure_count = 0;
} else {
score.last_probe_failure_at = Some(result.attempted_at);
score.probe_failure_count = score.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
score.hard_state = hard_state;
}
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
result.score_reason_patch.clone(),
);
score.updated_at = result.attempted_at;
updated += 1;
}
Ok(updated)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &feedback.identity)
|| !Self::matches_scope(score, feedback.scope.as_ref())
{
continue;
}
score.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => {
score.last_success_at = Some(feedback.scheduled_at);
}
Some(false) => {
score.last_failure_at = Some(feedback.scheduled_at);
score.failure_count = score.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
score.hard_state = hard_state;
}
score.score = score_with_delta(score.score, feedback.score_delta);
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
feedback.score_reason_patch.clone(),
);
score.updated_at = feedback.scheduled_at;
updated += 1;
}
Ok(updated)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if Self::matches_identity(score, identity) && Self::matches_scope(score, scope) {
score.hard_state = hard_state;
score.updated_at = updated_at;
updated += 1;
}
}
Ok(updated)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
let before = guard.len();
guard.retain(|_, score| !Self::matches_identity(score, identity));
Ok(before.saturating_sub(guard.len()))
}
}
fn probe_priority(score: &StoredPoolMemberScore) -> u8 {
if score.last_scheduled_at.is_some() && score.probe_status != PoolMemberProbeStatus::Ok {
return 0;
}
match score.hard_state {
PoolMemberHardState::QuotaExhausted => 1,
PoolMemberHardState::Unknown => 2,
_ if score.probe_status == PoolMemberProbeStatus::Stale => 3,
_ => 4,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::repository::pool_scores::{
POOL_KIND_PROVIDER_KEY_POOL, POOL_MEMBER_KIND_PROVIDER_API_KEY,
POOL_SCORE_CAPABILITY_API_FORMAT, POOL_SCORE_SCOPE_KIND_MODEL,
};
fn score(id: &str, member_id: &str, value: f64) -> StoredPoolMemberScore {
StoredPoolMemberScore {
id: id.to_string(),
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: member_id.to_string(),
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
score: value,
hard_state: PoolMemberHardState::Available,
score_version: 1,
score_reason: serde_json::json!({}),
last_ranked_at: Some(1),
last_scheduled_at: None,
last_success_at: None,
last_failure_at: None,
failure_count: 0,
last_probe_attempt_at: None,
last_probe_success_at: None,
last_probe_failure_at: None,
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Never,
updated_at: 1,
}
}
#[tokio::test]
async fn lists_ranked_members_by_score() {
let repository = InMemoryPoolMemberScoreRepository::seed(vec![
score("score-1", "key-1", 0.2),
score("score-2", "key-2", 0.9),
]);
let rows = repository
.list_ranked_pool_members(&ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
hard_states: vec![PoolMemberHardState::Available],
probe_statuses: None,
offset: 0,
limit: 10,
})
.await
.expect("list should succeed");
assert_eq!(
rows.into_iter()
.map(|row| row.member_id)
.collect::<Vec<_>>(),
vec!["key-2".to_string(), "key-1".to_string()]
);
}
#[tokio::test]
async fn marks_probe_in_progress_without_incrementing_failure_count() {
let repository =
InMemoryPoolMemberScoreRepository::seed(vec![score("score-1", "key-1", 0.2)]);
let updated = repository
.mark_pool_member_probe_in_progress(PoolMemberProbeAttempt {
identity: PoolMemberIdentity::provider_api_key("provider-1", "key-1"),
scope: None,
attempted_at: 100,
score_reason_patch: Some(serde_json::json!({ "last_probe": "in_progress" })),
})
.await
.expect("mark should succeed");
assert_eq!(updated, 1);
let rows = repository
.list_pool_member_scores(&ListPoolMemberScoresQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
capability: None,
scope_kind: None,
scope_id: None,
hard_states: Vec::new(),
probe_statuses: Some(vec![PoolMemberProbeStatus::InProgress]),
offset: 0,
limit: 10,
})
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].last_probe_attempt_at, Some(100));
assert_eq!(rows[0].probe_failure_count, 0);
}
}

View File

@@ -0,0 +1,54 @@
pub use aether_data_contracts::repository::pool_scores::*;
mod memory;
mod mysql;
mod postgres;
mod sqlite;
pub use memory::InMemoryPoolMemberScoreRepository;
pub use mysql::MysqlPoolMemberScoreRepository;
pub use postgres::PostgresPoolMemberScoreRepository;
pub use sqlite::SqlitePoolMemberScoreRepository;
fn merge_score_reason_patch(
mut current: serde_json::Value,
patch: Option<serde_json::Value>,
) -> serde_json::Value {
let Some(patch) = patch else {
return current;
};
match (current.as_object_mut(), patch) {
(Some(current), serde_json::Value::Object(patch)) => {
for (key, value) in patch {
current.insert(key, value);
}
serde_json::Value::Object(current.clone())
}
(_, patch) => patch,
}
}
fn score_with_delta(score: f64, delta_basis_points: Option<i32>) -> f64 {
let delta = delta_basis_points.unwrap_or_default() as f64 / 10_000.0;
(score + delta).clamp(0.0, 1.0)
}
fn i64_from_u64(value: u64, field: &str) -> Result<i64, crate::DataLayerError> {
i64::try_from(value).map_err(|_| {
crate::DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range"))
})
}
fn i64_opt_from_u64(value: Option<u64>, field: &str) -> Result<Option<i64>, crate::DataLayerError> {
value.map(|value| i64_from_u64(value, field)).transpose()
}
fn u64_from_i64(value: i64, field: &str) -> Result<u64, crate::DataLayerError> {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("{field} is negative: {value}"))
})
}
fn u64_opt_from_i64(value: Option<i64>, field: &str) -> Result<Option<u64>, crate::DataLayerError> {
value.map(|value| u64_from_i64(value, field)).transpose()
}

View File

@@ -0,0 +1,586 @@
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, MySql, MySqlPool, QueryBuilder, Row};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct MysqlPoolMemberScoreRepository {
pool: MySqlPool,
}
impl MysqlPoolMemberScoreRepository {
pub fn new(pool: MySqlPool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for MysqlPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for MysqlPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
let score_reason = serde_json::to_string(&stored.score_reason)
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
)
ON DUPLICATE KEY UPDATE
pool_kind = VALUES(pool_kind),
pool_id = VALUES(pool_id),
member_kind = VALUES(member_kind),
member_id = VALUES(member_id),
capability = VALUES(capability),
scope_kind = VALUES(scope_kind),
scope_id = VALUES(scope_id),
score = VALUES(score),
hard_state = VALUES(hard_state),
score_version = VALUES(score_version),
score_reason = VALUES(score_reason),
last_ranked_at = VALUES(last_ranked_at),
last_scheduled_at = COALESCE(VALUES(last_scheduled_at), last_scheduled_at),
last_success_at = COALESCE(VALUES(last_success_at), last_success_at),
last_failure_at = COALESCE(VALUES(last_failure_at), last_failure_at),
failure_count = VALUES(failure_count),
last_probe_attempt_at = COALESCE(VALUES(last_probe_attempt_at), last_probe_attempt_at),
last_probe_success_at = COALESCE(VALUES(last_probe_success_at), last_probe_success_at),
last_probe_failure_at = COALESCE(VALUES(last_probe_failure_at), last_probe_failure_at),
probe_failure_count = VALUES(probe_failure_count),
probe_status = VALUES(probe_status),
updated_at = VALUES(updated_at)
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(score_reason)
.bind(i64_opt_from_u64(
stored.last_ranked_at,
"pool score last_ranked_at",
)?)
.bind(i64_opt_from_u64(
stored.last_scheduled_at,
"pool score last_scheduled_at",
)?)
.bind(i64_opt_from_u64(
stored.last_success_at,
"pool score last_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_failure_at,
"pool score last_failure_at",
)?)
.bind(i64_from_u64(
stored.failure_count,
"pool score failure_count",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &MySqlRow) -> Result<StoredPoolMemberScore, DataLayerError> {
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_sql_err()?,
pool_kind: row.try_get("pool_kind").map_sql_err()?,
pool_id: row.try_get("pool_id").map_sql_err()?,
member_kind: row.try_get("member_kind").map_sql_err()?,
member_id: row.try_get("member_id").map_sql_err()?,
capability: row.try_get("capability").map_sql_err()?,
scope_kind: row.try_get("scope_kind").map_sql_err()?,
scope_id: row.try_get("scope_id").map_sql_err()?,
score: row.try_get("score").map_sql_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_sql_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_sql_err()?,
"pool_member_scores.score_version",
)?,
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_sql_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_sql_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_sql_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_sql_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_sql_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_sql_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_sql_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_sql_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_sql_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_sql_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_sql_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
last_ranked_at: score.last_ranked_at,
last_scheduled_at: score.last_scheduled_at,
last_success_at: score.last_success_at,
last_failure_at: score.last_failure_at,
failure_count: score.failure_count,
last_probe_attempt_at: score.last_probe_attempt_at,
last_probe_success_at: score.last_probe_success_at,
last_probe_failure_at: score.last_probe_failure_at,
probe_failure_count: score.probe_failure_count,
probe_status: score.probe_status,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}

View File

@@ -0,0 +1,588 @@
use async_trait::async_trait;
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlxResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct PostgresPoolMemberScoreRepository {
pool: PgPool,
}
impl PostgresPoolMemberScoreRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for PostgresPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.clone().into_stored();
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23
)
ON CONFLICT(id) DO UPDATE SET
pool_kind = EXCLUDED.pool_kind,
pool_id = EXCLUDED.pool_id,
member_kind = EXCLUDED.member_kind,
member_id = EXCLUDED.member_id,
capability = EXCLUDED.capability,
scope_kind = EXCLUDED.scope_kind,
scope_id = EXCLUDED.scope_id,
score = EXCLUDED.score,
hard_state = EXCLUDED.hard_state,
score_version = EXCLUDED.score_version,
score_reason = EXCLUDED.score_reason,
last_ranked_at = EXCLUDED.last_ranked_at,
last_scheduled_at = COALESCE(EXCLUDED.last_scheduled_at, pool_member_scores.last_scheduled_at),
last_success_at = COALESCE(EXCLUDED.last_success_at, pool_member_scores.last_success_at),
last_failure_at = COALESCE(EXCLUDED.last_failure_at, pool_member_scores.last_failure_at),
failure_count = EXCLUDED.failure_count,
last_probe_attempt_at = COALESCE(EXCLUDED.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
last_probe_success_at = COALESCE(EXCLUDED.last_probe_success_at, pool_member_scores.last_probe_success_at),
last_probe_failure_at = COALESCE(EXCLUDED.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
probe_failure_count = EXCLUDED.probe_failure_count,
probe_status = EXCLUDED.probe_status,
updated_at = EXCLUDED.updated_at
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(&stored.score_reason)
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = $1 AND pool_id = $2 AND member_kind = $3 AND member_id = $4
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &PgRow) -> Result<StoredPoolMemberScore, DataLayerError> {
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_postgres_err()?,
pool_kind: row.try_get("pool_kind").map_postgres_err()?,
pool_id: row.try_get("pool_id").map_postgres_err()?,
member_kind: row.try_get("member_kind").map_postgres_err()?,
member_id: row.try_get("member_id").map_postgres_err()?,
capability: row.try_get("capability").map_postgres_err()?,
scope_kind: row.try_get("scope_kind").map_postgres_err()?,
scope_id: row.try_get("scope_id").map_postgres_err()?,
score: row.try_get("score").map_postgres_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_postgres_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_postgres_err()?,
"pool_member_scores.score_version",
)?,
score_reason: row.try_get("score_reason").map_postgres_err()?,
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_postgres_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_postgres_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_postgres_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_postgres_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_postgres_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_postgres_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_postgres_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_postgres_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_postgres_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_postgres_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_postgres_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
last_ranked_at: score.last_ranked_at,
last_scheduled_at: score.last_scheduled_at,
last_success_at: score.last_success_at,
last_failure_at: score.last_failure_at,
failure_count: score.failure_count,
last_probe_attempt_at: score.last_probe_attempt_at,
last_probe_success_at: score.last_probe_success_at,
last_probe_failure_at: score.last_probe_failure_at,
probe_failure_count: score.probe_failure_count,
probe_status: score.probe_status,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}

View File

@@ -0,0 +1,571 @@
use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct SqlitePoolMemberScoreRepository {
pool: SqlitePool,
}
impl SqlitePoolMemberScoreRepository {
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for SqlitePoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for SqlitePoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
let score_reason = serde_json::to_string(&stored.score_reason)
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
)
ON CONFLICT(id) DO UPDATE SET
pool_kind = excluded.pool_kind,
pool_id = excluded.pool_id,
member_kind = excluded.member_kind,
member_id = excluded.member_id,
capability = excluded.capability,
scope_kind = excluded.scope_kind,
scope_id = excluded.scope_id,
score = excluded.score,
hard_state = excluded.hard_state,
score_version = excluded.score_version,
score_reason = excluded.score_reason,
last_ranked_at = excluded.last_ranked_at,
last_scheduled_at = COALESCE(excluded.last_scheduled_at, pool_member_scores.last_scheduled_at),
last_success_at = COALESCE(excluded.last_success_at, pool_member_scores.last_success_at),
last_failure_at = COALESCE(excluded.last_failure_at, pool_member_scores.last_failure_at),
failure_count = excluded.failure_count,
last_probe_attempt_at = COALESCE(excluded.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
last_probe_success_at = COALESCE(excluded.last_probe_success_at, pool_member_scores.last_probe_success_at),
last_probe_failure_at = COALESCE(excluded.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
probe_failure_count = excluded.probe_failure_count,
probe_status = excluded.probe_status,
updated_at = excluded.updated_at
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(score_reason)
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &SqliteRow) -> Result<StoredPoolMemberScore, DataLayerError> {
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_sql_err()?,
pool_kind: row.try_get("pool_kind").map_sql_err()?,
pool_id: row.try_get("pool_id").map_sql_err()?,
member_kind: row.try_get("member_kind").map_sql_err()?,
member_id: row.try_get("member_id").map_sql_err()?,
capability: row.try_get("capability").map_sql_err()?,
scope_kind: row.try_get("scope_kind").map_sql_err()?,
scope_id: row.try_get("scope_id").map_sql_err()?,
score: row.try_get("score").map_sql_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_sql_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_sql_err()?,
"pool_member_scores.score_version",
)?,
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_sql_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_sql_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_sql_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_sql_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_sql_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_sql_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_sql_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_sql_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_sql_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_sql_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_sql_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
last_ranked_at: score.last_ranked_at,
last_scheduled_at: score.last_scheduled_at,
last_success_at: score.last_success_at,
last_failure_at: score.last_failure_at,
failure_count: score.failure_count,
last_probe_attempt_at: score.last_probe_attempt_at,
last_probe_success_at: score.last_probe_success_at,
last_probe_failure_at: score.last_probe_failure_at,
probe_failure_count: score.probe_failure_count,
probe_status: score.probe_status,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}