Replace the GPU Lock with a queue-based approach

build.gradle

@@ -4,7 +4,7 @@ plugins {
 }
 
 group 'moe.lemonsh'
-version '0.1'
+version '0.2'
 
 jar {
     manifest {

src/main/java/moe/lemonsh/waifu2xlemon/App.java

@@ -6,6 +6,7 @@ import java.io.*;
 import java.net.URLConnection;
 import java.nio.charset.StandardCharsets;
 import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
 
 import static spark.Spark.*;
 
@@ -33,7 +34,7 @@ public class App {
             }
         }
         Waifu2x waifu2x = new Waifu2x(envexec);
-        threadPool(4, 1, -1);
+        new Thread(waifu2x::runQueue).start();
         try {
             port(Integer.parseInt(envport));
         } catch (NumberFormatException e) {
@@ -106,18 +107,29 @@ public class App {
                 case 2 -> outputExtension = ".webp";
                 default -> halt(400);
             }
-            Waifu2x.WaifuResult waifuResult;
+            final var latch = new CountDownLatch(1);
-            waifuResult = waifu2x.run(filePart.getInputStream(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'));
+            var resultRef = new Object() {
-            if (waifuResult.success) {
+                Waifu2x.WaifuResult waifuResult;
-                resp.raw().setContentType(URLConnection.guessContentTypeFromName(waifuResult.filename));
+            };
-                resp.raw().setHeader("Content-Disposition","attachment; filename="+waifuResult.filename);
+            var task = new Waifu2x.WaifuTask(filePart.getInputStream().readAllBytes(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'), result -> {
-                resp.raw().setContentLengthLong(waifuResult.size);
+                resultRef.waifuResult = result;
-                waifuResult.image.transferTo(resp.raw().getOutputStream());
+                latch.countDown();
+            });
+            switch (waifu2x.addTask(task)) {
+                case 1 -> { return "Invalid image (too big or in non-supported format)"; }
+                case 2 -> { return "The processing queue is full. Please try again later."; }
+            }
+            latch.await();
+            if (resultRef.waifuResult.success) {
+                resp.raw().setContentType(URLConnection.guessContentTypeFromName(resultRef.waifuResult.filename));
+                resp.raw().setHeader("Content-Disposition","attachment; filename="+resultRef.waifuResult.filename);
+                resp.raw().setContentLengthLong(resultRef.waifuResult.size);
+                resultRef.waifuResult.image.transferTo(resp.raw().getOutputStream());
                 resp.raw().getOutputStream().flush();
                 resp.raw().getOutputStream().close();
                 return null;
             }
-            return "Error.\n"+waifuResult.stdout;
+            return "Error.\n"+resultRef.waifuResult.stdout;
         });
     }
 }

src/main/java/moe/lemonsh/waifu2xlemon/Waifu2x.java

@@ -3,21 +3,24 @@ package moe.lemonsh.waifu2xlemon;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.imageio.ImageIO;
 import java.io.*;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.concurrent.BlockingDeque;
+import java.util.concurrent.LinkedBlockingDeque;
 
 public class Waifu2x {
-    public static class WaifuResult {
+    @FunctionalInterface
-        public boolean success;
+    public interface WaifuCallback {
-        public String stdout;
+        void exec(WaifuResult result);
-        public InputStream image;
-        public String filename;
-        public long size;
     }
 
-    private final Object GPULock;
+    public record WaifuTask(byte[] inputImage, String inputExtension, String outputExtension,
+                            char noise, char scale, WaifuCallback callback) {}
+
+    private final BlockingDeque<WaifuTask> taskQueue;
     private final Path ExecutablePath;
     private final Logger log = LoggerFactory.getLogger(Waifu2x.class);
 
@@ -27,11 +30,43 @@ public class Waifu2x {
             throw new IllegalArgumentException("'%s' is not a valid path.".formatted(wf2xExecutablePath));
         } else {
             ExecutablePath = executable.toPath().toAbsolutePath();
-            GPULock = new Object();
+            taskQueue = new LinkedBlockingDeque<>(2);
         }
     }
 
-    public WaifuResult run(InputStream inputImage, String inputExtension, String outputExtension, char noise, char scale) throws IOException, InterruptedException {
+    public synchronized void runQueue() {
+        while (true) {
+            try {
+                WaifuTask task = taskQueue.take();
+                log.info("New job! %s image with options noise=%c scale=%c output=%s".formatted(task.inputExtension, task.noise, task.scale, task.outputExtension));
+                var result = run(task.inputImage, task.inputExtension, task.outputExtension, task.noise, task.scale);
+                log.info("Job completed! Result:\n"+result.stdout);
+                task.callback.exec(result);
+            } catch (IOException | InterruptedException e) {
+                System.exit(10);
+            }
+        }
+    }
+
+    private boolean imageIllegal(byte[] fullImage) throws IOException {
+        // this can most likely be optimized
+        var image = ImageIO.read(new ByteArrayInputStream(fullImage));
+        return image == null || image.getWidth() > 2000 || image.getHeight() > 2000;
+    }
+
+    public int addTask(WaifuTask task) throws IOException {
+        return imageIllegal(task.inputImage) ? 1 : taskQueue.offer(task) ? 0 : 2;
+    }
+
+    static class WaifuResult {
+        public boolean success;
+        public String stdout;
+        public InputStream image;
+        public String filename;
+        public long size;
+    }
+
+    private WaifuResult run(byte[] inputImage, String inputExtension, String outputExtension, char noise, char scale) throws IOException, InterruptedException {
         var inputFile = new File("wfx_input" + inputExtension).toPath().toAbsolutePath();
         var outputFile = new File("wfx_output" + outputExtension);
         var outputFilePath = outputFile.toPath().toAbsolutePath();
@@ -43,19 +78,16 @@ public class Waifu2x {
                 return buf;
             }
         };
-        synchronized (GPULock) {
+        Files.deleteIfExists(inputFile);
-            log.info("Starting processing %s image with options noise=%c scale=%c output=%s".formatted(inputExtension, noise, scale, outputExtension));
+        Files.write(inputFile, inputImage);
-            Files.deleteIfExists(inputFile);
+        var waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n",
-            Files.copy(inputImage, inputFile);
+                Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start();
-            var waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n",
+        waifuProcess.waitFor();
-                    Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start();
+        Files.delete(inputFile);
-            waifuProcess.waitFor();
+        result.stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8);
-            Files.delete(inputFile);
+        if (!(result.success = outputFile.isFile())) return result;
-            log.info(result.stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8));
+        Files.copy(outputFilePath, outputBytes);
-            if (!(result.success = outputFile.isFile())) return result;
+        Files.delete(outputFilePath);
-            Files.copy(outputFilePath, outputBytes);
-            Files.delete(outputFilePath);
-        }
         result.image = new ByteArrayInputStream(outputBytes.toByteArray(), 0, outputBytes.size());
         result.size = outputBytes.size();
         return result;