session/
write_progress.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//! Process-local acknowledgements for a finite write request.
16
17use std::fmt;
18use std::sync::Arc;
19use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
20use tokio::sync::Notify;
21
22/// Counts only confirmed storage writes, including partial successes.
23#[derive(Default)]
24pub struct WriteProgress {
25    written: AtomicU64,
26    accepted: AtomicBool,
27    // A queued write must keep its decoder's memory reservation alive.
28    _reservation: Option<Arc<dyn Send + Sync>>,
29    on_write: Option<Box<dyn Fn(u64) + Send + Sync>>,
30    pending: AtomicUsize,
31    completed: Notify,
32}
33
34/// Keeps a queued submission pending until the worker reports its result.
35pub struct PendingWrite(Arc<WriteProgress>);
36
37impl Drop for PendingWrite {
38    fn drop(&mut self) {
39        if self.0.pending.fetch_sub(1, Ordering::AcqRel) == 1 {
40            self.0.completed.notify_waiters();
41        }
42    }
43}
44
45impl fmt::Debug for WriteProgress {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        f.debug_struct("WriteProgress")
48            .field("written", &self.written())
49            .field("accepted", &self.accepted())
50            .finish_non_exhaustive()
51    }
52}
53
54impl WriteProgress {
55    pub fn start(self: &Arc<Self>) -> PendingWrite {
56        self.pending.fetch_add(1, Ordering::Relaxed);
57        PendingWrite(self.clone())
58    }
59
60    /// Waits for previously queued submissions before returning a failed request's counts.
61    pub async fn wait(&self) {
62        loop {
63            let notified = self.completed.notified();
64            tokio::pin!(notified);
65            notified.as_mut().enable();
66            if self.pending.load(Ordering::Acquire) == 0 {
67                return;
68            }
69            notified.await;
70        }
71    }
72
73    pub fn with_reservation(reservation: Arc<dyn Send + Sync>) -> Self {
74        Self {
75            _reservation: Some(reservation),
76            ..Self::default()
77        }
78    }
79
80    pub fn add_written(&self, rows: usize) {
81        self.written.fetch_add(rows as u64, Ordering::Relaxed);
82        self.observe(rows as u64);
83    }
84
85    pub fn with_observer(mut self, observer: impl Fn(u64) + Send + Sync + 'static) -> Self {
86        self.on_write = Some(Box::new(observer));
87        self
88    }
89
90    /// A synchronous whole-request acknowledgement also covers handlers without region tracking.
91    pub fn confirm(&self, rows: u64) {
92        let previous = self.written.fetch_max(rows, Ordering::Relaxed);
93        self.observe(rows.saturating_sub(previous));
94    }
95
96    fn observe(&self, rows: u64) {
97        if rows > 0
98            && let Some(observer) = &self.on_write
99        {
100            observer(rows);
101        }
102    }
103
104    pub fn written(&self) -> u64 {
105        self.written.load(Ordering::Relaxed)
106    }
107
108    /// Records an acknowledgement returned before storage completion.
109    pub fn mark_accepted(&self) {
110        self.accepted.store(true, Ordering::Relaxed);
111    }
112
113    pub fn accepted(&self) -> bool {
114        self.accepted.load(Ordering::Relaxed)
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use std::sync::atomic::AtomicUsize;
121
122    use super::*;
123
124    #[test]
125    fn confirmations_do_not_double_count_and_contexts_keep_local_ownership() {
126        let observed = Arc::new(AtomicUsize::new(0));
127        let counter = observed.clone();
128        let progress = Arc::new(WriteProgress::default().with_observer(move |rows| {
129            counter.fetch_add(rows as usize, Ordering::Relaxed);
130        }));
131        let mut context = crate::context::QueryContext::with("greptime", "public");
132        context.set_write_progress(progress.clone());
133        let fork = context.with_write_admission();
134        assert!(Arc::ptr_eq(fork.write_progress().unwrap(), &progress));
135        let wire = api::v1::QueryContext::from(context);
136        let received = crate::context::QueryContext::try_from(wire).unwrap();
137        assert!(received.write_progress().is_none());
138        progress.add_written(2);
139        progress.confirm(3);
140        progress.confirm(3);
141        progress.mark_accepted();
142        assert_eq!(progress.written(), 3);
143        assert_eq!(observed.load(Ordering::Relaxed), 3);
144        assert!(progress.accepted());
145    }
146}
session/write_progress.rs

session/
write_progress.rs
