objectstore_service/backend/
counting.rs

1//! Defines [`CountingBackend`], a decorator for the [`Backend`] trait that emits Cost of Goods Sold
2//! (COGS) compute usage metrics to the `cogs.usage` counter.
3//!
4//! [`CountingBackend`] is meant to wrap the outer-most [`Backend`] implementation owned by
5//! [`StorageService`] so that every tracked backend operation, whether a single-object operation
6//! called by [`StorageService`] or a batched operation streamed by [`StreamExecutor`], is counted
7//! once. Notably, any operation that fails before it gets to [`StorageService`] (e.g. an auth or
8//! rate limit failure at a higher layer) is not counted.
9//!
10//! For COGS purposes we use operation count as a proxy for compute cost under the assumption that
11//! each request we serve has a basically flat CPU cost. Large payloads take longer, but they can be
12//! streamed in the background while other requests are served so they don't really cost more.
13//!
14//! [`StorageService`]: crate::service::StorageService
15//! [`StreamExecutor`]: crate::streaming::StreamExecutor
16
17use std::num::NonZeroU64;
18
19use objectstore_types::metadata::Metadata;
20use objectstore_types::range::ByteRange;
21use objectstore_types::resumable::UploadProgress;
22use objectstore_types::time::Timestamp;
23use objectstore_types::usecase::Usecase;
24
25use crate::backend::common::{
26    Backend, DeleteResponse, ExpiryUpdate, GetResponse, MetadataResponse, PutResponse,
27    SetExpiryResponse,
28};
29use crate::error::Result;
30use crate::id::ObjectId;
31use crate::resumable::{BackendToken, Session};
32use crate::stream::ClientStream;
33
34/// Increments `cogs.usage` by one operation for the given `usecase`.
35///
36/// Under the hood, the `usecase` is used as the `app_feature`. This allows to identify distinct
37/// products and map them in the for the COGs pipeline.
38fn count(usecase: &Usecase) {
39    objectstore_metrics::count!("cogs.usage" += 1, app_feature = usecase.to_string());
40}
41
42/// A [`Backend`] decorator that counts each operation performed for COGS.
43///
44/// See the [module documentation](self) for how it should be used.
45#[derive(Debug)]
46pub struct CountingBackend {
47    inner: Box<dyn Backend>,
48}
49
50impl CountingBackend {
51    /// Creates a [`CountingBackend`] that wraps `inner` and increments `cogs.usage`
52    /// before delegating operations to it.
53    pub fn new(inner: Box<dyn Backend>) -> Self {
54        Self { inner }
55    }
56}
57
58#[async_trait::async_trait]
59impl Backend for CountingBackend {
60    fn name(&self) -> &'static str {
61        self.inner.name()
62    }
63
64    fn upload_granularity(&self) -> u64 {
65        self.inner.upload_granularity()
66    }
67
68    async fn put_object(
69        &self,
70        id: &ObjectId,
71        metadata: &Metadata,
72        stream: ClientStream,
73        access_time: Timestamp,
74    ) -> Result<PutResponse> {
75        count(id.usecase());
76        self.inner
77            .put_object(id, metadata, stream, access_time)
78            .await
79    }
80
81    async fn get_object(
82        &self,
83        id: &ObjectId,
84        access_time: Timestamp,
85        range: Option<ByteRange>,
86    ) -> Result<GetResponse> {
87        count(id.usecase());
88        self.inner.get_object(id, access_time, range).await
89    }
90
91    async fn get_metadata(
92        &self,
93        id: &ObjectId,
94        access_time: Timestamp,
95    ) -> Result<MetadataResponse> {
96        count(id.usecase());
97        self.inner.get_metadata(id, access_time).await
98    }
99
100    async fn set_expiry(
101        &self,
102        id: &ObjectId,
103        target: ExpiryUpdate,
104        access_time: Timestamp,
105    ) -> Result<SetExpiryResponse> {
106        count(id.usecase());
107        self.inner.set_expiry(id, target, access_time).await
108    }
109
110    async fn delete_object(&self, id: &ObjectId, access_time: Timestamp) -> Result<DeleteResponse> {
111        count(id.usecase());
112        self.inner.delete_object(id, access_time).await
113    }
114
115    async fn join(&self) {
116        self.inner.join().await;
117    }
118
119    async fn create_upload_session(
120        &self,
121        id: &ObjectId,
122        metadata: &Metadata,
123        upload_length: NonZeroU64,
124    ) -> Result<Option<BackendToken>> {
125        count(id.usecase());
126        self.inner
127            .create_upload_session(id, metadata, upload_length)
128            .await
129    }
130
131    async fn put_chunk(
132        &self,
133        session: &Session,
134        offset: u64,
135        content_length: u64,
136        stream: ClientStream,
137    ) -> Result<UploadProgress> {
138        count(session.object_id.usecase());
139        self.inner
140            .put_chunk(session, offset, content_length, stream)
141            .await
142    }
143
144    async fn upload_offset(&self, session: &Session) -> Result<UploadProgress> {
145        count(session.object_id.usecase());
146        self.inner.upload_offset(session).await
147    }
148
149    async fn cancel_upload(&self, session: &Session) -> Result<()> {
150        count(session.object_id.usecase());
151        self.inner.cancel_upload(session).await
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use objectstore_types::scope::{Scope, Scopes};
158
159    use super::*;
160    use crate::backend::in_memory::InMemoryBackend;
161    use crate::id::ObjectContext;
162    use crate::stream;
163
164    fn object_id(usecase: &str) -> ObjectId {
165        ObjectId::new(
166            ObjectContext {
167                usecase: usecase.parse().unwrap(),
168                scopes: Scopes::from_iter([Scope::create("org", "1").unwrap()]),
169            },
170            "key".into(),
171        )
172    }
173
174    /// Runs `f` on a current-thread runtime while capturing emitted metrics.
175    ///
176    /// The capturing client is thread-local, so the futures must run on the same
177    /// thread that installs it.
178    fn capture(f: impl std::future::Future<Output = ()>) -> Vec<String> {
179        objectstore_metrics::with_capturing_test_client(|| {
180            tokio::runtime::Builder::new_current_thread()
181                .enable_all()
182                .build()
183                .unwrap()
184                .block_on(f);
185        })
186    }
187
188    #[test]
189    fn counts_each_core_operation_once() {
190        let captured = capture(async {
191            let backend = CountingBackend::new(Box::new(InMemoryBackend::new("in-memory")));
192            let id = object_id("attachments");
193
194            backend
195                .put_object(
196                    &id,
197                    &Metadata::default(),
198                    stream::single("hi"),
199                    Timestamp::now(),
200                )
201                .await
202                .unwrap();
203            backend
204                .get_object(&id, Timestamp::now(), None)
205                .await
206                .unwrap();
207            backend.get_metadata(&id, Timestamp::now()).await.unwrap();
208            backend.delete_object(&id, Timestamp::now()).await.unwrap();
209        });
210
211        let cogs = captured
212            .iter()
213            .filter(|m| m.starts_with("cogs.usage"))
214            .count();
215        assert_eq!(
216            cogs, 4,
217            "expected one count per operation, captured: {captured:?}"
218        );
219        assert!(
220            captured
221                .iter()
222                .all(|m| !m.starts_with("cogs.usage")
223                    || m == "cogs.usage:+1|c|#app_feature:attachments"),
224            "captured: {captured:?}"
225        );
226    }
227
228    #[test]
229    fn counts_missing_reads_on_dispatch() {
230        let captured = capture(async {
231            let backend = CountingBackend::new(Box::new(InMemoryBackend::new("in-memory")));
232            // Nothing stored: the read returns `None` but is still billed.
233            let result = backend
234                .get_object(&object_id("attachments"), Timestamp::now(), None)
235                .await
236                .unwrap();
237            assert!(result.is_none());
238        });
239
240        assert_eq!(
241            captured
242                .iter()
243                .filter(|m| m.starts_with("cogs.usage"))
244                .count(),
245            1,
246            "captured: {captured:?}"
247        );
248    }
249
250    #[test]
251    fn new_usecase_is_app_feature() {
252        let captured = capture(async {
253            let backend = CountingBackend::new(Box::new(InMemoryBackend::new("in-memory")));
254            backend
255                .get_object(&object_id("new_usecase"), Timestamp::now(), None)
256                .await
257                .unwrap();
258        });
259
260        assert!(
261            captured
262                .iter()
263                .any(|m| m == "cogs.usage:+1|c|#app_feature:new_usecase"),
264            "captured: {captured:?}"
265        );
266    }
267}
objectstore_service/backend/counting.rs

objectstore_service/backend/
counting.rs
