zksync_object_store 0.1.0

ZKsync implementation of object stores
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! Mirroring object store.

use async_trait::async_trait;

use crate::{file::FileBackedObjectStore, raw::ObjectStore, Bucket, ObjectStoreError};

#[derive(Debug)]
pub(crate) struct MirroringObjectStore<S> {
    inner: S,
    mirror_store: FileBackedObjectStore,
}

impl<S: ObjectStore> MirroringObjectStore<S> {
    pub async fn new(inner: S, mirror_path: String) -> Result<Self, ObjectStoreError> {
        tracing::info!("Initializing mirroring for store {inner:?} at `{mirror_path}`");
        let mirror_store = FileBackedObjectStore::new(mirror_path).await?;
        Ok(Self {
            inner,
            mirror_store,
        })
    }
}

#[async_trait]
impl<S: ObjectStore> ObjectStore for MirroringObjectStore<S> {
    #[tracing::instrument(skip(self))]
    async fn get_raw(&self, bucket: Bucket, key: &str) -> Result<Vec<u8>, ObjectStoreError> {
        match self.mirror_store.get_raw(bucket, key).await {
            Ok(object) => {
                tracing::trace!("obtained object from mirror");
                return Ok(object);
            }
            Err(err) => {
                if !matches!(err, ObjectStoreError::KeyNotFound(_)) {
                    tracing::warn!(
                        "unexpected error calling local mirror store: {:#}",
                        anyhow::Error::from(err)
                    );
                }
                let object = self.inner.get_raw(bucket, key).await?;
                tracing::trace!("obtained object from underlying store");
                if let Err(err) = self.mirror_store.put_raw(bucket, key, object.clone()).await {
                    tracing::warn!("failed mirroring object: {:#}", anyhow::Error::from(err));
                } else {
                    tracing::trace!("mirrored object");
                }
                Ok(object)
            }
        }
    }

    #[tracing::instrument(skip(self, value), fields(value.len = value.len()))]
    async fn put_raw(
        &self,
        bucket: Bucket,
        key: &str,
        value: Vec<u8>,
    ) -> Result<(), ObjectStoreError> {
        self.inner.put_raw(bucket, key, value.clone()).await?;
        // Only put the value into the mirror once it has been put in the underlying store
        if let Err(err) = self.mirror_store.put_raw(bucket, key, value).await {
            tracing::warn!("failed mirroring object: {:#}", anyhow::Error::from(err));
        } else {
            tracing::trace!("mirrored object");
        }
        Ok(())
    }

    #[tracing::instrument(skip(self))]
    async fn remove_raw(&self, bucket: Bucket, key: &str) -> Result<(), ObjectStoreError> {
        self.inner.remove_raw(bucket, key).await?;
        // Only remove the value from the mirror once it has been removed in the underlying store
        if let Err(err) = self.mirror_store.remove_raw(bucket, key).await {
            tracing::warn!(
                "failed removing object from mirror: {:#}",
                anyhow::Error::from(err)
            );
        } else {
            tracing::trace!("removed object from mirror");
        }
        Ok(())
    }

    fn storage_prefix_raw(&self, bucket: Bucket) -> String {
        self.inner.storage_prefix_raw(bucket)
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use tempfile::TempDir;

    use super::*;
    use crate::MockObjectStore;

    #[tokio::test]
    async fn mirroring_basics() {
        let dir = TempDir::new().unwrap();
        let path = dir.into_path().into_os_string().into_string().unwrap();

        let mock_store = MockObjectStore::default();
        mock_store
            .put_raw(Bucket::StorageSnapshot, "test", vec![1, 2, 3])
            .await
            .unwrap();
        let mirroring_store = MirroringObjectStore::new(mock_store, path).await.unwrap();

        let object = mirroring_store
            .get_raw(Bucket::StorageSnapshot, "test")
            .await
            .unwrap();
        assert_eq!(object, [1, 2, 3]);
        // Check that the object got mirrored.
        let object_in_mirror = mirroring_store
            .mirror_store
            .get_raw(Bucket::StorageSnapshot, "test")
            .await
            .unwrap();
        assert_eq!(object_in_mirror, [1, 2, 3]);
        let object = mirroring_store
            .get_raw(Bucket::StorageSnapshot, "test")
            .await
            .unwrap();
        assert_eq!(object, [1, 2, 3]);

        let err = mirroring_store
            .get_raw(Bucket::StorageSnapshot, "missing")
            .await
            .unwrap_err();
        assert_matches!(err, ObjectStoreError::KeyNotFound(_));

        mirroring_store
            .put_raw(Bucket::StorageSnapshot, "other", vec![3, 2, 1])
            .await
            .unwrap();
        // Check that the object got mirrored.
        let object_in_mirror = mirroring_store
            .mirror_store
            .get_raw(Bucket::StorageSnapshot, "other")
            .await
            .unwrap();
        assert_eq!(object_in_mirror, [3, 2, 1]);
        let object = mirroring_store
            .get_raw(Bucket::StorageSnapshot, "other")
            .await
            .unwrap();
        assert_eq!(object, [3, 2, 1]);
    }
}