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    /// Recursively deletes all objects under the given path.
146    pub async fn remove_all(&self, path: &str) -> Result<()> {
147        self.object_store
148            .delete_with(path)
149            .recursive(true)
150            .await
151            .context(OpenDalSnafu)
152    }
153}
154
155/// A wrapper around [`AsyncRead`] that adds instrumentation for monitoring
156#[pin_project]
157pub(crate) struct InstrumentedAsyncRead<'a, R> {
158    #[pin]
159    inner: R,
160    read_byte_count: CounterGuard<'a>,
161    read_count: CounterGuard<'a>,
162    seek_count: CounterGuard<'a>,
163}
164
165impl<'a, R> InstrumentedAsyncRead<'a, R> {
166    /// Create a new `InstrumentedAsyncRead`.
167    fn new(
168        inner: R,
169        read_byte_count: &'a IntCounter,
170        read_count: &'a IntCounter,
171        seek_count: &'a IntCounter,
172    ) -> Self {
173        Self {
174            inner,
175            read_byte_count: CounterGuard::new(read_byte_count),
176            read_count: CounterGuard::new(read_count),
177            seek_count: CounterGuard::new(seek_count),
178        }
179    }
180}
181
182impl<R: AsyncRead + Unpin + Send> AsyncRead for InstrumentedAsyncRead<'_, R> {
183    fn poll_read(
184        mut self: Pin<&mut Self>,
185        cx: &mut Context<'_>,
186        buf: &mut [u8],
187    ) -> Poll<io::Result<usize>> {
188        let poll = self.as_mut().project().inner.poll_read(cx, buf);
189        if let Poll::Ready(Ok(n)) = &poll {
190            self.read_count.inc_by(1);
191            self.read_byte_count.inc_by(*n);
192        }
193        poll
194    }
195}
196
197impl<R: AsyncSeek + Unpin + Send> AsyncSeek for InstrumentedAsyncRead<'_, R> {
198    fn poll_seek(
199        mut self: Pin<&mut Self>,
200        cx: &mut Context<'_>,
201        pos: io::SeekFrom,
202    ) -> Poll<io::Result<u64>> {
203        let poll = self.as_mut().project().inner.poll_seek(cx, pos);
204        if let Poll::Ready(Ok(_)) = &poll {
205            self.seek_count.inc_by(1);
206        }
207        poll
208    }
209}
210
211/// A wrapper around [`AsyncWrite`] that adds instrumentation for monitoring
212#[pin_project]
213pub(crate) struct InstrumentedAsyncWrite<'a, W> {
214    #[pin]
215    inner: W,
216    write_byte_count: CounterGuard<'a>,
217    write_count: CounterGuard<'a>,
218    flush_count: CounterGuard<'a>,
219}
220
221impl<'a, W> InstrumentedAsyncWrite<'a, W> {
222    /// Create a new `InstrumentedAsyncWrite`.
223    fn new(
224        inner: W,
225        write_byte_count: &'a IntCounter,
226        write_count: &'a IntCounter,
227        flush_count: &'a IntCounter,
228    ) -> Self {
229        Self {
230            inner,
231            write_byte_count: CounterGuard::new(write_byte_count),
232            write_count: CounterGuard::new(write_count),
233            flush_count: CounterGuard::new(flush_count),
234        }
235    }
236}
237
238impl<W: AsyncWrite + Unpin + Send> AsyncWrite for InstrumentedAsyncWrite<'_, W> {
239    fn poll_write(
240        mut self: Pin<&mut Self>,
241        cx: &mut Context<'_>,
242        buf: &[u8],
243    ) -> Poll<io::Result<usize>> {
244        let poll = self.as_mut().project().inner.poll_write(cx, buf);
245        if let Poll::Ready(Ok(n)) = &poll {
246            self.write_count.inc_by(1);
247            self.write_byte_count.inc_by(*n);
248        }
249        poll
250    }
251
252    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
253        let poll = self.as_mut().project().inner.poll_flush(cx);
254        if let Poll::Ready(Ok(())) = &poll {
255            self.flush_count.inc_by(1);
256        }
257        poll
258    }
259
260    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
261        self.project().inner.poll_close(cx)
262    }
263}
264
265/// Ranges fetched in parallel by one `read_vec`, as for parquet reads.
266const FETCH_CONCURRENCY: usize = 8;
267
268/// Implements `RangeReader` for `ObjectStore` and record metrics.
269pub(crate) struct InstrumentedRangeReader<'a> {
270    store: ObjectStore,
271    path: String,
272    read_byte_count: &'a IntCounter,
273    read_count: &'a IntCounter,
274    file_size_hint: Option<u64>,
275}
276
277impl SizeAwareRangeReader for InstrumentedRangeReader<'_> {
278    fn with_file_size_hint(&mut self, file_size_hint: u64) {
279        self.file_size_hint = Some(file_size_hint);
280    }
281}
282
283#[async_trait]
284impl RangeReader for InstrumentedRangeReader<'_> {
285    async fn metadata(&self) -> io::Result<Metadata> {
286        match self.file_size_hint {
287            Some(file_size_hint) => Ok(Metadata {
288                content_length: file_size_hint,
289            }),
290            None => {
291                let stat = self.store.stat(&self.path).await?;
292                Ok(Metadata {
293                    content_length: stat.content_length(),
294                })
295            }
296        }
297    }
298
299    async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
300        let buf = self.store.reader(&self.path).await?.read(range).await?;
301        self.read_byte_count.inc_by(buf.len() as _);
302        self.read_count.inc_by(1);
303        Ok(buf.to_bytes())
304    }
305
306    async fn read_into(&self, range: Range<u64>, buf: &mut (impl BufMut + Send)) -> io::Result<()> {
307        let reader = self.store.reader(&self.path).await?;
308        let size = reader.read_into(buf, range).await?;
309        self.read_byte_count.inc_by(size as _);
310        self.read_count.inc_by(1);
311        Ok(())
312    }
313
314    async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
315        // OpenDAL merges ranges less than 1 MiB apart by default without bounding the merged
316        // length, so sparse page reads could fetch most of the file. Only merge adjacent
317        // ranges, and fetch the rest in parallel like parquet reads do.
318        // TODO: decide whether to merge across small gaps based on object store latency,
319        // bytes read and query latency.
320        let bufs = self
321            .store
322            .reader_with(&self.path)
323            .gap(0)
324            .concurrent(FETCH_CONCURRENCY)
325            .await?
326            .fetch(ranges.to_owned())
327            .await?;
328        let total_size: usize = bufs.iter().map(|buf| buf.len()).sum();
329        self.read_byte_count.inc_by(total_size as _);
330        self.read_count.inc_by(1);
331        Ok(bufs.into_iter().map(|buf| buf.to_bytes()).collect())
332    }
333}
334
335/// A guard that increments a counter when dropped.
336struct CounterGuard<'a> {
337    count: usize,
338    counter: &'a IntCounter,
339}
340
341impl<'a> CounterGuard<'a> {
342    /// Create a new `CounterGuard`.
343    fn new(counter: &'a IntCounter) -> Self {
344        Self { count: 0, counter }
345    }
346
347    /// Increment the counter by `n`.
348    fn inc_by(&mut self, n: usize) {
349        self.count += n;
350    }
351}
352
353impl Drop for CounterGuard<'_> {
354    fn drop(&mut self) {
355        if self.count > 0 {
356            self.counter.inc_by(self.count as _);
357        }
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use futures::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
364    use object_store::services::Memory;
365
366    use super::*;
367
368    #[tokio::test]
369    async fn test_instrumented_store_read_write() {
370        let instrumented_store =
371            InstrumentedStore::new(ObjectStore::new(Memory::default()).unwrap());
372
373        let read_byte_count = IntCounter::new("read_byte_count", "read_byte_count").unwrap();
374        let read_count = IntCounter::new("read_count", "read_count").unwrap();
375        let seek_count = IntCounter::new("seek_count", "seek_count").unwrap();
376        let write_byte_count = IntCounter::new("write_byte_count", "write_byte_count").unwrap();
377        let write_count = IntCounter::new("write_count", "write_count").unwrap();
378        let flush_count = IntCounter::new("flush_count", "flush_count").unwrap();
379
380        let mut writer = instrumented_store
381            .writer("my_file", &write_byte_count, &write_count, &flush_count)
382            .await
383            .unwrap();
384        writer.write_all(b"hello").await.unwrap();
385        writer.flush().await.unwrap();
386        writer.close().await.unwrap();
387        drop(writer);
388
389        let mut reader = instrumented_store
390            .reader("my_file", &read_byte_count, &read_count, &seek_count)
391            .await
392            .unwrap();
393        let mut buf = vec![0; 5];
394        reader.read_exact(&mut buf).await.unwrap();
395        reader.seek(io::SeekFrom::Start(0)).await.unwrap();
396        reader.read_exact(&mut buf).await.unwrap();
397        drop(reader);
398
399        assert_eq!(read_byte_count.get(), 10);
400        assert_eq!(read_count.get(), 2);
401        assert_eq!(seek_count.get(), 1);
402        assert_eq!(write_byte_count.get(), 5);
403        assert_eq!(write_count.get(), 1);
404        assert_eq!(flush_count.get(), 1);
405    }
406}
mito2/sst/index/store.rs

mito2/sst/index/
store.rs
