bugfix(controller): override contentLength in InputStreamResource ret…

…urned by ModelService.getFileData so that the underlying InputStream won't be read twice (#2409)
star-whale · Jun 26, 2023 · b257518 · b257518
1 parent ffeccdb
commit b257518
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/server/controller/src/main/java/ai/starwhale/mlops/api/ModelController.java b/server/controller/src/main/java/ai/starwhale/mlops/api/ModelController.java
@@ -254,7 +254,12 @@ public ResponseEntity<ResponseMessage<ListFilesResult>> listFiles(
     @Override
     public ResponseEntity<InputStreamResource> getFileData(String project, String model, String version, String path) {
         var result = this.modelService.getFileData(project, model, version, path);
-        return ResponseEntity.ok(new InputStreamResource(result));
+        return ResponseEntity.ok(new InputStreamResource(result) {
+            @Override
+            public long contentLength() {
+                return result.getSize();
+            }
+        });
     }
 
     @Override

diff --git a/server/controller/src/main/java/ai/starwhale/mlops/domain/model/ModelService.java b/server/controller/src/main/java/ai/starwhale/mlops/domain/model/ModelService.java
@@ -73,6 +73,7 @@
 import ai.starwhale.mlops.exception.SwValidationException;
 import ai.starwhale.mlops.exception.SwValidationException.ValidSubject;
 import ai.starwhale.mlops.exception.api.StarwhaleApiException;
+import ai.starwhale.mlops.storage.LengthAbleInputStream;
 import cn.hutool.core.util.StrUtil;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.github.pagehelper.PageHelper;
@@ -602,20 +603,21 @@ private FileNode convert(ModelPackageStorage.File file) {
                 .build();
     }
 
-    public InputStream getFileData(String project, String model, String version, String path) {
+    public LengthAbleInputStream getFileData(String project, String model, String version, String path) {
         return this.getFileData(this.getModelMetaBlob(project, model, version, ""), path);
     }
 
-    private InputStream getFileData(ModelPackageStorage.MetaBlob metaBlob, String path) {
+    private LengthAbleInputStream getFileData(ModelPackageStorage.MetaBlob metaBlob, String path) {
         var files = new LinkedList<>(this.getFile(metaBlob, path));
-        if (files.get(0).getSize() == 0) {
-            return new ByteArrayInputStream(new byte[0]);
+        var size = files.get(0).getSize();
+        if (size == 0) {
+            return new LengthAbleInputStream(new ByteArrayInputStream(new byte[0]), 0);
         }
         var compressionAlgorithm = files.get(0).getCompressionAlgorithm();
         if (files.get(0).getBlobIdsCount() == 0) {
             files.removeFirst();
         }
-        return new SequenceInputStream(new Enumeration<>() {
+        var in = new SequenceInputStream(new Enumeration<>() {
             int index = 0;
 
             @Override
@@ -656,6 +658,7 @@ public InputStream nextElement() {
                 }
             }
         });
+        return new LengthAbleInputStream(in, size);
     }
 
     private InputStream readFileData(