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

objectstore_server/
state.rs
