fix(controller): oss list all keys

star-whale · Sep 20, 2022 · 175ba76 · 175ba76
1 parent b7a536c
commit 175ba76
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 8 deletions.
diff --git a/...ess-layer/src/main/java/ai/starwhale/mlops/storage/aliyun/StorageAccessServiceAliyun.java b/...ess-layer/src/main/java/ai/starwhale/mlops/storage/aliyun/StorageAccessServiceAliyun.java
@@ -27,11 +27,15 @@
 import com.aliyun.oss.OSSException;
 import com.aliyun.oss.model.GetObjectRequest;
 import com.aliyun.oss.model.HeadObjectRequest;
+import com.aliyun.oss.model.ListObjectsRequest;
 import com.aliyun.oss.model.OSSObjectSummary;
+import com.aliyun.oss.model.ObjectListing;
 import com.aliyun.oss.model.ObjectMetadata;
+import com.google.common.collect.Streams;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.ArrayList;
 import java.util.stream.Stream;
 
 public class StorageAccessServiceAliyun implements StorageAccessService {
@@ -85,12 +89,19 @@ public LengthAbleInputStream get(String path, Long offset, Long size) throws IOE
 
     @Override
     public Stream<String> list(String path) throws IOException {
+        Stream<String> files = Stream.empty();
         try {
-            var resp = this.ossClient.listObjects(this.bucket, path);
-            return resp.getObjectSummaries().stream().map(OSSObjectSummary::getKey);
+            var req = new ListObjectsRequest(this.bucket).withPrefix(path);
+            ObjectListing resp;
+            do {
+                resp = this.ossClient.listObjects(req);
+                files = Streams.concat(files, resp.getObjectSummaries().stream().map(OSSObjectSummary::getKey));
+                req.setMarker(resp.getNextMarker());
+            } while (resp.isTruncated());
+            return files;
         } catch (OSSException e) {
             if (e.getErrorCode().equals(OSSErrorCode.NO_SUCH_KEY)) {
-                return Stream.empty();
+                return files;
             }
             throw e;
         }

diff --git a/...rage-access-layer/src/main/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3.java b/...rage-access-layer/src/main/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3.java
@@ -20,6 +20,7 @@
 import ai.starwhale.mlops.storage.StorageAccessService;
 import ai.starwhale.mlops.storage.StorageObjectInfo;
 import ai.starwhale.mlops.storage.util.MetaHelper;
+import com.google.common.collect.Streams;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
@@ -193,12 +194,18 @@ public LengthAbleInputStream get(String path, Long offset, Long size) throws IOE
 
     @Override
     public Stream<String> list(String path) {
+        Stream<String> files = Stream.empty();
         try {
-            final ListObjectsResponse listObjectsResponse = s3client.listObjects(
-                    ListObjectsRequest.builder().bucket(s3Config.getBucket()).prefix(path).build());
-            return listObjectsResponse.contents().stream().map(S3Object::key);
+            ListObjectsResponse resp;
+            var reqBuilder = ListObjectsRequest.builder().bucket(s3Config.getBucket()).prefix(path);
+            do {
+                resp = s3client.listObjects(reqBuilder.build());
+                files = Streams.concat(files, resp.contents().stream().map(S3Object::key));
+                reqBuilder.marker(resp.nextMarker());
+            } while (resp.isTruncated());
+            return files;
         } catch (NoSuchKeyException e) {
-            return Stream.empty();
+            return files;
         }
     }
 

diff --git a/...-access-layer/src/test/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3Test.java b/...-access-layer/src/test/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3Test.java
@@ -27,6 +27,7 @@
 import java.io.IOException;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.stream.Collectors;
 import org.junit.jupiter.api.BeforeEach;
@@ -48,7 +49,7 @@ public class StorageAccessServiceS3Test {
     private static final S3MockContainer s3Mock =
             new S3MockContainer(System.getProperty("s3mock.version", "latest"))
                     .withValidKmsKeys("arn:aws:kms:us-east-1:1234567890:key/valid-test-key-ref")
-                    .withInitialBuckets("test");
+                    .withInitialBuckets("test,huge");
 
     private StorageAccessServiceS3 s3;
     private S3Client client;
@@ -161,4 +162,35 @@ public void testDelete() throws IOException {
         this.s3.delete("x");
         assertThat(this.s3.list("").collect(Collectors.toList()), containsInAnyOrder("t1", "t2", "t/1", "t/2"));
     }
+
+    @Test
+    public void testListHugeNumberOfFiles() {
+        var s3 = new StorageAccessServiceS3(
+                S3Config.builder()
+                        .bucket("huge")
+                        .accessKey("ak")
+                        .secretKey("sk")
+                        .region("us-west-1")
+                        .endpoint(s3Mock.getHttpEndpoint())
+                        .hugeFileThreshold(10 * 1024 * 1024)
+                        .hugeFilePartSize(5 * 1024 * 1024)
+                        .build());
+
+        final String prefix = "huge-number";
+        // https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjectsV2.html#API_ListObjectsV2_RequestSyntax
+        // By default the action returns up to 1,000 key names.
+        // The response might contain fewer keys but will never contain more
+        final int count = 1001;
+
+        var expectedFiles = new ArrayList<String>();
+        for (int i = 0; i < count; i++) {
+            var path = prefix + "/" + i;
+            expectedFiles.add(path);
+            s3.put(path, new byte[] {7});
+        }
+
+        var resp = s3.list(prefix + "/").collect(Collectors.toList());
+        assertThat(resp.size(), is(count));
+        assertThat(resp, containsInAnyOrder(expectedFiles.toArray()));
+    }
 }