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        #[cfg(feature = "storage-cogs")]
68        let streams = config
69            .storage_cogs
70            .as_ref()
71            .map(ChangeStreamFactory::new)
72            .unwrap_or_default();
73        #[cfg(not(feature = "storage-cogs"))]
74        let streams = ChangeStreamFactory::default();
75        let cipher = match config.service.cipher()? {
76            Some(cipher) => cipher,
77            None => {
78                let cipher = Cipher::ephemeral().context("failed to initialize encryption")?;
79                objectstore_log::warn!(
80                    "No persistent encryption keyring configured; using a process-local ephemeral key. Resumable uploads cannot reliably cross replicas or survive restarts. Configure service.encryption in production."
81                );
82                cipher
83            }
84        };
85        let backend = backend::from_config(config.storage.clone(), &streams).await?;
86        let concurrency = ConcurrencyLimiter::new(config.service.max_concurrency)
87            .with_queue(config.service.concurrency_queue)
88            .with_timeout(config.service.concurrency_timeout)
89            .with_bulk(config.service.bulk_concurrency_pct);
90        let mut service = StorageService::new(backend, cipher)
91            .with_concurrency(concurrency)
92            .with_background_queue_limit(config.service.background_queue);
93        service.start();
94
95        let key_directory = Arc::new(PublicKeyDirectory::from_config(&config.auth).await?);
96        if config.auth.enforce && key_directory.keys.is_empty() {
97            anyhow::bail!(
98                "Auth enforcement is enabled but no keys are configured. Either disable auth enforcement (dev/test environments) or configure a public key."
99            );
100        }
101        let rate_limiter = RateLimiter::new(config.rate_limits.clone());
102        rate_limiter.start();
103
104        let request_counter = RequestCounter::new(config.http.max_requests);
105        tokio::spawn(request_counter.clone().run_emitter());
106
107        Ok(Arc::new(Self {
108            config,
109            service,
110            key_directory,
111            rate_limiter,
112            request_counter,
113        }))
114    }
115
116    /// Wraps a byte stream with bandwidth metering for rate limiting.
117    ///
118    /// Works with any stream type — use for both [`objectstore_service::PayloadStream`]
119    /// (outgoing, `E = io::Error`) and [`objectstore_service::ClientStream`]
120    /// (incoming, `E = ClientError`).
121    pub(crate) fn meter_stream<S, E>(
122        &self,
123        stream: S,
124        context: &ObjectContext,
125    ) -> MeteredPayloadStream<S>
126    where
127        S: Stream<Item = Result<Bytes, E>> + Send + 'static,
128    {
129        MeteredPayloadStream::new(stream, self.rate_limiter.bandwidth_handle(context))
130    }
131
132    /// Records bandwidth usage for the given context without wrapping a stream.
133    ///
134    /// Used for cases where the payload size is known upfront (e.g. batch INSERT).
135    pub fn record_bandwidth(&self, context: &ObjectContext, bytes: u64) {
136        self.rate_limiter.record_bandwidth(context, bytes);
137    }
138}
139
140/// Periodically captures and reports jemalloc stats.
141#[cfg(target_os = "linux")]
142async fn track_allocator_metrics(interval: Duration) {
143    // INVARIANT: MIB resolution only fails if jemalloc is not the active allocator,
144    // which would be a misconfigured build. Panic early to surface the problem.
145    let epoch = tikv_jemalloc_ctl::epoch::mib().expect("jemalloc epoch MIB");
146    let allocated = tikv_jemalloc_ctl::stats::allocated::mib().expect("jemalloc allocated MIB");
147    let active = tikv_jemalloc_ctl::stats::active::mib().expect("jemalloc active MIB");
148    let resident = tikv_jemalloc_ctl::stats::resident::mib().expect("jemalloc resident MIB");
149    let mapped = tikv_jemalloc_ctl::stats::mapped::mib().expect("jemalloc mapped MIB");
150
151    let mut ticker = tokio::time::interval(interval);
152    loop {
153        ticker.tick().await;
154
155        let Ok(_) = epoch.advance() else {
156            continue;
157        };
158
159        if let Ok(allocated_bytes) = allocated.read() {
160            // Bytes currently allocated by the application.
161            objectstore_metrics::gauge!("jemalloc.allocated" = allocated_bytes);
162        }
163        if let Ok(active_bytes) = active.read() {
164            // Bytes in active jemalloc pages (≥ allocated).
165            objectstore_metrics::gauge!("jemalloc.active" = active_bytes);
166        }
167        if let Ok(resident_bytes) = resident.read() {
168            // Bytes in resident pages mapped from the OS (≥ active).
169            objectstore_metrics::gauge!("jemalloc.resident" = resident_bytes);
170        }
171        if let Ok(mapped_bytes) = mapped.read() {
172            // Bytes in chunks mapped from the OS (≥ resident).
173            objectstore_metrics::gauge!("jemalloc.mapped" = mapped_bytes);
174        }
175    }
176}
177
178/// Periodically captures and reports internal Tokio runtime metrics.
179async fn track_runtime_metrics(interval: Duration) {
180    let mut ticker = tokio::time::interval(interval);
181    let metrics = Handle::current().metrics();
182
183    loop {
184        ticker.tick().await;
185        objectstore_log::trace!("Capturing runtime metrics");
186
187        objectstore_metrics::gauge!("runtime.num_workers" = metrics.num_workers());
188        objectstore_metrics::gauge!("runtime.num_alive_tasks" = metrics.num_alive_tasks());
189        objectstore_metrics::gauge!("runtime.global_queue_depth" = metrics.global_queue_depth());
190        objectstore_metrics::gauge!(
191            "runtime.num_blocking_threads" = metrics.num_blocking_threads()
192        );
193        objectstore_metrics::gauge!(
194            "runtime.num_idle_blocking_threads" = metrics.num_idle_blocking_threads()
195        );
196        objectstore_metrics::gauge!(
197            "runtime.blocking_queue_depth" = metrics.blocking_queue_depth()
198        );
199
200        let registered_fds = metrics.io_driver_fd_registered_count();
201        let deregistered_fds = metrics.io_driver_fd_deregistered_count();
202        objectstore_metrics::gauge!(
203            "runtime.num_io_driver_fds" = registered_fds - deregistered_fds
204        );
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[tokio::test]
213    async fn enforce_without_keys_fails_startup() {
214        let config = Config {
215            auth: crate::config::AuthZ {
216                enforce: true,
217                ..Default::default()
218            },
219            ..Default::default()
220        };
221        let err = Services::spawn(config).await.unwrap_err();
222        assert!(
223            err.to_string()
224                .contains("Auth enforcement is enabled but no keys are configured"),
225        );
226    }
227
228    #[tokio::test]
229    async fn no_enforce_without_keys_starts_ok() {
230        let config = Config {
231            auth: crate::config::AuthZ {
232                enforce: false,
233                ..Default::default()
234            },
235            ..Default::default()
236        };
237        assert!(Services::spawn(config).await.is_ok());
238    }
239}
objectstore_server/state.rs

objectstore_server/
state.rs
