From 7ce7cb4248ed2cb4bb774858c4037582ebc48500 Mon Sep 17 00:00:00 2001
From: Paul Masurel <paul@quickwit.io>
Date: Mon, 21 Aug 2023 23:00:52 +0900
Subject: [PATCH] Simplification

---
 .../src/cached_splits_registry.rs             | 29 +++++++++++++++----
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/quickwit/quickwit-storage/src/cached_splits_registry.rs b/quickwit/quickwit-storage/src/cached_splits_registry.rs
index 6de44f4e78e..0ace0042b50 100644
--- a/quickwit/quickwit-storage/src/cached_splits_registry.rs
+++ b/quickwit/quickwit-storage/src/cached_splits_registry.rs
@@ -54,6 +54,7 @@ struct InnerCachedSplitRegistry {
     storage_config: CacheStorageConfig,
     counters: AtomicCacheStorageCounters,
     splits: RwLock<HashMap<String, SplitInfo>>,
+    cache_storage: tokio::sync::OnceCell<Option<Arc<dyn Storage>>>,
 }
 
 impl CachedSplitRegistry {
@@ -63,6 +64,7 @@ impl CachedSplitRegistry {
                 storage_config,
                 counters: Default::default(),
                 splits: RwLock::new(HashMap::new()),
+                cache_storage: Default::default(),
             }),
         }
     }
@@ -71,6 +73,23 @@ impl CachedSplitRegistry {
         &self.inner.counters
     }
 
+    // TODO why `Option`
+    async fn cache_storage(&self, storage_resolver: &StorageResolver) -> Option<Arc<dyn Storage>> {
+        self.inner.cache_storage.get_or_init(|| async {
+            let cache_uri = self.inner.storage_config.cache_uri().clone()?;
+            match storage_resolver.resolve(&cache_uri).await {
+                Ok(cache) => {
+                    Some(cache)
+                }
+                Err(err) => {
+                    error!("Failed to resolve cache storage. {:?}", err);
+                    None
+                }
+            }
+        }).await.clone()
+
+    }
+
     #[allow(dead_code)]
     pub async fn load(
         &self,
@@ -165,10 +184,9 @@ impl CachedSplitRegistry {
         output_path: &Path,
     ) -> anyhow::Result<()> {
         // TODO: Figure out an ealier way to handle these issues
-        if let Some(cache_uri) = self.inner.storage_config.cache_uri().clone() {
+        if let Some(cache_storage) = self.cache_storage(storage_resolver).await {
             let storage = storage_resolver.resolve(storage_uri).await?;
-            let cache = storage_resolver.resolve(&cache_uri).await?;
-            storage.copy_to_storage(path, cache, output_path).await?;
+            storage.copy_to_storage(path, cache_storage, output_path).await?;
         }
         Ok(())
     }
@@ -218,9 +236,8 @@ impl CachedSplitRegistry {
     ) -> anyhow::Result<()> {
         // TODO: This should've been handle earlier, but we need have a more graceful way of dealing
         // with possible issues
-        if let Some(cache_uri) = self.inner.storage_config.cache_uri() {
-            let cache = storage_resolver.resolve(&cache_uri).await?;
-            cache.delete(path).await?;
+        if let Some(cache_storage) = self.cache_storage(storage_resolver).await {
+            cache_storage.delete(path).await?;
         }
         Ok(())
     }