datanode/
store.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//! object storage utilities
16
17use std::time::{Duration, SystemTime, UNIX_EPOCH};
18
19use common_telemetry::{info, warn};
20use futures::TryStreamExt;
21use object_store::factory::new_raw_object_store;
22use object_store::util::{clean_temp_dir, join_dir, with_instrument_layers, with_retry_layers};
23use object_store::{ATOMIC_WRITE_DIR, Error as ObjectStoreError, ErrorKind, ObjectStore};
24use snafu::{IntoError, prelude::*};
25
26use crate::config::ObjectStoreConfig;
27use crate::error::{self, Result};
28
29/// Cleans up old LRU read cache directories that were removed.
30fn clean_old_read_cache(store: &ObjectStoreConfig, data_home: &str) {
31    if !store.is_object_storage() {
32        return;
33    }
34
35    let Some(cache_config) = store.cache_config() else {
36        return;
37    };
38
39    // Only cleans if read cache was enabled
40    if !cache_config.enable_read_cache {
41        return;
42    }
43
44    let cache_base_dir = if cache_config.cache_path.is_empty() {
45        data_home
46    } else {
47        &cache_config.cache_path
48    };
49
50    // Cleans up the old read cache directory
51    let old_read_cache_dir = join_dir(cache_base_dir, "cache/object/read");
52    info!(
53        "Cleaning up old read cache directory: {}",
54        old_read_cache_dir
55    );
56    if let Err(e) = clean_temp_dir(&old_read_cache_dir) {
57        warn!(e; "Failed to clean old read cache directory {}", old_read_cache_dir);
58    }
59
60    // Cleans up the atomic temp dir used by the cache layer
61    let cache_atomic_temp_dir = join_dir(cache_base_dir, ATOMIC_WRITE_DIR);
62    info!(
63        "Cleaning up old cache atomic temp directory: {}",
64        cache_atomic_temp_dir
65    );
66    if let Err(e) = clean_temp_dir(&cache_atomic_temp_dir) {
67        warn!(e; "Failed to clean old cache atomic temp directory {}", cache_atomic_temp_dir);
68    }
69}
70
71pub async fn new_object_store(store: ObjectStoreConfig, data_home: &str) -> Result<ObjectStore> {
72    // Cleans up old LRU read cache directories.
73    // TODO: Remove this line after the 1.0 release.
74    clean_old_read_cache(&store, data_home);
75
76    let object_store = new_raw_object_store(&store, data_home)
77        .await
78        .context(error::ObjectStoreSnafu)?;
79    // Enables retry layer for non-fs object storages
80    let object_store = if store.is_object_storage() {
81        // Adds retry layer
82        with_retry_layers(object_store)
83    } else {
84        object_store
85    };
86
87    let object_store = with_instrument_layers(object_store, true);
88    Ok(object_store)
89}
90
91/// The content written by [`probe_object_store`].
92const PROBE_CONTENT: &[u8] = b"greptimedb object store probe";
93
94/// The content appended by [`probe_object_store`] to exercise read-modify-write.
95const PROBE_SUFFIX: &[u8] = b" - read-modify-write";
96
97/// The directory that holds the objects created by [`probe_object_store`].
98const PROBE_DIR: &str = ".greptime_probe";
99
100/// The budget of one [`probe_object_store`] call, including the retries an object
101/// store makes on top of its 30s request timeout.
102pub(crate) const PROBE_TIMEOUT: Duration = Duration::from_secs(180);
103
104/// The maximum random delay before a probe starts. It is short enough that startup
105/// waits at most this delay, and it spreads out the probes of providers that share
106/// an object storage.
107const PROBE_STAGGER: Duration = Duration::from_millis(500);
108
109/// The budget of the best-effort cleanup after a timed-out probe, kept short so an
110/// unresponsive store cannot block startup a second time.
111const PROBE_CLEANUP_TIMEOUT: Duration = Duration::from_secs(10);
112
113/// Returns a random delay in `[0, PROBE_STAGGER)`.
114pub(crate) fn probe_start_delay() -> Duration {
115    Duration::from_nanos(rand::random_range(0..PROBE_STAGGER.as_nanos() as u64))
116}
117
118/// Returns a unique key prefix for [`probe_object_store`], so datanodes sharing
119/// the same bucket and root do not collide.
120pub(crate) fn probe_key_prefix(node_id: u64) -> String {
121    let nanos = SystemTime::now()
122        .duration_since(UNIX_EPOCH)
123        .unwrap_or_default()
124        .as_nanos();
125    format!("{PROBE_DIR}/{node_id}-{}-{nanos}", std::process::id())
126}
127
128/// Returns the object key [`probe_object_store`] probes under `prefix`.
129fn probe_object_key(prefix: &str) -> String {
130    format!("{prefix}/probe")
131}
132
133/// Checks that `store` supports the operations the storage engines rely on: write,
134/// read, overwrite, list and delete.
135///
136/// The probe object is removed before returning, after a failed step too. A caller
137/// that abandons the check early must call [`remove_probe_object`].
138pub(crate) async fn probe_object_store(
139    store: &ObjectStore,
140    name: &str,
141    prefix: &str,
142) -> Result<()> {
143    let key = probe_object_key(prefix);
144
145    let check = async {
146        store
147            .write(&key, PROBE_CONTENT)
148            .await
149            .context(error::ObjectStoreProbeSnafu {
150                name,
151                step: "write",
152            })?;
153        verify_content(store, name, &key, PROBE_CONTENT).await?;
154
155        // Object stores cannot append in place, so this reads the object, appends to
156        // it and overwrites it, exercising an overwrite of an existing object.
157        let mut content = store
158            .read(&key)
159            .await
160            .context(error::ObjectStoreProbeSnafu { name, step: "read" })?
161            .to_vec();
162        content.extend_from_slice(PROBE_SUFFIX);
163        store
164            .write(&key, content.clone())
165            .await
166            .context(error::ObjectStoreProbeSnafu {
167                name,
168                step: "overwrite",
169            })?;
170        verify_content(store, name, &key, &content).await?;
171
172        let mut lister = store
173            .lister_with(&format!("{prefix}/"))
174            .await
175            .context(error::ObjectStoreProbeSnafu { name, step: "list" })?;
176        while let Some(entry) = lister
177            .try_next()
178            .await
179            .context(error::ObjectStoreProbeSnafu { name, step: "list" })?
180        {
181            if entry.path() == key {
182                return Ok(());
183            }
184        }
185
186        Err(
187            error::ObjectStoreProbeSnafu { name, step: "list" }.into_error(ObjectStoreError::new(
188                ErrorKind::Unexpected,
189                format!("'{key}' is not visible in the listing"),
190            )),
191        )
192    }
193    .await;
194
195    match store.delete(&key).await {
196        Ok(()) => check,
197        Err(err) if check.is_err() => {
198            warn!(
199                err;
200                "Failed to remove the object storage probe object '{}' of provider '{}'",
201                key,
202                name
203            );
204            check
205        }
206        Err(err) => Err(err).context(error::ObjectStoreProbeSnafu {
207            name,
208            step: "delete",
209        }),
210    }
211}
212
213async fn verify_content(store: &ObjectStore, name: &str, key: &str, expected: &[u8]) -> Result<()> {
214    let got = store
215        .read(key)
216        .await
217        .context(error::ObjectStoreProbeSnafu { name, step: "read" })?;
218    if got.to_vec() != expected {
219        return Err(
220            error::ObjectStoreProbeSnafu { name, step: "read" }.into_error(ObjectStoreError::new(
221                ErrorKind::Unexpected,
222                format!("content of '{key}' differs from what was written"),
223            )),
224        );
225    }
226    Ok(())
227}
228
229/// Removes the probe object of a probe that did not finish.
230///
231/// It has its own short budget, and failures only log: the caller reports the timeout.
232pub(crate) async fn remove_probe_object(store: &ObjectStore, prefix: &str) {
233    let key = probe_object_key(prefix);
234    match tokio::time::timeout(PROBE_CLEANUP_TIMEOUT, store.delete(&key)).await {
235        Ok(Ok(())) => {}
236        Ok(Err(err)) => warn!(
237            err;
238            "Failed to remove the object storage probe object '{}'", key
239        ),
240        Err(_) => warn!(
241            "Timed out removing the object storage probe object '{}'",
242            key
243        ),
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use common_test_util::temp_dir::create_temp_dir;
250    use object_store::services::Fs;
251
252    use super::*;
253
254    fn new_fs_store(root: &str) -> ObjectStore {
255        ObjectStore::new(Fs::default().root(root)).unwrap()
256    }
257
258    #[tokio::test]
259    async fn test_probe_object_store() {
260        let dir = create_temp_dir("object-store-probe");
261        let store = new_fs_store(dir.path().to_str().unwrap());
262        let prefix = probe_key_prefix(1);
263
264        probe_object_store(&store, "fs", &prefix).await.unwrap();
265
266        // The lister also returns the listed directory itself, so only check that the
267        // probe object is gone.
268        let key = format!("{prefix}/probe");
269        let mut lister = store.lister_with(&format!("{prefix}/")).await.unwrap();
270        while let Some(entry) = lister.try_next().await.unwrap() {
271            assert_ne!(entry.path(), key);
272        }
273    }
274
275    #[tokio::test]
276    async fn test_probe_object_store_write_failure() {
277        let dir = create_temp_dir("object-store-probe-write-failure");
278        let root = dir.path().join("not-a-directory");
279        std::fs::write(&root, b"").unwrap();
280        let store = new_fs_store(root.to_str().unwrap());
281
282        let err = probe_object_store(&store, "broken", &probe_key_prefix(1))
283            .await
284            .unwrap_err();
285        assert!(
286            err.to_string()
287                .contains("probe for provider 'broken' at step 'write'")
288        );
289    }
290
291    #[tokio::test]
292    async fn test_remove_probe_object() {
293        let dir = create_temp_dir("object-store-probe-cleanup");
294        let store = new_fs_store(dir.path().to_str().unwrap());
295        let prefix = probe_key_prefix(1);
296        let key = probe_object_key(&prefix);
297        store.write(&key, PROBE_CONTENT).await.unwrap();
298
299        remove_probe_object(&store, &prefix).await;
300
301        assert!(store.read(&key).await.is_err());
302    }
303}
datanode/store.rs

datanode/
store.rs
