Skip to main content

brain_edge/state/
mod.rs

1//! Shared state: the [`BrainPool`] connection cache to Brain plus the injected
2//! [`CredentialResolver`] + [`MeteringSink`] ports.
3//!
4//! The resolver turns an incoming request into a Brain wire credential; the
5//! [`BrainPool`] reuses one handshake per credential so we don't re-connect on
6//! every request; the metering sink observes each op's outcome. Self-host uses
7//! the bearer/no-op defaults ([`EdgeState::new`]); the gateway injects its own
8//! ports ([`EdgeState::with_ports`]).
9
10use std::sync::Arc;
11use std::time::Duration;
12
13use axum::http::HeaderMap;
14use brain_db_sdk::BrainClient;
15
16use crate::config::EdgeConfig;
17use crate::error::ApiError;
18use crate::pool::{BrainPool, BrainPoolConfig};
19use crate::port::{
20    BearerResolver, CredentialResolver, MeterEvent, MeteringSink, NoopMeter, Outcome,
21    ResolvedCredential,
22};
23
24/// Cheap-to-clone application state.
25#[derive(Clone)]
26pub struct EdgeState {
27    inner: Arc<Inner>,
28}
29
30impl std::fmt::Debug for EdgeState {
31    /// The ports are trait objects and the pool holds live connections, so this
32    /// prints the shape rather than the contents.
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        f.debug_struct("EdgeState")
35            .field("pool", &self.inner.pool)
36            .finish_non_exhaustive()
37    }
38}
39
40struct Inner {
41    pool: BrainPool,
42    resolver: Arc<dyn CredentialResolver>,
43    meter: Arc<dyn MeteringSink>,
44}
45
46impl EdgeState {
47    /// The self-host default: bearer passthrough, no metering.
48    #[must_use]
49    pub fn new(config: EdgeConfig) -> Self {
50        Self::with_ports(config, Arc::new(BearerResolver), Arc::new(NoopMeter))
51    }
52
53    /// Build state with an injected resolver + metering sink (the hosted gateway
54    /// path). The gateway passes an API-key resolver and its analytics sink.
55    #[must_use]
56    pub fn with_ports(
57        config: EdgeConfig,
58        resolver: Arc<dyn CredentialResolver>,
59        meter: Arc<dyn MeteringSink>,
60    ) -> Self {
61        let pool = BrainPool::new(BrainPoolConfig {
62            brain_addr: config.brain_addr,
63            pool_size: config.pool_size,
64            max_credentials: config.max_credentials,
65            idle_ttl: Duration::from_secs(config.idle_ttl_secs),
66        });
67        Self {
68            inner: Arc::new(Inner {
69                pool,
70                resolver,
71                meter,
72            }),
73        }
74    }
75
76    /// Resolve a request to a Brain credential via the injected resolver.
77    ///
78    /// # Errors
79    /// Propagates the resolver's `401`/`403` when no valid credential is present.
80    pub async fn resolve(&self, headers: &HeaderMap) -> Result<ResolvedCredential, ApiError> {
81        self.inner.resolver.resolve(headers).await
82    }
83
84    /// Record one op's outcome via the injected sink (fire-and-forget).
85    pub fn record(&self, op: &str, ident: &ResolvedCredential, outcome: Outcome) {
86        self.inner.meter.record(&MeterEvent {
87            tenant: ident.tenant.as_deref(),
88            op,
89            outcome,
90        });
91    }
92
93    /// A Brain client authenticated as `credential`, from the [`BrainPool`]'s
94    /// cached (or freshly opened) per-credential pool.
95    ///
96    /// # Errors
97    /// Returns an engine error if the pool fails to open a connection.
98    pub async fn client_for(&self, credential: &str) -> Result<Arc<BrainClient>, ApiError> {
99        self.inner
100            .pool
101            .client_for(credential)
102            .await
103            .map_err(|e| ApiError::from_brain(&e))
104    }
105
106    /// The shared Brain connection pool (for the background idle sweeper and
107    /// readiness probing).
108    #[must_use]
109    pub fn pool(&self) -> &BrainPool {
110        &self.inner.pool
111    }
112
113    /// The injected metering sink, shared. The wire proxy reuses it so wire-path
114    /// ops meter through the same sink as the HTTP data plane.
115    #[must_use]
116    pub fn meter(&self) -> Arc<dyn MeteringSink> {
117        Arc::clone(&self.inner.meter)
118    }
119}