Attempt to fix memory leak

AzisabaNetwork · May 25, 2022 · 0815b42 · 0815b42
1 parent b892b59
commit 0815b42
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 2 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -6,7 +6,7 @@ plugins {
 }
 
 group = "net.azisaba.simpleProxy"
-version = "1.1.4"
+version = "1.1.5"
 
 extra.set("log4jVersion", "2.17.2")
 

diff --git a/proxy/src/main/java/net/azisaba/simpleProxy/proxy/Main.java b/proxy/src/main/java/net/azisaba/simpleProxy/proxy/Main.java
@@ -1,5 +1,6 @@
 package net.azisaba.simpleProxy.proxy;
 
+import net.azisaba.simpleProxy.proxy.util.MemoryReserve;
 import net.azisaba.simpleProxy.proxy.util.SignalUtil;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -11,6 +12,7 @@ public static void main(String[] args) {
         try {
             preload();
             SignalUtil.registerAll();
+            MemoryReserve.reserve();
             new ProxyInstance().start();
         } catch (Throwable throwable) {
             LOGGER.fatal("Failed to start proxy server", throwable);

diff --git a/proxy/src/main/java/net/azisaba/simpleProxy/proxy/connection/MessageForwarder.java b/proxy/src/main/java/net/azisaba/simpleProxy/proxy/connection/MessageForwarder.java
@@ -11,6 +11,8 @@
 import net.azisaba.simpleProxy.api.event.connection.RemoteConnectionActiveEvent;
 import net.azisaba.simpleProxy.proxy.ProxyInstance;
 import net.azisaba.simpleProxy.proxy.config.ProxyConfigInstance;
+import net.azisaba.simpleProxy.proxy.util.MemoryReserve;
+import net.azisaba.simpleProxy.proxy.util.Util;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.NotNull;
@@ -77,7 +79,8 @@ public void channelInactive(@NotNull ChannelHandlerContext ctx) throws Exception
         ctx.channel().close();
         if (remote != null) remote.close();
         if (ProxyConfigInstance.debug) {
-            LOGGER.info("Forwarder: Closed connection: " + ctx.channel());
+            int freed = Util.release(queue);
+            LOGGER.info("Forwarder: Closed connection: {} (freed {} objects)", ctx.channel(), freed);
         }
     }
 
@@ -125,6 +128,9 @@ public void channelRead(@NotNull ChannelHandlerContext ctx, @NotNull Object msg)
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
         LOGGER.warn("Caught exception in {}!", ctx.channel(), cause);
         ctx.channel().close();
+        if (cause instanceof OutOfMemoryError) {
+            MemoryReserve.tryShutdownGracefully();
+        }
     }
 
     void remoteActive() {

diff --git a/proxy/src/main/java/net/azisaba/simpleProxy/proxy/connection/MessageForwarderForwarder.java b/proxy/src/main/java/net/azisaba/simpleProxy/proxy/connection/MessageForwarderForwarder.java
@@ -12,6 +12,7 @@
 import net.azisaba.simpleProxy.api.config.Protocol;
 import net.azisaba.simpleProxy.api.config.ServerInfo;
 import net.azisaba.simpleProxy.proxy.config.ProxyConfigInstance;
+import net.azisaba.simpleProxy.proxy.util.MemoryReserve;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.NotNull;
@@ -121,5 +122,8 @@ public void channelInactive(@NotNull ChannelHandlerContext ctx) throws Exception
     public void exceptionCaught(@NotNull ChannelHandlerContext ctx, Throwable cause) {
         LOGGER.warn("Caught exception!", cause);
         ctx.channel().close();
+        if (cause instanceof OutOfMemoryError) {
+            MemoryReserve.tryShutdownGracefully();
+        }
     }
 }
diff --git a/proxy/src/main/java/net/azisaba/simpleProxy/proxy/util/MemoryReserve.java b/proxy/src/main/java/net/azisaba/simpleProxy/proxy/util/MemoryReserve.java
@@ -0,0 +1,28 @@
+package net.azisaba.simpleProxy.proxy.util;
+
+import net.azisaba.simpleProxy.proxy.ProxyInstance;
+
+public class MemoryReserve {
+    private static byte[] reserve = null; // 5 MB
+
+    public static void reserve() {
+        if (reserve == null) {
+            reserve = new byte[1024 * 1024 * 5];
+        }
+    }
+
+    public static void release() {
+        reserve = null;
+    }
+
+    public static void tryShutdownGracefully() {
+        try {
+            release();
+            System.gc();
+            ProxyInstance.getInstance().stop();
+            System.exit(0);
+        } catch (Throwable t) {
+            System.exit(0xc0000001);
+        }
+    }
+}
diff --git a/proxy/src/main/java/net/azisaba/simpleProxy/proxy/util/Util.java b/proxy/src/main/java/net/azisaba/simpleProxy/proxy/util/Util.java
@@ -1,7 +1,10 @@
 package net.azisaba.simpleProxy.proxy.util;
 
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
 import java.util.function.Supplier;
@@ -42,4 +45,28 @@ public static <T> T getOrGet(Supplier<T> valueSupplier, Supplier<T> defSupplier)
         if (value == null) return defSupplier.get();
         return value;
     }
+
+    /**
+     * Try to release the object and returns the number of freed objects.
+     * @param o the object, this can be an instance of ReferenceCounted or Iterable to free all elements.
+     * @return the number of freed objects.
+     */
+    public static int release(@Nullable Object o) {
+        if (o == null) {
+            return 0;
+        }
+        if (o instanceof ReferenceCounted) {
+            if (((ReferenceCounted) o).release()) {
+                return 1; // return 1 only if the object has been deallocated
+            }
+        } else if (o instanceof Iterable<?>) {
+            Iterable<?> itr = (Iterable<?>) o;
+            int freed = 0;
+            for (Object o1 : itr) {
+                freed += release(o1);
+            }
+            return freed;
+        }
+        return 0;
+    }
 }