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

objectstore_server/
state.rs
