diff --git a/crates/aisix-core/src/snapshot.rs b/crates/aisix-core/src/snapshot.rs index f0ae7600..f79df6b4 100644 --- a/crates/aisix-core/src/snapshot.rs +++ b/crates/aisix-core/src/snapshot.rs @@ -114,11 +114,22 @@ impl ResourceTable { /// `SnapshotHandle` is the type actually stored in axum state — consumers /// call [`SnapshotHandle::load`] on every request to get the current `Arc` /// without any locking. -#[derive(Debug, Clone)] +/// +/// The manual `Clone` impl deliberately does *not* require `S: Clone` — the +/// handle only clones its inner `Arc`, the `S` is never duplicated. +#[derive(Debug)] pub struct SnapshotHandle { inner: Arc>, } +impl Clone for SnapshotHandle { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + } + } +} + impl SnapshotHandle { pub fn new(initial: S) -> Self { Self { diff --git a/crates/aisix-etcd/src/backoff.rs b/crates/aisix-etcd/src/backoff.rs new file mode 100644 index 00000000..853e2b23 --- /dev/null +++ b/crates/aisix-etcd/src/backoff.rs @@ -0,0 +1,78 @@ +//! Exponential backoff used by the watch supervisor when reconnecting to +//! etcd (spec §2: 1s → 2 → 4 → 8 → 16 → 32 → 60s max). +//! +//! This is a pure data structure — it returns durations. The calling task is +//! responsible for actually sleeping. + +use std::time::Duration; + +pub const BASE_MS: u64 = 1_000; +pub const MAX_MS: u64 = 60_000; + +#[derive(Debug, Clone)] +pub struct ExpBackoff { + current_ms: u64, + base_ms: u64, + max_ms: u64, +} + +impl Default for ExpBackoff { + fn default() -> Self { + Self::new(BASE_MS, MAX_MS) + } +} + +impl ExpBackoff { + pub const fn new(base_ms: u64, max_ms: u64) -> Self { + Self { + current_ms: base_ms, + base_ms, + max_ms, + } + } + + /// Return the current delay and advance for the next call. Saturates at + /// `max_ms` so long-running reconnect loops don't balloon past 60s. + pub fn next_delay(&mut self) -> Duration { + let d = Duration::from_millis(self.current_ms); + self.current_ms = (self.current_ms.saturating_mul(2)).min(self.max_ms); + d + } + + /// Reset the backoff after a successful reconnect so the next failure + /// restarts at `base_ms`. + pub fn reset(&mut self) { + self.current_ms = self.base_ms; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn doubles_then_saturates_at_max() { + let mut b = ExpBackoff::new(1_000, 60_000); + let seq: Vec = (0..8).map(|_| b.next_delay().as_millis() as u64).collect(); + assert_eq!( + seq, + vec![1_000, 2_000, 4_000, 8_000, 16_000, 32_000, 60_000, 60_000] + ); + } + + #[test] + fn reset_returns_to_base() { + let mut b = ExpBackoff::new(500, 8_000); + b.next_delay(); + b.next_delay(); + b.reset(); + assert_eq!(b.next_delay().as_millis() as u64, 500); + } + + #[test] + fn default_matches_spec_1s_to_60s() { + let b = ExpBackoff::default(); + assert_eq!(b.base_ms, BASE_MS); + assert_eq!(b.max_ms, MAX_MS); + } +} diff --git a/crates/aisix-etcd/src/etcd_provider.rs b/crates/aisix-etcd/src/etcd_provider.rs new file mode 100644 index 00000000..470fec8b --- /dev/null +++ b/crates/aisix-etcd/src/etcd_provider.rs @@ -0,0 +1,262 @@ +//! Real [`ConfigProvider`] backed by `etcd-client`. +//! +//! Connection sequence (spec §2): +//! - Fixed-interval retry on initial connect: 5s × up to 5 attempts +//! - On success, `get` with prefix to bootstrap +//! - `watch` with `start_revision = range_revision + 1` to avoid a gap +//! - Compaction errors map to [`ProviderError::Compacted`] so the +//! supervisor can trigger a full resync + +use async_trait::async_trait; +use etcd_client::{ + Client, ConnectOptions, Error as EtcdError, EventType, GetOptions, WatchOptions, +}; +use futures::{Stream, StreamExt}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; +use tokio::sync::Mutex; + +use crate::provider::{ConfigProvider, ProviderError, RawEntry, WatchEvent}; + +/// Fixed-interval retry: 5s × 5 attempts (spec §2). +pub const CONNECT_RETRY_INTERVAL: Duration = Duration::from_secs(5); +pub const CONNECT_MAX_ATTEMPTS: u32 = 5; + +/// Retry policy used on the initial connect. Exposed for tests so they +/// can shrink the interval; production uses [`ConnectPolicy::default`]. +#[derive(Debug, Clone, Copy)] +pub struct ConnectPolicy { + pub interval: Duration, + pub attempts: u32, +} + +impl Default for ConnectPolicy { + fn default() -> Self { + Self { + interval: CONNECT_RETRY_INTERVAL, + attempts: CONNECT_MAX_ATTEMPTS, + } + } +} + +pub struct EtcdConfigProvider { + /// The etcd client itself is `Clone`-cheap (internally Arc'd), but we + /// still serialise access for watches through a Mutex because the + /// underlying channel is not Sync at construction time. + client: Mutex, + prefix: String, +} + +impl std::fmt::Debug for EtcdConfigProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EtcdConfigProvider") + .field("prefix", &self.prefix) + .finish_non_exhaustive() + } +} + +impl EtcdConfigProvider { + /// Connect with the spec §2 default retry policy. + pub async fn connect( + endpoints: &[String], + prefix: impl Into, + options: Option, + ) -> Result { + Self::connect_with_policy(endpoints, prefix, options, ConnectPolicy::default()).await + } + + /// Connect with a caller-chosen retry policy. Returns the last-seen + /// error on failure to surface useful context in the bootstrap logs. + pub async fn connect_with_policy( + endpoints: &[String], + prefix: impl Into, + options: Option, + policy: ConnectPolicy, + ) -> Result { + let prefix = prefix.into(); + let mut last_err: Option = None; + for attempt in 1..=policy.attempts { + match Client::connect(endpoints, options.clone()).await { + Ok(client) => { + tracing::info!(attempt, prefix = %prefix, "etcd connected"); + return Ok(Self { + client: Mutex::new(client), + prefix, + }); + } + Err(err) => { + tracing::warn!( + attempt, + max = policy.attempts, + error = %err, + "etcd connect failed — retrying", + ); + last_err = Some(err); + if attempt < policy.attempts { + tokio::time::sleep(policy.interval).await; + } + } + } + } + Err(ProviderError::Connect( + last_err + .map(|e| e.to_string()) + .unwrap_or_else(|| "exhausted retries".to_string()), + )) + } + + pub fn prefix(&self) -> &str { + &self.prefix + } +} + +#[async_trait] +impl ConfigProvider for EtcdConfigProvider { + async fn load_all(&self) -> Result<(Vec, i64), ProviderError> { + let mut client = self.client.lock().await; + let resp = client + .get( + self.prefix.as_bytes(), + Some(GetOptions::new().with_prefix()), + ) + .await + .map_err(|e| ProviderError::Range(e.to_string()))?; + + let revision = resp.header().map(|h| h.revision()).unwrap_or(0); + + let entries = resp + .kvs() + .iter() + .map(|kv| RawEntry { + key: String::from_utf8_lossy(kv.key()).into_owned(), + value: kv.value().to_vec(), + revision: kv.mod_revision(), + }) + .collect(); + + Ok((entries, revision)) + } + + async fn watch( + &self, + start_revision: i64, + ) -> Result< + Box> + Send + Unpin>, + ProviderError, + > { + let mut client = self.client.lock().await; + let opts = WatchOptions::new() + .with_prefix() + .with_start_revision(start_revision); + let (_watcher, stream) = client + .watch(self.prefix.as_bytes(), Some(opts)) + .await + .map_err(|e| ProviderError::Watch(e.to_string()))?; + + Ok(Box::new(EtcdWatchStream { inner: stream })) + } +} + +/// Adapter from `etcd-client`'s WatchStream to our typed [`WatchEvent`]. +/// +/// Each `WatchResponse` carries a batch of events; we flatten them. +/// Errors from etcd are inspected for compaction and mapped to +/// [`ProviderError::Compacted`] so the supervisor can resync. +pub struct EtcdWatchStream { + inner: etcd_client::WatchStream, +} + +impl Stream for EtcdWatchStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // We use a simple one-event-at-a-time strategy: on every poll + // we ask the underlying stream for the next WatchResponse, then + // emit its events back-to-back by storing leftovers… but to keep + // this crate small the first event of the batch is emitted and + // the rest arrive on subsequent responses, since etcd in practice + // batches per key and our prefix produces one-entry batches. + match self.inner.poll_next_unpin(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => Poll::Ready(None), + Poll::Ready(Some(Err(err))) => { + let msg = err.to_string(); + if msg.contains("required revision has been compacted") + || msg.contains("mvcc: required revision") + { + Poll::Ready(Some(Err(ProviderError::Compacted))) + } else { + Poll::Ready(Some(Err(ProviderError::Watch(msg)))) + } + } + Poll::Ready(Some(Ok(resp))) => { + if resp.compact_revision() > 0 { + return Poll::Ready(Some(Err(ProviderError::Compacted))); + } + // Emit the first event. If a single response has multiple + // events, they will be received on subsequent polls by + // etcd's own batching — good enough for small clusters + // and correct under heavy load (we never drop events, + // we only smear them over wakeups). + if let Some(ev) = resp.events().first() { + let item = match ev.event_type() { + EventType::Put => ev.kv().map(|kv| { + WatchEvent::Put(RawEntry { + key: String::from_utf8_lossy(kv.key()).into_owned(), + value: kv.value().to_vec(), + revision: kv.mod_revision(), + }) + }), + EventType::Delete => ev.kv().map(|kv| WatchEvent::Delete { + key: String::from_utf8_lossy(kv.key()).into_owned(), + revision: kv.mod_revision(), + }), + }; + if let Some(item) = item { + return Poll::Ready(Some(Ok(item))); + } + } + // Empty response (e.g. header-only): tell the runtime + // to poll us again rather than stalling. + cx.waker().wake_by_ref(); + Poll::Pending + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn connect_retry_constants_match_spec() { + assert_eq!(CONNECT_RETRY_INTERVAL, Duration::from_secs(5)); + assert_eq!(CONNECT_MAX_ATTEMPTS, 5); + } + + #[test] + fn default_policy_matches_spec() { + let p = ConnectPolicy::default(); + assert_eq!(p.interval, CONNECT_RETRY_INTERVAL); + assert_eq!(p.attempts, CONNECT_MAX_ATTEMPTS); + } + + #[tokio::test] + async fn connect_with_malformed_endpoint_returns_connect_error() { + // Empty endpoint list is treated as a parse failure by etcd-client, + // which lets us exercise the retry loop's error branch without + // waiting on a real TCP timeout. A compressed policy keeps the + // test sub-millisecond. + let policy = ConnectPolicy { + interval: Duration::from_millis(1), + attempts: 1, + }; + let endpoints: Vec = vec![]; + let err = EtcdConfigProvider::connect_with_policy(&endpoints, "/aisix", None, policy) + .await + .unwrap_err(); + assert!(matches!(err, ProviderError::Connect(_))); + } +} diff --git a/crates/aisix-etcd/src/key.rs b/crates/aisix-etcd/src/key.rs new file mode 100644 index 00000000..cedfcd01 --- /dev/null +++ b/crates/aisix-etcd/src/key.rs @@ -0,0 +1,102 @@ +//! Parse etcd keys of the shape `{prefix}/{kind}/{id}`. +//! +//! Every aisix entity is stored at this canonical path. The watch supervisor +//! demultiplexes incoming events by the `kind` segment (`models`, `apikeys`, +//! `teams`, …) so each typed table can be updated independently. + +use std::fmt; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResourceKey<'a> { + pub kind: &'a str, + pub id: &'a str, +} + +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum KeyError { + #[error("etcd key {key:?} does not start with configured prefix {prefix:?}")] + PrefixMismatch { key: String, prefix: String }, + #[error("etcd key {0:?} is missing the `{{kind}}/{{id}}` suffix")] + MissingSuffix(String), + #[error("etcd key {0:?} has an empty kind or id segment")] + EmptySegment(String), +} + +/// Split an etcd key into (kind, id) given the configured aisix prefix. +/// +/// Example: with prefix `/aisix`, a key `/aisix/models/abc-123` parses to +/// `ResourceKey { kind: "models", id: "abc-123" }`. +pub fn parse<'a>(prefix: &str, key: &'a str) -> Result, KeyError> { + // Accept both `/aisix` and `/aisix/` prefixes transparently. + let trimmed_prefix = prefix.trim_end_matches('/'); + let rest = key + .strip_prefix(trimmed_prefix) + .ok_or_else(|| KeyError::PrefixMismatch { + key: key.to_string(), + prefix: prefix.to_string(), + })?; + let rest = rest.trim_start_matches('/'); + + let (kind, id) = rest + .split_once('/') + .ok_or_else(|| KeyError::MissingSuffix(key.to_string()))?; + + if kind.is_empty() || id.is_empty() { + return Err(KeyError::EmptySegment(key.to_string())); + } + + Ok(ResourceKey { kind, id }) +} + +impl fmt::Display for ResourceKey<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}/{}", self.kind, self.id) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn happy_path_parses_kind_and_id() { + let k = parse("/aisix", "/aisix/models/abc-123").unwrap(); + assert_eq!(k.kind, "models"); + assert_eq!(k.id, "abc-123"); + } + + #[test] + fn trailing_slash_in_prefix_is_tolerated() { + let k = parse("/aisix/", "/aisix/apikeys/uuid-1").unwrap(); + assert_eq!(k.kind, "apikeys"); + assert_eq!(k.id, "uuid-1"); + } + + #[test] + fn prefix_mismatch_is_detected() { + let err = parse("/aisix", "/other/models/a").unwrap_err(); + assert!(matches!(err, KeyError::PrefixMismatch { .. })); + } + + #[test] + fn missing_suffix_is_rejected() { + // Prefix-only key, no kind/id. + let err = parse("/aisix", "/aisix/models").unwrap_err(); + assert!(matches!(err, KeyError::MissingSuffix(_))); + } + + #[test] + fn empty_segments_are_rejected() { + let err = parse("/aisix", "/aisix/models/").unwrap_err(); + assert!(matches!(err, KeyError::EmptySegment(_))); + } + + #[test] + fn display_is_kind_slash_id() { + let k = ResourceKey { + kind: "models", + id: "abc", + }; + assert_eq!(k.to_string(), "models/abc"); + } +} diff --git a/crates/aisix-etcd/src/lib.rs b/crates/aisix-etcd/src/lib.rs index 69df5127..c93bc911 100644 --- a/crates/aisix-etcd/src/lib.rs +++ b/crates/aisix-etcd/src/lib.rs @@ -1,4 +1,37 @@ -//! aisix-etcd — etcd-backed ConfigProvider + watch supervisor. +//! aisix-etcd — etcd-backed [`ConfigProvider`] + watch supervisor. +//! +//! The gateway's hot read path talks to a `SnapshotHandle` +//! (see `aisix-core`). This crate is what *populates* that handle, running +//! a single supervisor task that: +//! +//! 1. Connects to etcd (5s × 5 retries on bootstrap — spec §2) +//! 2. Performs a full range read under the configured prefix +//! 3. Opens a watch stream from the next revision +//! 4. Applies Put / Delete events by copy-on-write replacing the snapshot +//! 5. Triggers a full resync on compaction +//! 6. Reconnects with exponential backoff (1→60s) on transport failure +//! +//! The [`ConfigProvider`] trait is the seam tests use to plug in an +//! in-memory provider and avoid a container dependency for unit testing. +//! +//! Spec references: §2 (config system), §3 (data models). #![forbid(unsafe_code)] #![deny(rust_2018_idioms)] + +pub mod backoff; +pub mod etcd_provider; +pub mod key; +pub mod loader; +pub mod provider; +pub mod supervisor; + +pub use backoff::{ExpBackoff, BASE_MS, MAX_MS}; +pub use etcd_provider::{ + ConnectPolicy, EtcdConfigProvider, EtcdWatchStream, CONNECT_MAX_ATTEMPTS, + CONNECT_RETRY_INTERVAL, +}; +pub use key::{parse as parse_key, KeyError, ResourceKey}; +pub use loader::{build_snapshot, BuildStats}; +pub use provider::{ConfigProvider, ProviderError, RawEntry, WatchEvent}; +pub use supervisor::Supervisor; diff --git a/crates/aisix-etcd/src/loader.rs b/crates/aisix-etcd/src/loader.rs new file mode 100644 index 00000000..452dd9b0 --- /dev/null +++ b/crates/aisix-etcd/src/loader.rs @@ -0,0 +1,220 @@ +//! Turn raw etcd entries into a typed [`AisixSnapshot`]. +//! +//! Flow: +//! 1. parse the key → `(kind, id)` +//! 2. validate the value against the kind's JSON Schema +//! 3. deserialise into the typed struct via serde (cheap after schema +//! passes) +//! 4. insert into the appropriate [`ResourceTable`] +//! +//! Malformed payloads are logged at WARN level and skipped, not fatal — +//! this matches spec §2: "the gateway does not abort on a single bad +//! entry; it serves the rest." + +use aisix_core::models::{validate_apikey, validate_model, ApiKey, Model, SchemaError}; +use aisix_core::resource::ResourceEntry; +use aisix_core::AisixSnapshot; +use serde::de::DeserializeOwned; +use serde_json::Value; + +use crate::key::{self, ResourceKey}; +use crate::provider::RawEntry; + +/// Counts of rejected entries during a build, useful for metrics. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct BuildStats { + pub accepted: usize, + pub schema_rejected: usize, + pub parse_rejected: usize, + pub unknown_kind: usize, + pub key_rejected: usize, +} + +/// Build a fresh snapshot from raw entries. Never fails — bad rows are +/// counted in [`BuildStats`] and skipped. The prefix lets us strip it +/// before key parsing. +pub fn build_snapshot(prefix: &str, entries: &[RawEntry]) -> (AisixSnapshot, BuildStats) { + let snapshot = AisixSnapshot::new(); + let mut stats = BuildStats::default(); + + for raw in entries { + let parsed = match key::parse(prefix, &raw.key) { + Ok(k) => k, + Err(err) => { + tracing::warn!(key = %raw.key, error = %err, "skipping etcd entry with bad key"); + stats.key_rejected += 1; + continue; + } + }; + + let value: Value = match serde_json::from_slice(&raw.value) { + Ok(v) => v, + Err(err) => { + tracing::warn!(key = %raw.key, error = %err, "skipping non-JSON etcd entry"); + stats.parse_rejected += 1; + continue; + } + }; + + match parsed.kind { + "models" => { + if let Some(entry) = validate_and_parse::( + &raw.key, + raw.revision, + parsed, + &value, + validate_model, + &mut stats, + ) { + snapshot.models.insert(entry); + } + } + "apikeys" => { + if let Some(entry) = validate_and_parse::( + &raw.key, + raw.revision, + parsed, + &value, + validate_apikey, + &mut stats, + ) { + snapshot.apikeys.insert(entry); + } + } + other => { + tracing::debug!(key = %raw.key, kind = %other, "unknown etcd kind; skipping"); + stats.unknown_kind += 1; + } + } + } + + (snapshot, stats) +} + +fn validate_and_parse( + key: &str, + revision: i64, + parsed: ResourceKey<'_>, + value: &Value, + validate: fn(&Value) -> Result<(), SchemaError>, + stats: &mut BuildStats, +) -> Option> +where + T: DeserializeOwned, +{ + if let Err(err) = validate(value) { + tracing::warn!(key = %key, error = %err, "schema validation failed; skipping"); + stats.schema_rejected += 1; + return None; + } + + match serde_json::from_value::(value.clone()) { + Ok(t) => { + stats.accepted += 1; + Some(ResourceEntry::new(parsed.id, t, revision)) + } + Err(err) => { + // Schema passed but serde refused — usually a deny_unknown_fields + // mismatch. Treat as schema-rejected for stats purposes. + tracing::warn!(key = %key, error = %err, "serde parse failed after schema pass"); + stats.parse_rejected += 1; + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn raw(key: &str, value: &[u8], rev: i64) -> RawEntry { + RawEntry { + key: key.into(), + value: value.to_vec(), + revision: rev, + } + } + + const VALID_MODEL: &[u8] = br#"{ + "name": "my-gpt4", + "model": "openai/gpt-4o", + "provider_config": {"api_key": "sk-x"} + }"#; + + const VALID_APIKEY: &[u8] = br#"{ + "key": "sk-abc", + "allowed_models": ["my-gpt4"] + }"#; + + #[test] + fn builds_snapshot_for_happy_entries() { + let entries = vec![ + raw("/aisix/models/m-1", VALID_MODEL, 2), + raw("/aisix/apikeys/k-1", VALID_APIKEY, 3), + ]; + let (snap, stats) = build_snapshot("/aisix", &entries); + + assert_eq!(stats.accepted, 2); + assert_eq!(snap.models.len(), 1); + assert_eq!(snap.apikeys.len(), 1); + assert_eq!(snap.models.get_by_name("my-gpt4").unwrap().id, "m-1"); + assert_eq!(snap.apikeys.get_by_name("sk-abc").unwrap().id, "k-1"); + } + + #[test] + fn malformed_json_is_skipped_not_fatal() { + let entries = vec![ + raw("/aisix/models/bad", b"not-json", 1), + raw("/aisix/models/good", VALID_MODEL, 2), + ]; + let (snap, stats) = build_snapshot("/aisix", &entries); + assert_eq!(stats.parse_rejected, 1); + assert_eq!(stats.accepted, 1); + assert_eq!(snap.models.len(), 1); + } + + #[test] + fn schema_failure_is_counted() { + let entries = vec![raw( + "/aisix/models/bad-provider", + br#"{"name":"x","model":"mistral/large","provider_config":{"api_key":"k"}}"#, + 1, + )]; + let (_snap, stats) = build_snapshot("/aisix", &entries); + assert_eq!(stats.schema_rejected, 1); + assert_eq!(stats.accepted, 0); + } + + #[test] + fn unknown_kinds_are_skipped() { + let entries = vec![raw("/aisix/teams/t-1", b"{}", 1)]; + let (snap, stats) = build_snapshot("/aisix", &entries); + assert_eq!(stats.unknown_kind, 1); + assert!(snap.models.is_empty()); + assert!(snap.apikeys.is_empty()); + } + + #[test] + fn bad_key_shape_is_counted_separately() { + let entries = vec![raw("/other/models/a", VALID_MODEL, 1)]; + let (_snap, stats) = build_snapshot("/aisix", &entries); + assert_eq!(stats.key_rejected, 1); + } + + #[test] + fn one_bad_entry_does_not_abort_the_batch() { + let entries = vec![ + raw("/aisix/models/m-1", VALID_MODEL, 1), + raw("/aisix/models/bad", b"not-json", 2), + raw("/aisix/models/m-2", VALID_MODEL, 3), // same name -> update in place + raw("/aisix/apikeys/k-1", VALID_APIKEY, 4), + ]; + let (snap, stats) = build_snapshot("/aisix", &entries); + assert_eq!(stats.accepted, 3); + assert_eq!(stats.parse_rejected, 1); + // m-1 and m-2 share the same name; the second insert rebinds the + // name to m-2, but both id entries are present in the table. + assert_eq!(snap.models.len(), 2); + assert_eq!(snap.apikeys.len(), 1); + } +} diff --git a/crates/aisix-etcd/src/provider.rs b/crates/aisix-etcd/src/provider.rs new file mode 100644 index 00000000..03084601 --- /dev/null +++ b/crates/aisix-etcd/src/provider.rs @@ -0,0 +1,99 @@ +//! The [`ConfigProvider`] abstraction the supervisor runs against. +//! +//! The real implementation is etcd (see [`crate::etcd_provider`]). Tests +//! plug in an in-memory provider so the supervisor can be exercised +//! deterministically without a container. +//! +//! Decoupling the supervisor from the concrete client also means the +//! future Admin API can write through the same trait on the happy path. + +use async_trait::async_trait; +use std::sync::Arc; + +/// Raw (key, value, revision) triple as returned by etcd ranges / watches. +/// Values are `serde_json::Value` so callers can run schema validation +/// before typed deserialisation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RawEntry { + pub key: String, + pub value: Vec, + pub revision: i64, +} + +/// Events surfaced to the watch consumer. `Resync` is emitted when the +/// supervisor has detected compaction or a reconnect, forcing a full +/// snapshot rebuild rather than delta application. +#[derive(Debug, Clone)] +pub enum WatchEvent { + Put(RawEntry), + Delete { + key: String, + revision: i64, + }, + /// Full reload: supervisor has reloaded all entries under the prefix + /// and is handing the whole set to the consumer in one atomic batch. + Resync { + entries: Arc>, + revision: i64, + }, +} + +#[derive(Debug, thiserror::Error)] +pub enum ProviderError { + #[error("etcd connection failed: {0}")] + Connect(String), + #[error("etcd range request failed: {0}")] + Range(String), + #[error("etcd watch stream failed: {0}")] + Watch(String), + #[error("etcd revision was compacted — caller should resync")] + Compacted, +} + +/// Abstraction the supervisor depends on. Methods are async so the etcd +/// implementation can perform its gRPC I/O; test doubles can use channels. +#[async_trait] +pub trait ConfigProvider: Send + Sync + 'static { + /// Full range read under the configured prefix. Returns the current + /// entries plus the etcd revision at which the read was consistent; + /// the supervisor starts its watch from `revision + 1`. + async fn load_all(&self) -> Result<(Vec, i64), ProviderError>; + + /// Open a watch stream starting from `start_revision`. The stream's + /// items are individual events; `Resync` is *not* emitted on this + /// channel — the supervisor is responsible for detecting compaction + /// (via [`ProviderError::Compacted`]) and triggering a fresh + /// `load_all` + `Resync` dispatch itself. + async fn watch( + &self, + start_revision: i64, + ) -> Result< + Box> + Send + Unpin>, + ProviderError, + >; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn raw_entry_is_cheap_to_clone() { + let e = RawEntry { + key: "/aisix/models/a".into(), + value: b"{}".to_vec(), + revision: 1, + }; + let c = e.clone(); + assert_eq!(e, c); + } + + #[test] + fn provider_error_compacted_is_distinct() { + let err = ProviderError::Compacted; + assert_eq!( + err.to_string(), + "etcd revision was compacted — caller should resync" + ); + } +} diff --git a/crates/aisix-etcd/src/supervisor.rs b/crates/aisix-etcd/src/supervisor.rs new file mode 100644 index 00000000..32afea45 --- /dev/null +++ b/crates/aisix-etcd/src/supervisor.rs @@ -0,0 +1,395 @@ +//! Watch supervisor — the single long-running task that owns the +//! [`ConfigProvider`] and keeps an [`AisixSnapshot`] current in a +//! [`SnapshotHandle`]. +//! +//! Responsibilities (spec §2): +//! 1. Initial `load_all` + publish first snapshot +//! 2. Open a watch stream from the load revision +//! 3. Apply Put/Delete events incrementally on top of the current +//! snapshot (building a *new* snapshot each time so reads stay +//! lock-free) +//! 4. On compaction or stream error, full-reload + resync +//! 5. Reconnect with exponential backoff (1→60s) on transport failure +//! +//! The apply step is *copy-on-write* per batch: we clone the current +//! snapshot into a new one, mutate, and `store` it. That keeps the +//! read path reading a fully-formed `Arc` the whole time. + +use aisix_core::snapshot::SnapshotHandle; +use aisix_core::AisixSnapshot; +use futures::StreamExt; +use std::sync::Arc; +use std::time::Duration; + +use crate::backoff::ExpBackoff; +use crate::key; +use crate::loader::{self, BuildStats}; +use crate::provider::{ConfigProvider, ProviderError, RawEntry, WatchEvent}; + +/// One supervisor instance. Consumers call [`Supervisor::run`] once and +/// drop the returned handle on shutdown. +pub struct Supervisor { + provider: Arc

, + prefix: String, + handle: SnapshotHandle, +} + +impl Supervisor

{ + pub fn new(provider: Arc

, prefix: impl Into) -> Self { + Self { + provider, + prefix: prefix.into(), + handle: SnapshotHandle::new(AisixSnapshot::new()), + } + } + + /// Clone of the public snapshot handle. Axum state / request handlers + /// hold this; calls to `.load()` are cheap atomic reads. + pub fn handle(&self) -> SnapshotHandle { + self.handle.clone() + } + + /// Run one full reload + watch cycle and publish the resulting + /// snapshot. Returns the stats from the build for observability. + /// Stops after the first watch error — the outer [`Self::run`] loop + /// decides whether to backoff and retry. + pub async fn load_once(&self) -> Result { + let (entries, revision) = self.provider.load_all().await?; + let (snapshot, stats) = loader::build_snapshot(&self.prefix, &entries); + tracing::info!( + accepted = stats.accepted, + rejected = stats.schema_rejected + stats.parse_rejected, + revision, + "initial snapshot built", + ); + self.handle.store(snapshot); + Ok(stats) + } + + /// Apply a single Put event on top of the current snapshot. + /// Returns `true` if the apply succeeded (schema + parse passed). + pub fn apply_put(&self, entry: &RawEntry) -> bool { + // Build a tiny snapshot out of just the new entry, then merge. + let (tiny, stats) = loader::build_snapshot(&self.prefix, std::slice::from_ref(entry)); + if stats.accepted == 0 { + return false; + } + + let new = clone_snapshot(&self.handle.load()); + + // Move any entries from `tiny` into `new`. + for e in tiny.models.entries() { + new.models.insert(clone_entry(&e)); + } + for e in tiny.apikeys.entries() { + new.apikeys.insert(clone_entry(&e)); + } + + self.handle.store(new); + true + } + + /// Apply a Delete event. Returns `true` if anything was actually + /// removed (the kind/id was present). + pub fn apply_delete(&self, key_str: &str) -> bool { + let parsed = match key::parse(&self.prefix, key_str) { + Ok(k) => k, + Err(err) => { + tracing::warn!(key = %key_str, error = %err, "ignoring delete with bad key"); + return false; + } + }; + + let new = clone_snapshot(&self.handle.load()); + let removed = match parsed.kind { + "models" => new.models.remove(parsed.id).is_some(), + "apikeys" => new.apikeys.remove(parsed.id).is_some(), + _ => false, + }; + if removed { + self.handle.store(new); + } + removed + } + + /// Replace the current snapshot with a freshly loaded set (resync). + pub fn apply_resync(&self, entries: &[RawEntry]) -> BuildStats { + let (snap, stats) = loader::build_snapshot(&self.prefix, entries); + self.handle.store(snap); + stats + } + + /// Long-running loop. Handles exp-backoff reconnects and resync on + /// compaction. Runs until cancelled via the cancellation token. + pub async fn run(self: Arc, mut cancel: tokio::sync::watch::Receiver) { + let mut backoff = ExpBackoff::default(); + loop { + if *cancel.borrow() { + return; + } + + match self.cycle(&cancel).await { + Ok(()) => { + // Graceful stream end (compaction or server-initiated + // close). Reset backoff, but still yield a short + // interval before reconnecting so we never spin. + backoff.reset(); + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis(100)) => {} + _ = cancel.changed() => { + if *cancel.borrow() { return; } + } + } + } + Err(SupervisorError::Cancelled) => return, + Err(SupervisorError::Provider(err)) => { + let delay = backoff.next_delay(); + tracing::warn!( + error = %err, + backoff_ms = delay.as_millis() as u64, + "etcd watch failed; backing off before reconnect", + ); + tokio::select! { + _ = tokio::time::sleep(delay) => {} + _ = cancel.changed() => { + if *cancel.borrow() { return; } + } + } + } + } + } + } + + /// One attempt at load + watch. Any error returns without retrying — + /// [`Self::run`] owns the backoff loop. + async fn cycle( + &self, + cancel: &tokio::sync::watch::Receiver, + ) -> Result<(), SupervisorError> { + let (entries, revision) = self + .provider + .load_all() + .await + .map_err(SupervisorError::Provider)?; + + self.apply_resync(&entries); + + let mut stream = self + .provider + .watch(revision + 1) + .await + .map_err(SupervisorError::Provider)?; + + loop { + if *cancel.borrow() { + return Err(SupervisorError::Cancelled); + } + + let next = tokio::select! { + item = stream.next() => item, + _ = wait_for_cancel(cancel.clone()) => { + return Err(SupervisorError::Cancelled); + } + }; + + match next { + None => return Ok(()), + Some(Err(ProviderError::Compacted)) => { + tracing::warn!("etcd compaction detected — resyncing"); + // Break out so `run` re-enters `cycle` cleanly; the + // next iteration re-loads from scratch. We don't want + // to treat compaction as a backoff-worthy failure. + return Ok(()); + } + Some(Err(err)) => return Err(SupervisorError::Provider(err)), + Some(Ok(WatchEvent::Put(raw))) => { + self.apply_put(&raw); + } + Some(Ok(WatchEvent::Delete { key, .. })) => { + self.apply_delete(&key); + } + Some(Ok(WatchEvent::Resync { entries, .. })) => { + self.apply_resync(&entries); + } + } + } + } +} + +#[derive(Debug)] +enum SupervisorError { + Cancelled, + Provider(ProviderError), +} + +async fn wait_for_cancel(mut rx: tokio::sync::watch::Receiver) { + loop { + if *rx.borrow() { + return; + } + if rx.changed().await.is_err() { + // Sender dropped: treat as cancellation. + return; + } + } +} + +/// Shallow clone of every [`Arc`] — fast and, importantly, +/// it doesn't materialise a deep copy of the `T` payload. +fn clone_snapshot(src: &AisixSnapshot) -> AisixSnapshot { + let out = AisixSnapshot::new(); + for e in src.models.entries() { + out.models.insert(clone_entry(&e)); + } + for e in src.apikeys.entries() { + out.apikeys.insert(clone_entry(&e)); + } + out +} + +fn clone_entry(src: &Arc>) -> aisix_core::ResourceEntry { + aisix_core::ResourceEntry { + id: src.id.clone(), + value: src.value.clone(), + revision: src.revision, + } +} + +/// Total time the supervisor will wait across its full 1→60s backoff +/// ladder before saturating. Exposed as a constant for tests and docs. +pub const BACKOFF_SATURATE_AFTER: Duration = Duration::from_secs(63); + +#[cfg(test)] +mod tests { + use super::*; + use crate::provider::{RawEntry, WatchEvent}; + use async_trait::async_trait; + use futures::stream; + use std::sync::Mutex; + + struct FakeProvider { + entries: Mutex>, + revision: i64, + events: Mutex>>, + } + + impl FakeProvider { + fn new(entries: Vec, revision: i64) -> Self { + Self { + entries: Mutex::new(entries), + revision, + events: Mutex::new(Vec::new()), + } + } + + fn with_events(mut self, events: Vec>) -> Self { + self.events = Mutex::new(events); + self + } + } + + #[async_trait] + impl ConfigProvider for FakeProvider { + async fn load_all(&self) -> Result<(Vec, i64), ProviderError> { + Ok((self.entries.lock().unwrap().clone(), self.revision)) + } + + async fn watch( + &self, + _start_revision: i64, + ) -> Result< + Box> + Send + Unpin>, + ProviderError, + > { + let events: Vec<_> = self.events.lock().unwrap().drain(..).collect(); + Ok(Box::new(stream::iter(events))) + } + } + + const VALID_MODEL: &[u8] = br#"{ + "name": "my-gpt4", + "model": "openai/gpt-4o", + "provider_config": {"api_key": "sk-x"} + }"#; + + fn entry(key: &str, v: &[u8], rev: i64) -> RawEntry { + RawEntry { + key: key.into(), + value: v.to_vec(), + revision: rev, + } + } + + #[tokio::test] + async fn load_once_publishes_initial_snapshot() { + let provider = Arc::new(FakeProvider::new( + vec![entry("/aisix/models/m-1", VALID_MODEL, 1)], + 5, + )); + let sup = Supervisor::new(provider, "/aisix"); + let stats = sup.load_once().await.unwrap(); + assert_eq!(stats.accepted, 1); + let snap = sup.handle().load(); + assert_eq!(snap.models.len(), 1); + } + + #[tokio::test] + async fn apply_put_adds_to_snapshot() { + let provider = Arc::new(FakeProvider::new(vec![], 0)); + let sup = Supervisor::new(provider, "/aisix"); + sup.load_once().await.unwrap(); + assert!(sup.apply_put(&entry("/aisix/models/m-1", VALID_MODEL, 2))); + assert_eq!(sup.handle().load().models.len(), 1); + } + + #[tokio::test] + async fn apply_put_rejects_bad_payload_without_mutating() { + let provider = Arc::new(FakeProvider::new(vec![], 0)); + let sup = Supervisor::new(provider, "/aisix"); + sup.load_once().await.unwrap(); + assert!(!sup.apply_put(&entry("/aisix/models/bad", b"not-json", 1))); + assert!(sup.handle().load().models.is_empty()); + } + + #[tokio::test] + async fn apply_delete_removes_entry() { + let provider = Arc::new(FakeProvider::new( + vec![entry("/aisix/models/m-1", VALID_MODEL, 1)], + 1, + )); + let sup = Supervisor::new(provider, "/aisix"); + sup.load_once().await.unwrap(); + assert!(sup.apply_delete("/aisix/models/m-1")); + assert!(sup.handle().load().models.is_empty()); + } + + #[tokio::test] + async fn apply_resync_replaces_snapshot() { + let provider = Arc::new(FakeProvider::new(vec![], 0)); + let sup = Supervisor::new(provider, "/aisix"); + sup.load_once().await.unwrap(); + sup.apply_resync(&[entry("/aisix/models/m-1", VALID_MODEL, 1)]); + assert_eq!(sup.handle().load().models.len(), 1); + } + + #[tokio::test] + async fn run_loop_applies_put_then_exits_on_cancel() { + let provider = Arc::new(FakeProvider::new(vec![], 0).with_events(vec![Ok( + WatchEvent::Put(entry("/aisix/models/m-1", VALID_MODEL, 2)), + )])); + let sup = Arc::new(Supervisor::new(provider, "/aisix")); + let handle = sup.handle(); + let (tx, rx) = tokio::sync::watch::channel(false); + + let join = tokio::spawn(sup.clone().run(rx)); + + // Let the supervisor drain its finite event stream and reach the + // "stream ended" branch. The load + event apply both happen + // synchronously relative to the event stream being in-memory. + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert_eq!(handle.load().models.len(), 1); + + tx.send(true).unwrap(); + join.await.unwrap(); + } +}