refactor(性能测试): 性能测试下载zip使用流下载

framework/sdk-parent/sdk/src/main/java/io/metersphere/engine/EngineContext.java

@@ -1,5 +1,6 @@
 package io.metersphere.engine;
 
+import java.io.InputStream;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -14,7 +15,7 @@ public class EngineContext {
     private Integer resourceIndex;
     private double[] ratios;
     private Map<String, Object> properties = new HashMap<>();
-    private Map<String, byte[]> testResourceFiles = new HashMap<>();
+    private Map<String, InputStream> testResourceFiles = new HashMap<>();
     private Map<String, Boolean> splitFlag = new HashMap<>();
     private boolean checkBackendListener;
 
@@ -103,11 +104,11 @@ public class EngineContext {
         this.ratios = ratios;
     }
 
-    public Map<String, byte[]> getTestResourceFiles() {
+    public Map<String, InputStream> getTestResourceFiles() {
         return testResourceFiles;
     }
 
-    public void setTestResourceFiles(Map<String, byte[]> testResourceFiles) {
+    public void setTestResourceFiles(Map<String, InputStream> testResourceFiles) {
         this.testResourceFiles = testResourceFiles;
     }
 

performance-test/backend/src/main/java/io/metersphere/controller/JmeterFileController.java

@@ -3,17 +3,14 @@ package io.metersphere.controller;
 
 import io.metersphere.commons.utils.WeakConcurrentHashMap;
 import io.metersphere.controller.handler.annotation.NoResultHolder;
-import io.metersphere.dto.ZipDTO;
 import io.metersphere.service.JmeterFileService;
-import org.springframework.http.HttpHeaders;
+import jakarta.annotation.Resource;
-import org.springframework.http.MediaType;
+import jakarta.servlet.http.HttpServletResponse;
-import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.RestController;
 
-import jakarta.annotation.Resource;
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -45,15 +42,11 @@ public class JmeterFileController {
     }
 
     @GetMapping("download")
-    public ResponseEntity<byte[]> downloadJmeterFiles(@RequestParam("ratio") String ratio,
+    public void downloadJmeterFiles(@RequestParam("ratio") String ratio,
-                                                      @RequestParam("reportId") String reportId,
+                                    @RequestParam("reportId") String reportId,
-                                                      @RequestParam("resourceIndex") int resourceIndex) {
+                                    @RequestParam("resourceIndex") int resourceIndex, HttpServletResponse response) {
         double[] ratios = Arrays.stream(ratio.split(",")).mapToDouble(Double::parseDouble).toArray();
-        ZipDTO zipDTO = jmeterFileService.downloadZip(reportId, ratios, resourceIndex);
+        jmeterFileService.downloadZip(reportId, ratios, resourceIndex, response);
-        return ResponseEntity.ok()
-                .contentType(MediaType.parseMediaType("application/octet-stream"))
-                .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + zipDTO.getTestId() + ".zip\"")
-                .body(zipDTO.getContent());
     }
 
 }

performance-test/backend/src/main/java/io/metersphere/engine/EngineFactory.java

@@ -172,16 +172,16 @@ public class EngineFactory {
         /*
         {"timeout":10,"statusCode":["302","301"],"params":[{"name":"param1","enable":true,"value":"0","edit":false}],"domains":[{"domain":"baidu.com","enable":true,"ip":"127.0.0.1","edit":false}]}
          */
-        Map<String, byte[]> testResourceFiles = new HashMap<>();
+        Map<String, InputStream> testResourceFiles = new HashMap<>();
         byte[] props = getJMeterProperties(loadTestReport, engineContext);
         byte[] sysProps = getSystemProperties(loadTestReport, engineContext);
         byte[] hosts = getDNSConfig(loadTestReport, engineContext);
         // JMeter Properties
-        testResourceFiles.put("ms.properties", props);
+        testResourceFiles.put("ms.properties", new ByteArrayInputStream(props));
         // System Properties
-        testResourceFiles.put("sys.properties", sysProps);
+        testResourceFiles.put("sys.properties", new ByteArrayInputStream(sysProps));
         // DNS
-        testResourceFiles.put("hosts", hosts);
+        testResourceFiles.put("hosts", new ByteArrayInputStream(hosts));
 
         final EngineSourceParser engineSourceParser = EngineSourceParserFactory.createEngineSourceParser(engineContext.getFileType());
 
@@ -191,8 +191,8 @@ public class EngineFactory {
 
         if (CollectionUtils.isNotEmpty(resourceFiles)) {
             resourceFiles.forEach(cf -> {
-                byte[] csvContent = fileMetadataService.loadFileAsBytes(cf.getId());
+                InputStream in = fileMetadataService.getFileAsStream(cf.getId());
-                testResourceFiles.put(cf.getName(), csvContent);
+                testResourceFiles.put(cf.getName(), in);
             });
         }
         engineContext.setTestResourceFiles(testResourceFiles);

performance-test/backend/src/main/java/io/metersphere/parse/xml/reader/JmeterDocumentParser.java

@@ -3,6 +3,7 @@ package io.metersphere.parse.xml.reader;
 import io.metersphere.base.domain.TestResourcePool;
 import io.metersphere.commons.exception.MSException;
 import io.metersphere.commons.utils.CommonBeanFactory;
+import io.metersphere.commons.utils.IOUtils;
 import io.metersphere.config.KafkaProperties;
 import io.metersphere.engine.EngineContext;
 import io.metersphere.i18n.Translator;
@@ -18,6 +19,8 @@ import org.dom4j.Element;
 import org.dom4j.Node;
 import org.dom4j.tree.BaseElement;
 
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.*;
@@ -75,7 +78,7 @@ public class JmeterDocumentParser implements EngineSourceParser {
     }
 
 
-    private void parseHashTree(Element hashTree) {
+    private void parseHashTree(Element hashTree) throws IOException {
         if (invalid(hashTree)) {
             return;
         }
@@ -269,7 +272,7 @@ public class JmeterDocumentParser implements EngineSourceParser {
         item.setText(filename);
     }
 
-    private void processCsvDataSet(Element element) {
+    private void processCsvDataSet(Element element) throws IOException {
         List<Element> childNodes = element.elements();
         String fileName = null;
         for (Element item : childNodes) {
@@ -366,7 +369,7 @@ public class JmeterDocumentParser implements EngineSourceParser {
             item.setText(String.valueOf(BooleanUtils.toBoolean(recycle)));
         }
     }
-    private void splitCsvFile(Node item) {
+    private void splitCsvFile(Node item) throws IOException {
         String filename = item.getText();
         filename = StringUtils.removeStart(filename, "/test/");
         // 已经分割过的不再二次分割
@@ -379,11 +382,11 @@ public class JmeterDocumentParser implements EngineSourceParser {
         }
         double[] ratios = context.getRatios();
         int resourceIndex = context.getResourceIndex();
-        byte[] content = context.getTestResourceFiles().get(filename);
+        InputStream content = context.getTestResourceFiles().get(filename);
         if (content == null) {
             return;
         }
-        StringTokenizer tokenizer = new StringTokenizer(new String(content), StringUtils.LF);
+        StringTokenizer tokenizer = new StringTokenizer(IOUtils.toString(content, StandardCharsets.UTF_8), StringUtils.LF);
         if (!tokenizer.hasMoreTokens()) {
             return;
         }
@@ -436,7 +439,7 @@ public class JmeterDocumentParser implements EngineSourceParser {
             index++;
         }
         // 替换文件
-        context.getTestResourceFiles().put(filename, csv.toString().getBytes(StandardCharsets.UTF_8));
+        context.getTestResourceFiles().put(filename, new ByteArrayInputStream(csv.toString().getBytes(StandardCharsets.UTF_8)));
         context.getSplitFlag().put(filename, true);
     }
 

performance-test/backend/src/main/java/io/metersphere/service/JmeterFileService.java

@@ -5,15 +5,16 @@ import io.metersphere.base.mapper.LoadTestReportMapper;
 import io.metersphere.base.mapper.ext.ExtLoadTestReportMapper;
 import io.metersphere.commons.exception.MSException;
 import io.metersphere.commons.utils.LogUtil;
-import io.metersphere.dto.ZipDTO;
 import io.metersphere.engine.EngineContext;
 import io.metersphere.engine.EngineFactory;
+import jakarta.annotation.Resource;
+import jakarta.servlet.http.HttpServletResponse;
 import org.apache.commons.collections4.MapUtils;
 import org.springframework.stereotype.Service;
 
-import jakarta.annotation.Resource;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -29,7 +30,7 @@ public class JmeterFileService {
     @Resource
     private LoadTestReportMapper loadTestReportMapper;
 
-    public ZipDTO downloadZip(String reportId, double[] ratios, int resourceIndex) {
+    public void downloadZip(String reportId, double[] ratios, int resourceIndex, HttpServletResponse response) {
         try {
             LoadTestReportWithBLOBs loadTestReport = loadTestReportMapper.selectByPrimaryKey(reportId);
             int wait = 0;
@@ -45,52 +46,47 @@ public class JmeterFileService {
                 MSException.throwException("测试报告不存在或还没产生");
             }
             EngineContext context = EngineFactory.createContext(loadTestReport, ratios, resourceIndex);
-            byte[] bytes = zipFilesToByteArray(context);
+            zipFilesToByteArray(context, response);
-            return new ZipDTO(loadTestReport.getTestId(), bytes);
         } catch (Exception e) {
             LogUtil.error(e.getMessage(), e);
             MSException.throwException(e);
         }
-        return null;
     }
 
-    private byte[] zipFilesToByteArray(EngineContext context) throws IOException {
+    private void zipFilesToByteArray(EngineContext context, HttpServletResponse response) throws IOException {
         String testId = context.getTestId();
         String fileName = testId + ".jmx";
 
-        Map<String, byte[]> files = new HashMap<>();
 
         //  每个测试生成一个文件夹
-        files.put(fileName, context.getContent());
         // 保存jmx
         LoadTestReportWithBLOBs record = new LoadTestReportWithBLOBs();
         record.setId(context.getReportId());
         record.setJmxContent(new String(context.getContent()));
         extLoadTestReportMapper.updateJmxContentIfAbsent(record);
-        // 保存 byte[]
+        // 关联文件
-        Map<String, byte[]> jarFiles = context.getTestResourceFiles();
+        Map<String, InputStream> testResourceFiles = context.getTestResourceFiles();
-        if (!MapUtils.isEmpty(jarFiles)) {
+        try (ZipOutputStream zos = new ZipOutputStream(response.getOutputStream())) {
-            for (String k : jarFiles.keySet()) {
+            // jmx 本身
-                byte[] v = jarFiles.get(k);
+            zos.putNextEntry(new ZipEntry(fileName));
-                files.put(k, v);
+            zos.write(context.getContent());
+            zos.closeEntry();
+            // 关联文件
+            if (!MapUtils.isEmpty(testResourceFiles)) {
+                for (String name : testResourceFiles.keySet()) {
+                    InputStream in = testResourceFiles.get(name);
+                    zos.putNextEntry(new ZipEntry(name));
+                    byte[] bytes = new byte[4096];
+                    int len;
+                    while ((len = in.read(bytes)) != -1) {
+                        zos.write(bytes, 0, len);
+                    }
+                    in.close();
+                    zos.closeEntry();
+                }
             }
+            response.setHeader("Content-Disposition", "attachment;filename=" + testId + ".zip");
+            response.flushBuffer();
         }
-
-        return listBytesToZip(files);
-    }
-
-
-    private byte[] listBytesToZip(Map<String, byte[]> mapReport) throws IOException {
-        ByteArrayOutputStream baos = new ByteArrayOutputStream();
-        ZipOutputStream zos = new ZipOutputStream(baos);
-        for (Map.Entry<String, byte[]> report : mapReport.entrySet()) {
-            ZipEntry entry = new ZipEntry(report.getKey());
-            entry.setSize(report.getValue().length);
-            zos.putNextEntry(entry);
-            zos.write(report.getValue());
-        }
-        zos.closeEntry();
-        zos.close();
-        return baos.toByteArray();
     }
 }