objectstore_server/
state.rs

1//! Shared server state passed to all HTTP request handlers.
2//!
3//! [`Services`] is constructed once during startup by [`Services::spawn`] and then shared
4//! across all request handlers as [`ServiceState`] (an `Arc<Services>`).
5
6use std::sync::Arc;
7use std::time::Duration;
8
9use anyhow::{Context, Result};
10use bytes::Bytes;
11use futures_util::Stream;
12use objectstore_service::change_stream::ChangeStreamFactory;
13use objectstore_service::concurrency::ConcurrencyLimiter;
14use objectstore_service::encryption::Cipher;
15use objectstore_service::id::ObjectContext;
16use objectstore_service::{StorageService, backend};
17use tokio::runtime::Handle;
18
19use crate::auth::PublicKeyDirectory;
20use crate::config::Config;
21use crate::rate_limits::{MeteredPayloadStream, RateLimiter};
22use crate::web::RequestCounter;
23
24/// Shared reference to the objectstore [`Services`].
25pub type ServiceState = Arc<Services>;
26
27/// Reference to the objectstore business logic.
28///
29/// This structure is created during server startup and shared with all HTTP request handlers. It
30/// can be used to access the configured storage backends and other shared resources.
31///
32/// In request handlers, use `axum::extract::State<ServiceState>` to retrieve a shared reference to
33/// this structure.
34#[derive(Debug)]
35pub struct Services {
36    /// The server configuration.
37    pub config: Config,
38    /// Raw handle to the underlying storage service that does not enforce authorization checks.
39    ///
40    /// Consider using [`crate::auth::AuthAwareService`] for auth-checked access.
41    pub service: StorageService,
42    /// Directory for EdDSA public keys.
43    ///
44    /// The `kid` header field from incoming authorization tokens should correspond to a public key
45    /// in this directory that can be used to verify the token.
46    pub key_directory: Arc<PublicKeyDirectory>,
47    /// Stateful admission-based rate limiter for incoming requests.
48    pub rate_limiter: RateLimiter,
49    /// In-flight HTTP request counter with the configured limit.
50    ///
51    /// Shared with the web layer so the concurrency-limit middleware, the tracking
52    /// layer, and any endpoint that reads the count all see the same atomic.
53    pub request_counter: RequestCounter,
54}
55
56impl Services {
57    /// Spawns all services and background tasks for objectstore.
58    ///
59    /// This returns a [`ServiceState`], which is a shared reference to the services suitable for
60    /// use in the web server.
61    pub async fn spawn(config: Config) -> Result<ServiceState> {
62        tokio::spawn(track_runtime_metrics(config.runtime.metrics_interval));
63        #[cfg(target_os = "linux")]
64        tokio::spawn(track_allocator_metrics(config.runtime.metrics_interval));
65
66        // Absent config means no transport, and therefore no reporting.
67        let streams = config
68            .storage_cogs
69            .as_ref()
70            .map(ChangeStreamFactory::new)
71            .unwrap_or_default();
72        let cipher = match config.service.cipher()? {
73            Some(cipher) => cipher,
74            None => Cipher::ephemeral().context("failed to initialize encryption")?,
75        };
76        let backend = backend::from_config(config.storage.clone(), &streams).await?;
77        let concurrency = ConcurrencyLimiter::new(config.service.max_concurrency)
78            .with_queue(config.service.concurrency_queue)
79            .with_timeout(config.service.concurrency_timeout)
80            .with_bulk(config.service.bulk_concurrency_pct);
81        let mut service = StorageService::new(backend, cipher)
82            .with_concurrency(concurrency)
83            .with_background_queue_limit(config.service.background_queue);
84        service.start();
85
86        let key_directory = Arc::new(PublicKeyDirectory::from_config(&config.auth).await?);
87        if config.auth.enforce && key_directory.keys.is_empty() {
88            anyhow::bail!(
89                "Auth enforcement is enabled but no keys are configured. Either disable auth enforcement (dev/test environments) or configure a public key."
90            );
91        }
92        let rate_limiter = RateLimiter::new(config.rate_limits.clone());
93        rate_limiter.start();
94
95        let request_counter = RequestCounter::new(config.http.max_requests);
96        tokio::spawn(request_counter.clone().run_emitter());
97
98        Ok(Arc::new(Self {
99            config,
100            service,
101            key_directory,
102            rate_limiter,
103            request_counter,
104        }))
105    }
106
107    /// Wraps a byte stream with bandwidth metering for rate limiting.
108    ///
109    /// Works with any stream type — use for both [`objectstore_service::PayloadStream`]
110    /// (outgoing, `E = io::Error`) and [`objectstore_service::ClientStream`]
111    /// (incoming, `E = ClientError`).
112    pub(crate) fn meter_stream<S, E>(
113        &self,
114        stream: S,
115        context: &ObjectContext,
116    ) -> MeteredPayloadStream<S>
117    where
118        S: Stream<Item = Result<Bytes, E>> + Send + 'static,
119    {
120        MeteredPayloadStream::new(stream, self.rate_limiter.bandwidth_handle(context))
121    }
122
123    /// Records bandwidth usage for the given context without wrapping a stream.
124    ///
125    /// Used for cases where the payload size is known upfront (e.g. batch INSERT).
126    pub fn record_bandwidth(&self, context: &ObjectContext, bytes: u64) {
127        self.rate_limiter.record_bandwidth(context, bytes);
128    }
129}
130
131/// Periodically captures and reports jemalloc stats.
132#[cfg(target_os = "linux")]
133async fn track_allocator_metrics(interval: Duration) {
134    // INVARIANT: MIB resolution only fails if jemalloc is not the active allocator,
135    // which would be a misconfigured build. Panic early to surface the problem.
136    let epoch = tikv_jemalloc_ctl::epoch::mib().expect("jemalloc epoch MIB");
137    let allocated = tikv_jemalloc_ctl::stats::allocated::mib().expect("jemalloc allocated MIB");
138    let active = tikv_jemalloc_ctl::stats::active::mib().expect("jemalloc active MIB");
139    let resident = tikv_jemalloc_ctl::stats::resident::mib().expect("jemalloc resident MIB");
140    let mapped = tikv_jemalloc_ctl::stats::mapped::mib().expect("jemalloc mapped MIB");
141
142    let mut ticker = tokio::time::interval(interval);
143    loop {
144        ticker.tick().await;
145
146        let Ok(_) = epoch.advance() else {
147            continue;
148        };
149
150        if let Ok(allocated_bytes) = allocated.read() {
151            // Bytes currently allocated by the application.
152            objectstore_metrics::gauge!("jemalloc.allocated" = allocated_bytes);
153        }
154        if let Ok(active_bytes) = active.read() {
155            // Bytes in active jemalloc pages (≥ allocated).
156            objectstore_metrics::gauge!("jemalloc.active" = active_bytes);
157        }
158        if let Ok(resident_bytes) = resident.read() {
159            // Bytes in resident pages mapped from the OS (≥ active).
160            objectstore_metrics::gauge!("jemalloc.resident" = resident_bytes);
161        }
162        if let Ok(mapped_bytes) = mapped.read() {
163            // Bytes in chunks mapped from the OS (≥ resident).
164            objectstore_metrics::gauge!("jemalloc.mapped" = mapped_bytes);
165        }
166    }
167}
168
169/// Periodically captures and reports internal Tokio runtime metrics.
170async fn track_runtime_metrics(interval: Duration) {
171    let mut ticker = tokio::time::interval(interval);
172    let metrics = Handle::current().metrics();
173
174    loop {
175        ticker.tick().await;
176        objectstore_log::trace!("Capturing runtime metrics");
177
178        objectstore_metrics::gauge!("runtime.num_workers" = metrics.num_workers());
179        objectstore_metrics::gauge!("runtime.num_alive_tasks" = metrics.num_alive_tasks());
180        objectstore_metrics::gauge!("runtime.global_queue_depth" = metrics.global_queue_depth());
181        objectstore_metrics::gauge!(
182            "runtime.num_blocking_threads" = metrics.num_blocking_threads()
183        );
184        objectstore_metrics::gauge!(
185            "runtime.num_idle_blocking_threads" = metrics.num_idle_blocking_threads()
186        );
187        objectstore_metrics::gauge!(
188            "runtime.blocking_queue_depth" = metrics.blocking_queue_depth()
189        );
190
191        let registered_fds = metrics.io_driver_fd_registered_count();
192        let deregistered_fds = metrics.io_driver_fd_deregistered_count();
193        objectstore_metrics::gauge!(
194            "runtime.num_io_driver_fds" = registered_fds - deregistered_fds
195        );
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[tokio::test]
204    async fn enforce_without_keys_fails_startup() {
205        let config = Config {
206            auth: crate::config::AuthZ {
207                enforce: true,
208                ..Default::default()
209            },
210            ..Default::default()
211        };
212        let err = Services::spawn(config).await.unwrap_err();
213        assert!(
214            err.to_string()
215                .contains("Auth enforcement is enabled but no keys are configured"),
216        );
217    }
218
219    #[tokio::test]
220    async fn no_enforce_without_keys_starts_ok() {
221        let config = Config {
222            auth: crate::config::AuthZ {
223                enforce: false,
224                ..Default::default()
225            },
226            ..Default::default()
227        };
228        assert!(Services::spawn(config).await.is_ok());
229    }
230}
objectstore_server/state.rs

objectstore_server/
state.rs
