Optimize the task to create the zip while processing, not only when i…

…t will be sent.
sepinf-inc · Feb 19, 2025 · d53b2a9 · d53b2a9
1 parent d795d91
commit d53b2a9
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 48 deletions.
diff --git a/iped-app/resources/config/conf/RemoteImageClassifier.txt b/iped-app/resources/config/conf/RemoteImageClassifier.txt
@@ -3,7 +3,7 @@
 #################################
 
 # url running the classifier
-host = http://localhost:8000
+url = http://192.168.160.208:8000/zip
 
 # max size number of thumbs to be included in the zip file sent to the server
 batch_size = 50

diff --git a/iped-engine/src/main/java/iped/engine/task/RemoteImageClassifier.java b/iped-engine/src/main/java/iped/engine/task/RemoteImageClassifier.java
@@ -30,11 +30,52 @@
 import iped.engine.config.RemoteImageClassifierConfig;
 import iped.engine.task.index.IndexItem;
 
+class zipfile {
+    TemporaryResources tmp;
+    File zipFile;
+    FileOutputStream fos;
+    ZipOutputStream zos;
+    int size;
+
+    public zipfile() throws IOException {
+        tmp = new TemporaryResources();
+        zipFile = tmp.createTemporaryFile();
+        fos = new FileOutputStream(zipFile);
+        zos = new ZipOutputStream(fos);
+        size = 0;
+    }
+
+    public int size() {
+        return size;
+    }
+
+    public void addFileToZip(String filename, byte[] dados) throws IOException {
+
+        ZipEntry zipEntry = new ZipEntry(filename);
+        zos.putNextEntry(zipEntry);
+        zos.write(dados, 0, dados.length);
+        zos.closeEntry();
+        size++;
+
+    }
+
+    public File closeAndGetZip() throws IOException {
+        zos.close();
+        fos.close();
+        return zipFile;
+    }
+
+    public void clean() throws IOException {
+        tmp.close();
+    }
+}
+
 public class RemoteImageClassifier extends AbstractTask {
 
     private String url = "http://localhost:8000/zip";
     private int batch_size = 50;
     private Map<String, IItem> queue = new TreeMap<>();
+    private zipfile zip = null;
     private RemoteImageClassifierConfig config;
 
     @Override
@@ -51,6 +92,10 @@ public void init(ConfigurationManager configurationManager) throws Exception {
         config = configurationManager.findObject(RemoteImageClassifierConfig.class);
         url = config.getUrl();
         batch_size = config.getBatchSize();
+        if (zip == null) {
+            zip = new zipfile();
+        }
+
     }
 
     @Override
@@ -59,41 +104,37 @@ public void finish() throws Exception {
 
     }
 
-    private void addFileToZip(String filename, byte[] dados, ZipOutputStream zos) throws IOException {
+    private void sendItemsToNextTask() throws Exception {
+        for (IItem item : queue.values()) {
+            if (item != null) {
+                super.sendToNextTask(item);
+            }
+        }
+    }
 
-        ZipEntry zipEntry = new ZipEntry(filename);
-        zos.putNextEntry(zipEntry);
-        zos.write(dados, 0, dados.length);
-        zos.closeEntry();
 
-    }
+    protected void sendToNextTask(IItem item) throws Exception {
+        if (!isEnabled()) {
+            super.sendToNextTask(item);
+            return;
+        }
 
-    private void sendItemsToNextTask() throws Exception {
-        // criar zip um com i.getThumb() de todos os itens
-        System.out.println("Envia fila de tamanho" + queue.size());
-        try (TemporaryResources tmp = new TemporaryResources()) {
-            if (queue.size() > 0) {
-                File zipFile = tmp.createTemporaryFile();
-                try (FileOutputStream fos = new FileOutputStream(zipFile);
-                        ZipOutputStream zos = new ZipOutputStream(fos)) {
-                    for (String name : queue.keySet()) {
-                        addFileToZip(name, queue.get(name).getThumb(), zos);
-                    }
+        if (zip.size() > 0 && (zip.size() >= batch_size || item.isQueueEnd())) {
+            sendZipFile(zip.closeAndGetZip());
+            zip = new zipfile();
+            sendItemsToNextTask();
+        }
+
 
-                }
-                sendZipFile(zipFile);
-            }
-        } finally {
-            for (IItem item : queue.values()) {
-                if (item != null) {
-                    super.sendToNextTask(item);
-                }
-            }
-            queue.clear();
+        if (!queue.containsValue(item) || item.isQueueEnd()) {
+            super.sendToNextTask(item);
         }
+
     }
 
+
     private void sendZipFile(File zipFile) throws IOException {
+        System.out.println("Envia zip de tamanho" + zip.size());
         try (CloseableHttpClient client = HttpClients.createDefault()) {
             HttpPost post = new HttpPost(url);
 
@@ -143,30 +184,17 @@ private void sendZipFile(File zipFile) throws IOException {
         }
     }
 
-    protected void sendToNextTask(IItem item) throws Exception {
-        if (!isEnabled()) {
-            super.sendToNextTask(item);
-            return;
-        }
-        if (queue.size() > 0 && (queue.size() >= batch_size || item.isQueueEnd())) {
-            sendItemsToNextTask();
-        }
-        if (!queue.containsValue(item) || item.isQueueEnd()) {
-            super.sendToNextTask(item);
-        }
-
-
-    }
-
-    @Override
+      @Override
     protected boolean processQueueEnd() {
         // TODO Auto-generated method stub
         return true;
     }
 
     @Override
     protected void process(IItem evidence) throws Exception {
-        if (evidence.isQueueEnd() && queue.size() > 0) {
+        if (evidence.isQueueEnd() && zip.size() > 0) {
+            sendZipFile(zip.closeAndGetZip());
+            zip = new zipfile(); 
             sendItemsToNextTask();
             return;
         }
@@ -175,8 +203,7 @@ protected void process(IItem evidence) throws Exception {
                 || evidence.getThumb() == null || evidence.isQueueEnd()) {
             return;
         }
-        queue.put(evidence.getExtraAttribute(IndexItem.TRACK_ID).toString() + ".jpg", evidence);
-
+        String name = evidence.getExtraAttribute(IndexItem.TRACK_ID).toString() + ".jpg";        zip.addFileToZip(name, evidence.getThumb());        queue.put(name, evidence);
     }
-
+        
 }