support fallback blobstore (#707)

Netflix · Sep 6, 2024 · fa55cb1 · fa55cb1
1 parent 46e8c23
commit fa55cb1
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 0 deletions.
diff --git a/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/BlobStore.java b/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/BlobStore.java
@@ -23,6 +23,7 @@
 import java.util.concurrent.locks.Lock;
 import lombok.AccessLevel;
 import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 import net.lingala.zip4j.ZipFile;
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.io.FilenameUtils;
@@ -43,6 +44,10 @@ default BlobStore withPrefix(URI prefixUri) {
         return new PrefixedBlobStore(prefixUri, this);
     }
 
+    default BlobStore withFallbackStore(BlobStore fallbackStore) {
+        return new FallbackEnabledBlobStore(this, fallbackStore);
+    }
+
     /**
      * blob store that when downloading zip files, also unpacks them and returns the unpacked file/directory to the caller.
      *
@@ -67,6 +72,23 @@ static BlobStore forHadoopFileSystem(URI clusterStoragePath, File localStoreDir)
                 .withThreadSafeBlobStore();
     }
 
+    static BlobStore forHadoopFileSystem(URI clusterStoragePath, URI fallbackStoragePath, File localStoreDir)
+        throws Exception {
+        final org.apache.hadoop.fs.FileSystem fileSystem =
+            FileSystemInitializer.create(clusterStoragePath);
+
+        final org.apache.hadoop.fs.FileSystem fallbackFileSystem =
+            FileSystemInitializer.create(fallbackStoragePath);
+
+        return
+            new HadoopFileSystemBlobStore(fileSystem, localStoreDir)
+                .withPrefix(clusterStoragePath)
+                .withFallbackStore(
+                    new HadoopFileSystemBlobStore(fallbackFileSystem, localStoreDir).withPrefix(fallbackStoragePath))
+                .withZipCapabilities()
+                .withThreadSafeBlobStore();
+    }
+
     @RequiredArgsConstructor(access = AccessLevel.PACKAGE)
     class PrefixedBlobStore implements BlobStore {
         private final URI rootUri;
@@ -84,6 +106,32 @@ public void close() throws IOException {
         }
     }
 
+    @RequiredArgsConstructor(access = AccessLevel.PACKAGE)
+    @Slf4j
+    class FallbackEnabledBlobStore implements BlobStore {
+        private final BlobStore blobStore;
+        private final BlobStore fallbackBlobStore;
+
+        @Override
+        public File get(URI blobUrl) throws IOException {
+            try
+            {
+                return blobStore.get(blobUrl);
+            }
+            catch (IOException e) {
+                log.error("Get blob error, fallback to next blobstore", e);
+            }
+
+            return fallbackBlobStore.get(blobUrl);
+        }
+
+        @Override
+        public void close() throws IOException {
+            blobStore.close();
+            fallbackBlobStore.close();
+        }
+    }
+
     @RequiredArgsConstructor(access = AccessLevel.PACKAGE)
     class ZipHandlingBlobStore implements BlobStore {
 

diff --git a/mantis-server/mantis-server-agent/src/test/java/io/mantisrx/server/agent/BlobStoreTest.java b/mantis-server/mantis-server-agent/src/test/java/io/mantisrx/server/agent/BlobStoreTest.java
@@ -23,6 +23,7 @@
 import static org.mockito.Mockito.when;
 
 import java.io.File;
+import java.io.IOException;
 import java.net.URI;
 import org.junit.Test;
 import org.mockito.Matchers;
@@ -46,4 +47,36 @@ public void testPrefixedBlobStore() throws Exception {
             "https://mantisrx.region.prod.io.net/mantis-artifacts/sananthanarayanan-mantis-jobs-sine-function-thin-0.1.0.zip"));
         verify(blobStore, times(2)).get(Matchers.eq(expectedUri));
     }
+
+    @Test
+    public void testFallbackBlobStore() throws Exception {
+        final BlobStore blobStore = mock(BlobStore.class);
+        final BlobStore fallbackBlobStore = mock(BlobStore.class);
+        final File file = mock(File.class);
+        when(fallbackBlobStore.get(any())).thenReturn(file);
+        when(blobStore.get(any())).thenThrow(new IOException("mock io err")).thenReturn(file);
+
+        final BlobStore prefixedBlobStore =
+            new BlobStore.PrefixedBlobStore(new URI("s3://mantisrx.s3.store/mantis/jobs/"), blobStore);
+
+        final BlobStore fallbackPrefixedBlobStore =
+            new BlobStore.PrefixedBlobStore(new URI("s3://mantisrx.s3.store.fallback/mantis/jobs/"), fallbackBlobStore);
+
+        final BlobStore fallbackEnabledBlobStore = prefixedBlobStore.withFallbackStore(fallbackPrefixedBlobStore);
+        fallbackEnabledBlobStore.get(new URI("http://sananthanarayanan-mantis-jobs-sine-function-thin-0.1.0.zip"));
+
+        final URI expectedUri =
+            new URI("s3://mantisrx.s3.store/mantis/jobs/sananthanarayanan-mantis-jobs-sine-function-thin-0.1.0.zip");
+        verify(blobStore, times(1)).get(Matchers.eq(expectedUri));
+
+        final URI expectedFallbackUri =
+            new URI("s3://mantisrx.s3.store.fallback/mantis/jobs/sananthanarayanan-mantis-jobs-sine-function-thin-0.1"
+                + ".0.zip");
+        verify(fallbackBlobStore, times(1)).get(Matchers.eq(expectedFallbackUri));
+
+        // test non-fallback path
+        fallbackEnabledBlobStore.get(new URI("http://sananthanarayanan-mantis-jobs-sine-function-thin-0.1.0.zip"));
+        verify(blobStore, times(2)).get(Matchers.eq(expectedUri));
+        verify(fallbackBlobStore, times(1)).get(Matchers.eq(expectedFallbackUri));
+    }
 }