fix(routing): harden routed pool scheduling

This commit is contained in:
elky
2026-07-27 22:06:28 +08:00
parent 550cc36760
commit 4148ab1931
62 changed files with 5510 additions and 850 deletions
+130 -32
View File
@@ -14,6 +14,8 @@ pub(crate) enum GatewayRoutingSelectionError {
Disabled(String),
#[error("routing group was explicitly requested but is not allowed for this principal: {0}")]
Forbidden(String),
#[error("routing group repository lookup failed: {0}")]
Repository(String),
}
#[derive(Debug, Clone, Default)]
@@ -42,22 +44,21 @@ pub(crate) async fn select_gateway_routing_group(
let group = repository
.find_routing_group(RoutingGroupLookupKey::Id(explicit))
.await
.ok()
.flatten()
.or({
let group: Option<StoredRoutingGroup> = repository
.find_routing_group(RoutingGroupLookupKey::Name(explicit))
.await
.unwrap_or_default();
group
});
.map_err(repository_selection_error)?;
let group = match group {
Some(group) => Some(group),
None => repository
.find_routing_group(RoutingGroupLookupKey::Name(explicit))
.await
.map_err(repository_selection_error)?,
};
let Some(group) = group else {
return Err(GatewayRoutingSelectionError::NotFound(explicit.to_string()));
};
if !group.enabled {
return Err(GatewayRoutingSelectionError::Disabled(group.id));
}
if !explicit_group_allowed(repository, &group.id, &input).await {
if !explicit_group_allowed(repository, &group, &input).await? {
return Err(GatewayRoutingSelectionError::Forbidden(group.id));
}
return Ok(GatewayRoutingGroupSelection {
@@ -70,13 +71,15 @@ pub(crate) async fn select_gateway_routing_group(
// produce a group. The data-state repository answers this with a cached
// existence query, so the common "routing configured but unused" case
// does not materialize the binding table per API key/user.
let has_bindings = repository.has_any_routing_group_binding().await;
if matches!(has_bindings, Ok(false)) {
let has_bindings = repository
.has_any_routing_group_binding()
.await
.map_err(repository_selection_error)?;
if !has_bindings {
let system_default = repository
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
.await
.ok()
.flatten()
.map_err(repository_selection_error)?
.filter(|group| group.enabled);
return Ok(GatewayRoutingGroupSelection {
group: system_default,
@@ -92,13 +95,12 @@ pub(crate) async fn select_gateway_routing_group(
subject_id: Some(subject_id.to_string()),
})
.await
.unwrap_or_default();
.map_err(repository_selection_error)?;
for binding in bindings.into_iter().filter(|binding| binding.is_default) {
let group = repository
.find_routing_group(RoutingGroupLookupKey::Id(&binding.group_id))
.await
.ok()
.flatten();
.map_err(repository_selection_error)?;
if let Some(group) = group.filter(|group| group.enabled) {
return Ok(GatewayRoutingGroupSelection {
group: Some(group),
@@ -111,8 +113,7 @@ pub(crate) async fn select_gateway_routing_group(
let system_default = repository
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
.await
.ok()
.flatten()
.map_err(repository_selection_error)?
.filter(|group| group.enabled);
Ok(GatewayRoutingGroupSelection {
group: system_default,
@@ -122,31 +123,30 @@ pub(crate) async fn select_gateway_routing_group(
async fn explicit_group_allowed(
repository: &(impl RoutingGroupReadRepository + ?Sized),
group_id: &str,
group: &StoredRoutingGroup,
input: &GatewayRoutingSelectionInput<'_>,
) -> bool {
if let Ok(Some(group)) = repository
.find_routing_group(RoutingGroupLookupKey::Id(group_id))
.await
{
if group.is_system_default {
return true;
}
) -> Result<bool, GatewayRoutingSelectionError> {
if group.is_system_default {
return Ok(true);
}
for (subject_type, subject_id, _) in default_binding_candidates(input) {
let bindings = repository
.list_routing_group_bindings(&RoutingGroupBindingQuery {
group_id: Some(group_id.to_string()),
group_id: Some(group.id.clone()),
subject_type: Some(subject_type),
subject_id: Some(subject_id.to_string()),
})
.await
.unwrap_or_default();
.map_err(repository_selection_error)?;
if bindings.iter().any(|binding| binding.allow_explicit_select) {
return true;
return Ok(true);
}
}
false
Ok(false)
}
fn repository_selection_error(error: impl std::fmt::Display) -> GatewayRoutingSelectionError {
GatewayRoutingSelectionError::Repository(error.to_string())
}
fn default_binding_candidates<'a>(
@@ -178,11 +178,57 @@ mod tests {
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
use aether_data_contracts::repository::routing_profiles::{
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, RoutingGroupWriteRepository,
StoredRoutingGroupBinding, StoredRoutingGroupVersion,
};
use aether_data_contracts::DataLayerError;
use async_trait::async_trait;
use serde_json::json;
use super::*;
struct FailingRoutingGroupRepository {
id_lookup_is_missing: bool,
}
impl FailingRoutingGroupRepository {
fn failure<T>() -> Result<T, DataLayerError> {
Err(DataLayerError::Sql(
"routing repository unavailable".to_string(),
))
}
}
#[async_trait]
impl RoutingGroupReadRepository for FailingRoutingGroupRepository {
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
Self::failure()
}
async fn find_routing_group(
&self,
lookup: RoutingGroupLookupKey<'_>,
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
if self.id_lookup_is_missing && matches!(lookup, RoutingGroupLookupKey::Id(_)) {
return Ok(None);
}
Self::failure()
}
async fn list_routing_group_bindings(
&self,
_query: &RoutingGroupBindingQuery,
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
Self::failure()
}
async fn list_routing_group_versions(
&self,
_group_id: &str,
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
Self::failure()
}
}
#[tokio::test]
async fn selects_api_key_default_binding() {
let repository = InMemoryRoutingGroupRepository::default();
@@ -336,6 +382,58 @@ mod tests {
);
}
#[tokio::test]
async fn propagates_explicit_name_lookup_failure_after_missing_id() {
let repository = FailingRoutingGroupRepository {
id_lookup_is_missing: true,
};
let error = select_gateway_routing_group(
&repository,
GatewayRoutingSelectionInput {
explicit_group: Some("group-name"),
user_id: Some("user-1"),
api_key_id: Some("api-key-1"),
user_group_ids: &[],
},
)
.await
.unwrap_err();
assert_eq!(
error,
GatewayRoutingSelectionError::Repository(
"sql error: routing repository unavailable".to_string()
)
);
}
#[tokio::test]
async fn propagates_implicit_binding_lookup_failure() {
let repository = FailingRoutingGroupRepository {
id_lookup_is_missing: false,
};
let error = select_gateway_routing_group(
&repository,
GatewayRoutingSelectionInput {
explicit_group: None,
user_id: Some("user-1"),
api_key_id: Some("api-key-1"),
user_group_ids: &[],
},
)
.await
.unwrap_err();
assert_eq!(
error,
GatewayRoutingSelectionError::Repository(
"sql error: routing repository unavailable".to_string()
)
);
}
#[tokio::test]
async fn rejects_explicit_disabled_group() {
let repository = InMemoryRoutingGroupRepository::default();