mz_persist/
intercept.rs

Help
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//! Test utilities for trapping and injecting responses in external storage.
11
12use std::fmt::Debug;
13use std::sync::{Arc, Mutex};
14
15use async_trait::async_trait;
16use bytes::Bytes;
17use mz_ore::bytes::SegmentedBytes;
18
19use crate::location::{Blob, BlobMetadata, ExternalError};
20
21/// Post-op closure for [Blob::delete].
22pub type PostDeleteFn = Arc<
23    dyn Fn(&str, Result<Option<usize>, ExternalError>) -> Result<Option<usize>, ExternalError>
24        + Send
25        + Sync,
26>;
27
28#[derive(Default)]
29struct InterceptCore {
30    post_delete: Option<PostDeleteFn>,
31}
32
33impl Debug for InterceptCore {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        let InterceptCore { post_delete } = self;
36        f.debug_struct("InterceptCore")
37            .field(
38                "post_delete",
39                &post_delete.as_ref().map(|x| format!("{:p}", x)),
40            )
41            .finish()
42    }
43}
44
45/// A handle for controlling the behavior of an intercept delegate.
46#[derive(Clone, Debug, Default)]
47pub struct InterceptHandle {
48    core: Arc<Mutex<InterceptCore>>,
49}
50
51impl InterceptHandle {
52    /// Sets a new [PostDeleteFn].
53    ///
54    /// Returns the previous closure, if any.
55    pub fn set_post_delete(&self, f: Option<PostDeleteFn>) -> Option<PostDeleteFn> {
56        let mut core = self.core.lock().expect("lock should not be poisoned");
57        std::mem::replace(&mut core.post_delete, f)
58    }
59}
60
61/// An intercept delegate to [Blob].
62///
63/// TODO: Tune this pattern to be most useful and then extend it to consensus
64/// and the rest of blob.
65#[derive(Debug)]
66pub struct InterceptBlob {
67    handle: InterceptHandle,
68    blob: Arc<dyn Blob>,
69}
70
71impl InterceptBlob {
72    /// Returns a new [InterceptBlob].
73    pub fn new(blob: Arc<dyn Blob>, handle: InterceptHandle) -> Self {
74        InterceptBlob { handle, blob }
75    }
76}
77
78#[async_trait]
79impl Blob for InterceptBlob {
80    async fn get(&self, key: &str) -> Result<Option<SegmentedBytes>, ExternalError> {
81        self.blob.get(key).await
82    }
83
84    async fn list_keys_and_metadata(
85        &self,
86        key_prefix: &str,
87        f: &mut (dyn FnMut(BlobMetadata) + Send + Sync),
88    ) -> Result<(), ExternalError> {
89        self.blob.list_keys_and_metadata(key_prefix, f).await
90    }
91
92    async fn set(&self, key: &str, value: Bytes) -> Result<(), ExternalError> {
93        self.blob.set(key, value).await
94    }
95
96    async fn delete(&self, key: &str) -> Result<Option<usize>, ExternalError> {
97        let ret = self.blob.delete(key).await;
98        let post_delete = self
99            .handle
100            .core
101            .lock()
102            .expect("lock should not be poisoned")
103            .post_delete
104            .clone();
105        match post_delete {
106            Some(x) => x(key, ret),
107            None => ret,
108        }
109    }
110
111    async fn restore(&self, key: &str) -> Result<(), ExternalError> {
112        self.blob.restore(key).await
113    }
114}
mz_persist/intercept.rs

mz_persist/
intercept.rs
