mz_adapter/
frontend_transaction.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Completion of read-only transactions whose resources belong to the session.
11
12use std::time::{Duration, Instant};
13
14use mz_repr::Timestamp;
15use mz_sql::session::vars::{EndTransactionAction, IsolationLevel};
16use mz_timestamp_oracle::TimestampOracle;
17use tokio::sync::watch;
18
19use crate::session::{ReadOnlyCompletion, Session};
20use crate::{AdapterError, ExecuteResponse, PeekClient};
21
22pub(crate) async fn complete_read_only_transaction(
23    session: &mut Session,
24    client: &mut PeekClient,
25    cancel: &mut watch::Receiver<()>,
26    completion: ReadOnlyCompletion,
27) -> Result<ExecuteResponse, AdapterError> {
28    // Cancellation while idle must not poison a later completion. Keep the session
29    // intact across awaits so dropping this future still leaves teardown with its owner.
30    cancel.borrow_and_update();
31    client
32        .coordinator_client()
33        .metrics()
34        .frontend_transaction_completions
35        .inc();
36    let result = async {
37        session.apply_external_metadata_updates();
38        let catalog = client.catalog_snapshot("end_transaction").await;
39        mz_sql::rbac::check_session_roles(&catalog.for_session(session), session)?;
40
41        if let Some(context) = &completion.timestamp {
42            match session.vars().transaction_isolation() {
43                IsolationLevel::StrictSerializable => {
44                    let started = Instant::now();
45                    let immediately_handled =
46                        if let Some((timeline, timestamp)) = context.timestamp_to_linearize() {
47                            client
48                                .coordinator_client()
49                                .metrics()
50                                .frontend_transaction_waits
51                                .inc();
52                            let oracle = client.ensure_oracle(timeline.clone()).await?;
53                            wait_for_read_timestamp(&**oracle, *timestamp, cancel).await?
54                        } else {
55                            client
56                                .coordinator_client()
57                                .metrics()
58                                .linearize_read_fast_path
59                                .inc();
60                            true
61                        };
62                    let metrics = client.coordinator_client().metrics();
63                    let duration = if immediately_handled {
64                        &metrics.linearize_read_immediate_seconds
65                    } else {
66                        &metrics.linearize_read_delayed_seconds
67                    };
68                    duration.observe(started.elapsed().as_secs_f64());
69                }
70                IsolationLevel::StrongSessionSerializable => {
71                    if let Some((timeline, timestamp)) = context.timeline_timestamp() {
72                        session
73                            .ensure_timestamp_oracle(timeline.clone())
74                            .apply_write(*timestamp);
75                    }
76                }
77                IsolationLevel::ReadUncommitted
78                | IsolationLevel::ReadCommitted
79                | IsolationLevel::RepeatableRead
80                | IsolationLevel::Serializable
81                | IsolationLevel::BoundedStaleness(_) => {}
82            }
83        }
84        Ok::<_, AdapterError>(())
85    }
86    .await;
87
88    let action = if result.is_ok() {
89        completion.action
90    } else {
91        EndTransactionAction::Rollback
92    };
93    let _ = session.clear_transaction();
94    let params = session.vars_mut().end_transaction(action);
95    result?;
96    Ok(match action {
97        EndTransactionAction::Commit => ExecuteResponse::TransactionCommitted { params },
98        EndTransactionAction::Rollback => ExecuteResponse::TransactionRolledBack { params },
99    })
100}
101
102async fn wait_for_read_timestamp(
103    oracle: &(dyn TimestampOracle<Timestamp> + Send + Sync),
104    timestamp: Timestamp,
105    cancel: &mut watch::Receiver<()>,
106) -> Result<bool, AdapterError> {
107    let wait = async {
108        let mut immediately_handled = true;
109        loop {
110            #[cfg(debug_assertions)]
111            {
112                fail::fail_point!("frontend_transaction_before_oracle_read");
113                while fail::eval("frontend_transaction_hold_oracle_read", |_| ()).is_some() {
114                    tokio::time::sleep(Duration::from_millis(10)).await;
115                }
116            }
117            let current = oracle.read_ts().await;
118            if timestamp <= current {
119                break immediately_handled;
120            }
121            immediately_handled = false;
122            let delay = Duration::from_millis(timestamp.saturating_sub(current).into());
123            tokio::time::sleep(delay.min(Duration::from_secs(1))).await;
124        }
125    };
126    tokio::select! {
127        immediately_handled = wait => Ok(immediately_handled),
128        _ = cancel.changed() => Err(AdapterError::Canceled),
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
136
137    use async_trait::async_trait;
138    use futures::FutureExt;
139    use mz_timestamp_oracle::WriteTimestamp;
140
141    #[derive(Debug, Default)]
142    struct ReadOracle {
143        timestamp: AtomicU64,
144        calls: AtomicUsize,
145    }
146
147    #[async_trait]
148    impl TimestampOracle<Timestamp> for ReadOracle {
149        async fn read_ts(&self) -> Timestamp {
150            self.calls.fetch_add(1, Ordering::SeqCst);
151            self.timestamp.load(Ordering::SeqCst).into()
152        }
153        async fn write_ts(&self) -> WriteTimestamp {
154            unreachable!("read-only completion")
155        }
156        async fn peek_write_ts(&self) -> Timestamp {
157            unreachable!("must use a linearized read")
158        }
159        async fn apply_write(&self, _: Timestamp) {
160            unreachable!("read-only completion")
161        }
162    }
163
164    #[mz_ore::test(tokio::test(start_paused = true))]
165    async fn completion_waits_for_the_shared_oracle() {
166        let oracle = ReadOracle::default();
167        let (_tx, mut cancel) = watch::channel(());
168        let wait = wait_for_read_timestamp(&oracle, 10.into(), &mut cancel);
169        tokio::pin!(wait);
170        assert!(wait.as_mut().now_or_never().is_none());
171        assert_eq!(oracle.calls.load(Ordering::SeqCst), 1);
172        tokio::time::advance(Duration::from_millis(10)).await;
173        assert!(
174            wait.as_mut().now_or_never().is_none(),
175            "wall time alone cannot release the read"
176        );
177        oracle.timestamp.store(10, Ordering::SeqCst);
178        tokio::time::advance(Duration::from_millis(10)).await;
179        assert!(!wait.await.expect("oracle advancement completes the read"));
180    }
181
182    #[mz_ore::test(tokio::test)]
183    async fn completion_wait_is_cancellable_and_old_cancels_are_ignored() {
184        let oracle = ReadOracle::default();
185        let (tx, mut cancel) = watch::channel(());
186        tx.send_replace(());
187        cancel.borrow_and_update();
188        {
189            let wait = wait_for_read_timestamp(&oracle, 10.into(), &mut cancel);
190            tokio::pin!(wait);
191            assert!(wait.as_mut().now_or_never().is_none());
192            tx.send_replace(());
193            assert!(matches!(wait.await, Err(AdapterError::Canceled)));
194        }
195        cancel.borrow_and_update();
196        oracle.timestamp.store(10, Ordering::SeqCst);
197        assert!(
198            wait_for_read_timestamp(&oracle, 10.into(), &mut cancel)
199                .await
200                .is_ok()
201        );
202    }
203
204    #[mz_ore::test(tokio::test)]
205    async fn completion_wait_ends_when_the_connection_is_gone() {
206        let oracle = ReadOracle::default();
207        let (tx, mut cancel) = watch::channel(());
208        drop(tx);
209        assert!(matches!(
210            wait_for_read_timestamp(&oracle, 10.into(), &mut cancel).await,
211            Err(AdapterError::Canceled)
212        ));
213    }
214}
mz_adapter/frontend_transaction.rs

mz_adapter/
frontend_transaction.rs
