diff --git a/build.gradle.kts b/build.gradle.kts index 985b17a..92d77af 100644 --- 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 index 9480ab7..4798357 100644 --- 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 index 809779f..ed0d8d0 100644 --- 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 index 329af7d..e67df03 100644 --- 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 new file mode 100644 index 0000000..53e3c7b --- /dev/null +++ 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 index 270621a..8a6b2c5 100644 --- 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 getOrGet(Supplier valueSupplier, Supplier 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; + } }