mito2/cache/range_cache/
shared_range.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
15//! A single range scan with independently advancing readers.
16
17use std::collections::VecDeque;
18use std::sync::atomic::{AtomicUsize, Ordering};
19use std::sync::{Arc, Weak};
20
21use datatypes::arrow::record_batch::RecordBatch;
22use futures::future::{BoxFuture, Shared};
23use futures::{FutureExt, TryStreamExt};
24use tokio::sync::{Mutex, OwnedSemaphorePermit};
25
26use crate::cache::RangeResultMemoryLimiter;
27use crate::cache::range_cache::{self, CachedBatchSlice};
28use crate::error::Error;
29use crate::read::BoxedRecordBatchStream;
30
31type SharedBatch = Result<Option<RecordBatch>, Arc<Error>>;
32type NextBatch = Shared<BoxFuture<'static, SharedBatch>>;
33
34/// The registry must hold only a weak reference: readers own the scan lifetime.
35/// The source is polled on demand, so an unused range does not run ahead.
36/// Cache admission stops when retention reservations run out. Unread batches
37/// remain owned by active readers; this is not a hard limit on their backlog.
38pub(crate) struct SharedRange {
39    source: Arc<Mutex<Option<BoxedRecordBatchStream>>>,
40    buffer: Mutex<RangeBuffer>,
41    limiter: Arc<RangeResultMemoryLimiter>,
42}
43
44/// A reader's cursor refers to original batch boundaries, including empty batches.
45pub(crate) struct RangeCursor {
46    batch: Arc<AtomicUsize>,
47}
48
49pub(crate) enum RangeRead {
50    Batch(RecordBatch),
51    Finished,
52    Failed(Arc<Error>),
53}
54
55impl SharedRange {
56    /// Wraps one source stream and its full-result retention budget.
57    pub(crate) fn new(
58        source: BoxedRecordBatchStream,
59        limiter: Arc<RangeResultMemoryLimiter>,
60    ) -> Arc<Self> {
61        Arc::new(Self {
62            source: Arc::new(Mutex::new(Some(source))),
63            buffer: Mutex::new(RangeBuffer::default()),
64            limiter,
65        })
66    }
67
68    /// Register before consuming any data. Once full-result retention stops,
69    /// new queries use their own ordinary scan rather than joining mid-range.
70    pub(crate) async fn subscribe(&self) -> Option<RangeCursor> {
71        let mut buffer = self.buffer.lock().await;
72        if buffer.streaming_only || matches!(buffer.finished, Some(Err(_))) {
73            return None;
74        }
75        let batch = Arc::new(AtomicUsize::new(0));
76        buffer.readers.push(Arc::downgrade(&batch));
77        Some(RangeCursor { batch })
78    }
79
80    /// Exactly one reader hands the complete result to the existing compactor.
81    pub(crate) async fn take_complete_batches(&self) -> Option<Vec<CachedBatchSlice>> {
82        let mut buffer = self.buffer.lock().await;
83        if buffer.streaming_only
84            || buffer.finalization_started
85            || !matches!(buffer.finished, Some(Ok(())))
86        {
87            return None;
88        }
89        buffer.finalization_started = true;
90        let mut batches: Vec<CachedBatchSlice> = Vec::new();
91        let mut previous = None;
92        for entry in &buffer.batches {
93            let allocation = Arc::as_ptr(&entry.allocation);
94            if previous == Some(allocation) {
95                if let Some(last) = batches.last_mut() {
96                    last.slice_lengths.push(entry.batch.num_rows());
97                }
98            } else {
99                batches.push(CachedBatchSlice {
100                    batch: entry.allocation.batch.clone(),
101                    slice_lengths: vec![entry.batch.num_rows()],
102                });
103                previous = Some(allocation);
104            }
105        }
106        Some(batches)
107    }
108
109    /// Returns the next batch for this reader, polling the shared source as needed.
110    pub(crate) async fn read(&self, cursor: &mut RangeCursor) -> RangeRead {
111        loop {
112            let (index, next) = {
113                let mut buffer = self.buffer.lock().await;
114                buffer.release_consumed();
115                let position = cursor.batch.load(Ordering::Relaxed);
116                if let Some(batch) = buffer.batches.get(position - buffer.first_batch) {
117                    let batch = batch.batch.clone();
118                    cursor.batch.store(position + 1, Ordering::Relaxed);
119                    buffer.release_consumed();
120                    return RangeRead::Batch(batch);
121                }
122                if let Some(result) = &buffer.finished {
123                    return match result {
124                        Ok(()) => RangeRead::Finished,
125                        Err(error) => RangeRead::Failed(error.clone()),
126                    };
127                }
128                let index = buffer.first_batch + buffer.batches.len();
129                let next = buffer
130                    .next
131                    .get_or_insert_with(|| {
132                        let source = self.source.clone();
133                        async move {
134                            let mut source = source.lock().await;
135                            let result = match source.as_mut() {
136                                Some(stream) => stream.try_next().await.map_err(Arc::new),
137                                None => Ok(None),
138                            };
139                            if !matches!(&result, Ok(Some(_))) {
140                                // Release file readers as soon as the source terminates.
141                                *source = None;
142                            }
143                            result
144                        }
145                        .boxed()
146                        .shared()
147                    })
148                    .clone();
149                (index, next)
150            };
151
152            // No buffer lock is held across source I/O. A slower reader can still
153            // consume buffered batches while another reader awaits the next one.
154            let result = next.await;
155            let mut buffer = self.buffer.lock().await;
156            if index != buffer.first_batch + buffer.batches.len() || buffer.finished.is_some() {
157                continue;
158            }
159            buffer.next = None;
160            match result {
161                Ok(Some(batch)) => {
162                    let bytes = batch.get_array_memory_size();
163                    let permit = if buffer.streaming_only {
164                        None
165                    } else {
166                        self.limiter.try_acquire(bytes)
167                    };
168                    buffer.retained_bytes += bytes;
169                    crate::metrics::RANGE_CACHE_STAGING_BYTES.add(bytes as i64);
170                    if buffer.streaming_only {
171                        crate::metrics::RANGE_CACHE_STREAMING_BYTES.add(bytes as i64);
172                    }
173                    let missing_permit = permit.is_none();
174                    let allocation = Arc::new(BufferedAllocation {
175                        batch: batch.clone(),
176                        _permit: permit,
177                    });
178                    buffer
179                        .batches
180                        .push_back(BufferedBatch { batch, allocation });
181                    if !buffer.streaming_only {
182                        buffer.pending_batches += 1;
183                        buffer.pending_rows +=
184                            buffer.batches.back().map_or(0, |b| b.batch.num_rows());
185                        buffer.pending_bytes += bytes;
186                        let compact = missing_permit
187                            || range_cache::should_compact_batches(
188                                buffer.pending_batches,
189                                buffer.pending_rows,
190                                buffer.pending_bytes,
191                            );
192                        let retained = !compact || buffer.compact_pending(&self.limiter);
193                        if !retained {
194                            crate::metrics::RANGE_CACHE_WORKER_EVENTS
195                                .with_label_values(&["budget_fallback"])
196                                .inc();
197                            buffer.streaming_only = true;
198                            crate::metrics::RANGE_CACHE_STREAMING_BYTES
199                                .add(buffer.retained_bytes as i64);
200                        }
201                    }
202                    buffer.release_consumed();
203                }
204                Ok(None) => buffer.finished = Some(Ok(())),
205                Err(error) => buffer.finished = Some(Err(error)),
206            }
207        }
208    }
209}
210
211#[derive(Default)]
212struct RangeBuffer {
213    batches: VecDeque<BufferedBatch>,
214    first_batch: usize,
215    pending_batches: usize,
216    pending_rows: usize,
217    pending_bytes: usize,
218    retained_bytes: usize,
219    next: Option<NextBatch>,
220    finished: Option<Result<(), Arc<Error>>>,
221    streaming_only: bool,
222    finalization_started: bool,
223    readers: Vec<Weak<AtomicUsize>>,
224}
225
226impl RangeBuffer {
227    /// Detaches the pending slices before deciding whether to stop cache admission.
228    /// Each compacted allocation owns one reservation, shared by its batch views.
229    fn compact_pending(&mut self, limiter: &RangeResultMemoryLimiter) -> bool {
230        let start = self.batches.len() - self.pending_batches;
231        let input = self
232            .batches
233            .iter()
234            .skip(start)
235            .map(|entry| entry.batch.clone())
236            .collect();
237        let compacted = match range_cache::compact_record_batches(input) {
238            Ok(compacted) => compacted,
239            Err(error) => {
240                common_telemetry::warn!(error; "Failed to compact shared range buffer");
241                return false;
242            }
243        };
244        let old_bytes = self.pending_bytes;
245        let bytes = compacted.batch.get_array_memory_size();
246        // Release the old reservations before the nonblocking replacement attempt.
247        // Reader-owned batches retain their own Arrow references independently.
248        self.batches.truncate(start);
249        let permit = limiter.try_acquire(bytes);
250        let retained = permit.is_some();
251        let allocation = Arc::new(BufferedAllocation {
252            batch: compacted.batch,
253            _permit: permit,
254        });
255        let mut offset = 0;
256        for len in compacted.slice_lengths {
257            self.batches.push_back(BufferedBatch {
258                batch: allocation.batch.slice(offset, len),
259                allocation: allocation.clone(),
260            });
261            offset += len;
262        }
263        self.retained_bytes = self.retained_bytes - old_bytes + bytes;
264        crate::metrics::RANGE_CACHE_STAGING_BYTES.add(bytes as i64 - old_bytes as i64);
265        self.pending_batches = 0;
266        self.pending_rows = 0;
267        self.pending_bytes = 0;
268        retained
269    }
270
271    fn release_consumed(&mut self) {
272        if !self.streaming_only {
273            return;
274        }
275        let mut consumed = self.first_batch + self.batches.len();
276        self.readers.retain(|reader| {
277            if let Some(position) = reader.upgrade() {
278                consumed = consumed.min(position.load(Ordering::Relaxed));
279                true
280            } else {
281                false
282            }
283        });
284        while self.first_batch < consumed {
285            if let Some(batch) = self.batches.pop_front()
286                && Arc::strong_count(&batch.allocation) == 1
287            {
288                let bytes = batch.allocation.batch.get_array_memory_size();
289                self.retained_bytes -= bytes;
290                crate::metrics::RANGE_CACHE_STAGING_BYTES.sub(bytes as i64);
291                crate::metrics::RANGE_CACHE_STREAMING_BYTES.sub(bytes as i64);
292            }
293            self.first_batch += 1;
294        }
295    }
296}
297
298impl Drop for RangeBuffer {
299    fn drop(&mut self) {
300        crate::metrics::RANGE_CACHE_STAGING_BYTES.sub(self.retained_bytes as i64);
301        if self.streaming_only {
302            crate::metrics::RANGE_CACHE_STREAMING_BYTES.sub(self.retained_bytes as i64);
303        }
304    }
305}
306
307struct BufferedBatch {
308    batch: RecordBatch,
309    allocation: Arc<BufferedAllocation>,
310}
311
312struct BufferedAllocation {
313    batch: RecordBatch,
314    _permit: Option<OwnedSemaphorePermit>,
315}
316
317#[cfg(test)]
318mod tests {
319    use std::sync::atomic::{AtomicUsize, Ordering};
320
321    use datatypes::arrow::array::{ArrayRef, Int64Array};
322    use futures::stream;
323    use tokio::sync::Notify;
324
325    use super::*;
326
327    fn batch(value: i64) -> RecordBatch {
328        let values = Arc::new(Int64Array::from(vec![value])) as ArrayRef;
329        RecordBatch::try_from_iter([("value", values)]).unwrap()
330    }
331
332    fn value(read: RangeRead) -> i64 {
333        match read {
334            RangeRead::Batch(batch) => batch
335                .column(0)
336                .as_any()
337                .downcast_ref::<Int64Array>()
338                .unwrap()
339                .value(0),
340            _ => panic!("expected a batch"),
341        }
342    }
343
344    #[tokio::test]
345    async fn test_pending_source_does_not_block_buffered_readers() {
346        let produced = Arc::new(AtomicUsize::new(0));
347        let counter = produced.clone();
348        let gate = Arc::new(Notify::new());
349        let ready = gate.clone();
350        let source = async_stream::try_stream! {
351            counter.fetch_add(1, Ordering::Relaxed);
352            yield batch(1);
353            ready.notified().await;
354            counter.fetch_add(1, Ordering::Relaxed);
355            yield batch(2);
356        };
357        let limiter = Arc::new(RangeResultMemoryLimiter::new(4096, 1));
358        let shared = SharedRange::new(Box::pin(source), limiter.clone());
359        let mut first = shared.subscribe().await.unwrap();
360        let mut second = shared.subscribe().await.unwrap();
361        assert_eq!(value(shared.read(&mut first).await), 1);
362        {
363            let next = shared.read(&mut first);
364            tokio::pin!(next);
365            assert!(futures::poll!(&mut next).is_pending());
366            assert_eq!(value(shared.read(&mut second).await), 1);
367        }
368        drop(first);
369        gate.notify_one();
370        assert_eq!(value(shared.read(&mut second).await), 2);
371        assert!(matches!(
372            shared.read(&mut second).await,
373            RangeRead::Finished
374        ));
375        assert_eq!(produced.load(Ordering::Relaxed), 2);
376        drop(shared);
377        assert_eq!(limiter.available_permits(), 4096);
378    }
379
380    #[tokio::test]
381    async fn test_slow_reader_backlog_grows_and_is_reclaimed() {
382        let batches = (0..256).map(batch).collect::<Vec<_>>();
383        let expected_bytes = batches
384            .iter()
385            .map(RecordBatch::get_array_memory_size)
386            .sum::<usize>();
387        let shared = SharedRange::new(
388            Box::pin(stream::iter(batches.into_iter().map(Ok))),
389            Arc::new(RangeResultMemoryLimiter::new(1, 1)),
390        );
391        let mut fast = shared.subscribe().await.unwrap();
392        let mut slow = shared.subscribe().await.unwrap();
393        // A paused subscriber must not stop the other reader or discard its data.
394        for expected in 0..256 {
395            assert_eq!(value(shared.read(&mut fast).await), expected);
396        }
397        assert!(matches!(shared.read(&mut fast).await, RangeRead::Finished));
398        {
399            let buffer = shared.buffer.lock().await;
400            assert!(buffer.streaming_only);
401            assert_eq!(buffer.retained_bytes, expected_bytes);
402            assert_eq!(buffer.batches.len(), 256);
403        }
404        assert!(shared.subscribe().await.is_none());
405        for expected in 0..256 {
406            assert_eq!(value(shared.read(&mut slow).await), expected);
407        }
408        assert!(matches!(shared.read(&mut slow).await, RangeRead::Finished));
409        let buffer = shared.buffer.lock().await;
410        assert!(buffer.batches.is_empty());
411        assert_eq!(buffer.retained_bytes, 0);
412    }
413}
mito2/cache/range_cache/shared_range.rs

mito2/cache/range_cache/
shared_range.rs
