diff --git a/build.gradle b/build.gradle index 82eb029..8a3adce 100644 --- a/build.gradle +++ b/build.gradle @@ -4,7 +4,7 @@ plugins { } group 'moe.lemonsh' -version '0.1' +version '0.2' jar { manifest { diff --git a/src/main/java/moe/lemonsh/waifu2xlemon/App.java b/src/main/java/moe/lemonsh/waifu2xlemon/App.java index 67f122c..079bb7b 100644 --- a/src/main/java/moe/lemonsh/waifu2xlemon/App.java +++ b/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; - waifuResult = waifu2x.run(filePart.getInputStream(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0')); - if (waifuResult.success) { - resp.raw().setContentType(URLConnection.guessContentTypeFromName(waifuResult.filename)); - resp.raw().setHeader("Content-Disposition","attachment; filename="+waifuResult.filename); - resp.raw().setContentLengthLong(waifuResult.size); - waifuResult.image.transferTo(resp.raw().getOutputStream()); + final var latch = new CountDownLatch(1); + var resultRef = new Object() { + Waifu2x.WaifuResult waifuResult; + }; + var task = new Waifu2x.WaifuTask(filePart.getInputStream().readAllBytes(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'), result -> { + resultRef.waifuResult = result; + 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; }); } } diff --git a/src/main/java/moe/lemonsh/waifu2xlemon/Waifu2x.java b/src/main/java/moe/lemonsh/waifu2xlemon/Waifu2x.java index bd7752b..d11fb52 100644 --- a/src/main/java/moe/lemonsh/waifu2xlemon/Waifu2x.java +++ b/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 { - public boolean success; - public String stdout; - public InputStream image; - public String filename; - public long size; + @FunctionalInterface + public interface WaifuCallback { + void exec(WaifuResult result); } - private final Object GPULock; + public record WaifuTask(byte[] inputImage, String inputExtension, String outputExtension, + char noise, char scale, WaifuCallback callback) {} + + private final BlockingDeque 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) { - log.info("Starting processing %s image with options noise=%c scale=%c output=%s".formatted(inputExtension, noise, scale, outputExtension)); - Files.deleteIfExists(inputFile); - Files.copy(inputImage, inputFile); - var waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n", - Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start(); - waifuProcess.waitFor(); - Files.delete(inputFile); - log.info(result.stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8)); - if (!(result.success = outputFile.isFile())) return result; - Files.copy(outputFilePath, outputBytes); - Files.delete(outputFilePath); - } + Files.deleteIfExists(inputFile); + Files.write(inputFile, inputImage); + var waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n", + Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start(); + waifuProcess.waitFor(); + Files.delete(inputFile); + result.stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8); + if (!(result.success = outputFile.isFile())) return result; + Files.copy(outputFilePath, outputBytes); + Files.delete(outputFilePath); result.image = new ByteArrayInputStream(outputBytes.toByteArray(), 0, outputBytes.size()); result.size = outputBytes.size(); return result;