mito2/sst/index/
store.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::io;
16use std::ops::Range;
17use std::pin::Pin;
18use std::task::{Context, Poll};
19
20use async_trait::async_trait;
21use bytes::{BufMut, Bytes};
22use common_base::range_read::{Metadata, RangeReader, SizeAwareRangeReader};
23use futures::{AsyncRead, AsyncSeek, AsyncWrite};
24use object_store::ObjectStore;
25use pin_project::pin_project;
26use prometheus::IntCounter;
27use snafu::ResultExt;
28
29use crate::error::{OpenDalSnafu, Result};
30
31/// A wrapper around [`ObjectStore`] that adds instrumentation for monitoring
32/// metrics such as bytes read, bytes written, and the number of seek operations.
33///
34/// TODO: Consider refactor InstrumentedStore to use async in trait instead of AsyncRead.
35#[derive(Clone)]
36pub(crate) struct InstrumentedStore {
37    /// The underlying object store.
38    object_store: ObjectStore,
39    /// The size of the write buffer.
40    write_buffer_size: Option<usize>,
41}
42
43impl InstrumentedStore {
44    /// Create a new `InstrumentedStore`.
45    pub fn new(object_store: ObjectStore) -> Self {
46        Self {
47            object_store,
48            write_buffer_size: None,
49        }
50    }
51
52    pub fn store(&self) -> &ObjectStore {
53        &self.object_store
54    }
55
56    /// Set the size of the write buffer.
57    pub fn with_write_buffer_size(mut self, write_buffer_size: Option<usize>) -> Self {
58        self.write_buffer_size = write_buffer_size.filter(|&size| size > 0);
59        self
60    }
61
62    /// Returns an [`InstrumentedRangeReader`] for the given path.
63    /// Metrics like the number of bytes read are recorded using the provided `IntCounter`.
64    pub async fn range_reader<'a>(
65        &self,
66        path: &str,
67        read_byte_count: &'a IntCounter,
68        read_count: &'a IntCounter,
69    ) -> Result<InstrumentedRangeReader<'a>> {
70        Ok(InstrumentedRangeReader {
71            store: self.object_store.clone(),
72            path: path.to_string(),
73            read_byte_count,
74            read_count,
75            file_size_hint: None,
76        })
77    }
78
79    /// Returns an [`InstrumentedAsyncRead`] for the given path.
80    /// Metrics like the number of bytes read, read and seek operations
81    /// are recorded using the provided `IntCounter`s.
82    pub async fn reader<'a>(
83        &self,
84        path: &str,
85        read_byte_count: &'a IntCounter,
86        read_count: &'a IntCounter,
87        seek_count: &'a IntCounter,
88    ) -> Result<InstrumentedAsyncRead<'a, object_store::FuturesAsyncReader>> {
89        let meta = self.object_store.stat(path).await.context(OpenDalSnafu)?;
90        let reader = self
91            .object_store
92            .reader(path)
93            .await
94            .context(OpenDalSnafu)?
95            .into_futures_async_read(0..meta.content_length())
96            .await
97            .context(OpenDalSnafu)?;
98        Ok(InstrumentedAsyncRead::new(
99            reader,
100            read_byte_count,
101            read_count,
102            seek_count,
103        ))
104    }
105
106    /// Returns an [`InstrumentedAsyncWrite`] for the given path.
107    /// Metrics like the number of bytes written, write and flush operations
108    /// are recorded using the provided `IntCounter`s.
109    pub async fn writer<'a>(
110        &self,
111        path: &str,
112        write_byte_count: &'a IntCounter,
113        write_count: &'a IntCounter,
114        flush_count: &'a IntCounter,
115    ) -> Result<InstrumentedAsyncWrite<'a, object_store::FuturesAsyncWriter>> {
116        let writer = match self.write_buffer_size {
117            Some(size) => self
118                .object_store
119                .writer_with(path)
120                .chunk(size)
121                .await
122                .context(OpenDalSnafu)?
123                .into_futures_async_write(),
124            None => self
125                .object_store
126                .writer(path)
127                .await
128                .context(OpenDalSnafu)?
129                .into_futures_async_write(),
130        };
131        Ok(InstrumentedAsyncWrite::new(
132            writer,
133            write_byte_count,
134            write_count,
135            flush_count,
136        ))
137    }
138
139    /// Proxies to [`ObjectStore::list`].
140    pub async fn list(&self, path: &str) -> Result<Vec<object_store::Entry>> {
141        let list = self.object_store.list(path).await.context(OpenDalSnafu)?;
142        Ok(list)
143    }
144
145    /// Proxies to [`ObjectStore::remove_all`].
146    pub async fn remove_all(&self, path: &str) -> Result<()> {
147        self.object_store
148            .remove_all(path)
149            .await
150            .context(OpenDalSnafu)
151    }
152}
153
154/// A wrapper around [`AsyncRead`] that adds instrumentation for monitoring
155#[pin_project]
156pub(crate) struct InstrumentedAsyncRead<'a, R> {
157    #[pin]
158    inner: R,
159    read_byte_count: CounterGuard<'a>,
160    read_count: CounterGuard<'a>,
161    seek_count: CounterGuard<'a>,
162}
163
164impl<'a, R> InstrumentedAsyncRead<'a, R> {
165    /// Create a new `InstrumentedAsyncRead`.
166    fn new(
167        inner: R,
168        read_byte_count: &'a IntCounter,
169        read_count: &'a IntCounter,
170        seek_count: &'a IntCounter,
171    ) -> Self {
172        Self {
173            inner,
174            read_byte_count: CounterGuard::new(read_byte_count),
175            read_count: CounterGuard::new(read_count),
176            seek_count: CounterGuard::new(seek_count),
177        }
178    }
179}
180
181impl<R: AsyncRead + Unpin + Send> AsyncRead for InstrumentedAsyncRead<'_, R> {
182    fn poll_read(
183        mut self: Pin<&mut Self>,
184        cx: &mut Context<'_>,
185        buf: &mut [u8],
186    ) -> Poll<io::Result<usize>> {
187        let poll = self.as_mut().project().inner.poll_read(cx, buf);
188        if let Poll::Ready(Ok(n)) = &poll {
189            self.read_count.inc_by(1);
190            self.read_byte_count.inc_by(*n);
191        }
192        poll
193    }
194}
195
196impl<R: AsyncSeek + Unpin + Send> AsyncSeek for InstrumentedAsyncRead<'_, R> {
197    fn poll_seek(
198        mut self: Pin<&mut Self>,
199        cx: &mut Context<'_>,
200        pos: io::SeekFrom,
201    ) -> Poll<io::Result<u64>> {
202        let poll = self.as_mut().project().inner.poll_seek(cx, pos);
203        if let Poll::Ready(Ok(_)) = &poll {
204            self.seek_count.inc_by(1);
205        }
206        poll
207    }
208}
209
210/// A wrapper around [`AsyncWrite`] that adds instrumentation for monitoring
211#[pin_project]
212pub(crate) struct InstrumentedAsyncWrite<'a, W> {
213    #[pin]
214    inner: W,
215    write_byte_count: CounterGuard<'a>,
216    write_count: CounterGuard<'a>,
217    flush_count: CounterGuard<'a>,
218}
219
220impl<'a, W> InstrumentedAsyncWrite<'a, W> {
221    /// Create a new `InstrumentedAsyncWrite`.
222    fn new(
223        inner: W,
224        write_byte_count: &'a IntCounter,
225        write_count: &'a IntCounter,
226        flush_count: &'a IntCounter,
227    ) -> Self {
228        Self {
229            inner,
230            write_byte_count: CounterGuard::new(write_byte_count),
231            write_count: CounterGuard::new(write_count),
232            flush_count: CounterGuard::new(flush_count),
233        }
234    }
235}
236
237impl<W: AsyncWrite + Unpin + Send> AsyncWrite for InstrumentedAsyncWrite<'_, W> {
238    fn poll_write(
239        mut self: Pin<&mut Self>,
240        cx: &mut Context<'_>,
241        buf: &[u8],
242    ) -> Poll<io::Result<usize>> {
243        let poll = self.as_mut().project().inner.poll_write(cx, buf);
244        if let Poll::Ready(Ok(n)) = &poll {
245            self.write_count.inc_by(1);
246            self.write_byte_count.inc_by(*n);
247        }
248        poll
249    }
250
251    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
252        let poll = self.as_mut().project().inner.poll_flush(cx);
253        if let Poll::Ready(Ok(())) = &poll {
254            self.flush_count.inc_by(1);
255        }
256        poll
257    }
258
259    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
260        self.project().inner.poll_close(cx)
261    }
262}
263
264/// Implements `RangeReader` for `ObjectStore` and record metrics.
265pub(crate) struct InstrumentedRangeReader<'a> {
266    store: ObjectStore,
267    path: String,
268    read_byte_count: &'a IntCounter,
269    read_count: &'a IntCounter,
270    file_size_hint: Option<u64>,
271}
272
273impl SizeAwareRangeReader for InstrumentedRangeReader<'_> {
274    fn with_file_size_hint(&mut self, file_size_hint: u64) {
275        self.file_size_hint = Some(file_size_hint);
276    }
277}
278
279#[async_trait]
280impl RangeReader for InstrumentedRangeReader<'_> {
281    async fn metadata(&self) -> io::Result<Metadata> {
282        match self.file_size_hint {
283            Some(file_size_hint) => Ok(Metadata {
284                content_length: file_size_hint,
285            }),
286            None => {
287                let stat = self.store.stat(&self.path).await?;
288                Ok(Metadata {
289                    content_length: stat.content_length(),
290                })
291            }
292        }
293    }
294
295    async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
296        let buf = self.store.reader(&self.path).await?.read(range).await?;
297        self.read_byte_count.inc_by(buf.len() as _);
298        self.read_count.inc_by(1);
299        Ok(buf.to_bytes())
300    }
301
302    async fn read_into(&self, range: Range<u64>, buf: &mut (impl BufMut + Send)) -> io::Result<()> {
303        let reader = self.store.reader(&self.path).await?;
304        let size = reader.read_into(buf, range).await?;
305        self.read_byte_count.inc_by(size as _);
306        self.read_count.inc_by(1);
307        Ok(())
308    }
309
310    async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
311        let bufs = self
312            .store
313            .reader(&self.path)
314            .await?
315            .fetch(ranges.to_owned())
316            .await?;
317        let total_size: usize = bufs.iter().map(|buf| buf.len()).sum();
318        self.read_byte_count.inc_by(total_size as _);
319        self.read_count.inc_by(1);
320        Ok(bufs.into_iter().map(|buf| buf.to_bytes()).collect())
321    }
322}
323
324/// A guard that increments a counter when dropped.
325struct CounterGuard<'a> {
326    count: usize,
327    counter: &'a IntCounter,
328}
329
330impl<'a> CounterGuard<'a> {
331    /// Create a new `CounterGuard`.
332    fn new(counter: &'a IntCounter) -> Self {
333        Self { count: 0, counter }
334    }
335
336    /// Increment the counter by `n`.
337    fn inc_by(&mut self, n: usize) {
338        self.count += n;
339    }
340}
341
342impl Drop for CounterGuard<'_> {
343    fn drop(&mut self) {
344        if self.count > 0 {
345            self.counter.inc_by(self.count as _);
346        }
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use futures::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
353    use object_store::services::Memory;
354
355    use super::*;
356
357    #[tokio::test]
358    async fn test_instrumented_store_read_write() {
359        let instrumented_store =
360            InstrumentedStore::new(ObjectStore::new(Memory::default()).unwrap().finish());
361
362        let read_byte_count = IntCounter::new("read_byte_count", "read_byte_count").unwrap();
363        let read_count = IntCounter::new("read_count", "read_count").unwrap();
364        let seek_count = IntCounter::new("seek_count", "seek_count").unwrap();
365        let write_byte_count = IntCounter::new("write_byte_count", "write_byte_count").unwrap();
366        let write_count = IntCounter::new("write_count", "write_count").unwrap();
367        let flush_count = IntCounter::new("flush_count", "flush_count").unwrap();
368
369        let mut writer = instrumented_store
370            .writer("my_file", &write_byte_count, &write_count, &flush_count)
371            .await
372            .unwrap();
373        writer.write_all(b"hello").await.unwrap();
374        writer.flush().await.unwrap();
375        writer.close().await.unwrap();
376        drop(writer);
377
378        let mut reader = instrumented_store
379            .reader("my_file", &read_byte_count, &read_count, &seek_count)
380            .await
381            .unwrap();
382        let mut buf = vec![0; 5];
383        reader.read_exact(&mut buf).await.unwrap();
384        reader.seek(io::SeekFrom::Start(0)).await.unwrap();
385        reader.read_exact(&mut buf).await.unwrap();
386        drop(reader);
387
388        assert_eq!(read_byte_count.get(), 10);
389        assert_eq!(read_count.get(), 2);
390        assert_eq!(seek_count.get(), 1);
391        assert_eq!(write_byte_count.get(), 5);
392        assert_eq!(write_count.get(), 1);
393        assert_eq!(flush_count.get(), 1);
394    }
395}
mito2/sst/index/store.rs

mito2/sst/index/
store.rs
