Implement a strategy to handle OOM in direct memory

docs/index.asciidoc

@@ -221,6 +221,7 @@ This plugin supports the following configuration options plus the <<plugins-{typ
 | <<plugins-{type}s-{plugin}-host>> |<<string,string>>|No
 | <<plugins-{type}s-{plugin}-include_codec_tag>> |<<boolean,boolean>>|__Deprecated__
 | <<plugins-{type}s-{plugin}-port>> |<<number,number>>|Yes
+| <<plugins-{type}s-{plugin}-protect_direct_memory>> |<<boolean,boolean>>|No
 | <<plugins-{type}s-{plugin}-ssl>> |<<boolean,boolean>>|__Deprecated__
 | <<plugins-{type}s-{plugin}-ssl_certificate>> |a valid filesystem path|No
 | <<plugins-{type}s-{plugin}-ssl_certificate_authorities>> |<<array,array>>|No
@@ -384,6 +385,17 @@ deprecated[6.5.0, Replaced by <<plugins-{type}s-{plugin}-enrich>>]
 
 The port to listen on.
 
+[id="plugins-{type}s-{plugin}-protect_direct_memory"]
+===== `protect_direct_memory`
+
+  * Value type is <<boolean,boolean>>
+  * Default value is `true`
+
+If enabled, actively check native memory used by network part to do parsing and avoid
+out of memory conditions. When the consumption of native memory used is close to
+the maximum limit, connections are being closed in undetermined order until the safe
+memory condition is reestablished.
+
 [id="plugins-{type}s-{plugin}-ssl"]
 ===== `ssl`
 deprecated[6.6.0, Replaced by <<plugins-{type}s-{plugin}-ssl_enabled>>]

lib/logstash/inputs/beats.rb

@@ -74,6 +74,10 @@ class LogStash::Inputs::Beats < LogStash::Inputs::Base
   # The port to listen on.
   config :port, :validate => :number, :required => true
 
+  # Proactive checks that keep the beats input active when the memory used by protocol parser and network
+  # related operations is going to terminate.
+  config :protect_direct_memory, :validate => :boolean, :default => true
+
   # Events are by default sent in plain text. You can
   # enable encryption by setting `ssl` to true and configuring
   # the `ssl_certificate` and `ssl_key` options.
@@ -243,9 +247,11 @@ def register
   end # def register
 
   def create_server
-    server = org.logstash.beats.Server.new(@host, @port, @client_inactivity_timeout, @executor_threads)
+    server = org.logstash.beats.Server.new(@host, @port, @client_inactivity_timeout, @executor_threads, @protect_direct_memory)
     server.setSslHandlerProvider(new_ssl_handshake_provider(new_ssl_context_builder)) if @ssl_enabled
     server
+  rescue java.lang.IllegalArgumentException => e
+    configuration_error e.message
   end
 
   def run(output_queue)

spec/inputs/beats_spec.rb

@@ -14,6 +14,7 @@
   let(:port) { BeatsInputTest.random_port }
   let(:client_inactivity_timeout) { 400 }
   let(:threads) { 1 + rand(9) }
+  let(:protect_direct_memory) { true }
   let(:queue)  { Queue.new }
   let(:config)   do
     {
@@ -36,7 +37,7 @@
       let(:port) { 9001 }
 
       it "sends the required options to the server" do
-        expect(org.logstash.beats.Server).to receive(:new).with(host, port, client_inactivity_timeout, threads)
+        expect(org.logstash.beats.Server).to receive(:new).with(host, port, client_inactivity_timeout, threads, protect_direct_memory)
         subject.register
       end
     end
@@ -529,8 +530,8 @@
     subject(:plugin) { LogStash::Inputs::Beats.new(config) }
 
     before do
-      @server = org.logstash.beats.Server.new(host, port, client_inactivity_timeout, threads)
-      expect( org.logstash.beats.Server ).to receive(:new).with(host, port, client_inactivity_timeout, threads).and_return @server
+      @server = org.logstash.beats.Server.new(host, port, client_inactivity_timeout, threads, protect_direct_memory)
+      expect( org.logstash.beats.Server ).to receive(:new).with(host, port, client_inactivity_timeout, threads, protect_direct_memory).and_return @server
       expect( @server ).to receive(:listen)
 
       subject.register

src/main/java/org/logstash/beats/BeatsHandler.java

@@ -1,5 +1,7 @@
 package org.logstash.beats;
 
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.PooledByteBufAllocator;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.SimpleChannelInboundHandler;
 import org.apache.logging.log4j.LogManager;
@@ -92,6 +94,9 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E
                     logger.info(format("closing (" + cause.getMessage() + ")"));
                 }
             } else {
+                PooledByteBufAllocator allocator = (PooledByteBufAllocator) ByteBufAllocator.DEFAULT;
+                OOMConnectionCloser.DirectMemoryUsage usageSnapshot = OOMConnectionCloser.DirectMemoryUsage.capture(allocator);
+                logger.info("Connection {}, memory status used: {}, pinned: {}, ratio {}", ctx.channel(), usageSnapshot.used, usageSnapshot.pinned, usageSnapshot.ratio);
                 final Throwable realCause = extractCause(cause, 0);
                 if (logger.isDebugEnabled()){
                     logger.info(format("Handling exception: " + cause + " (caused by: " + realCause + ")"), cause);

src/main/java/org/logstash/beats/BeatsParser.java

@@ -2,9 +2,12 @@
 
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.ByteBufOutputStream;
+import io.netty.buffer.PooledByteBufAllocator;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageDecoder;
+import io.netty.handler.codec.DecoderException;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -14,6 +17,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.zip.Inflater;
 import java.util.zip.InflaterOutputStream;
 
@@ -48,8 +52,8 @@ private enum States {
 
     @Override
     protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws InvalidFrameProtocolException, IOException {
-        if(!hasEnoughBytes(in)) {
-            if (decodingCompressedBuffer){
+        if (!hasEnoughBytes(in)) {
+            if (decodingCompressedBuffer) {
                 throw new InvalidFrameProtocolException("Insufficient bytes in compressed content to decode: " + currentState);
             }
             return;
@@ -182,6 +186,7 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
 
             case READ_COMPRESSED_FRAME: {
                 logger.trace("Running: READ_COMPRESSED_FRAME");
+
                 inflateCompressedFrame(ctx, in, (buffer) -> {
                     transition(States.READ_HEADER);
 
@@ -199,9 +204,18 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
             }
             case READ_JSON: {
                 logger.trace("Running: READ_JSON");
-                ((V2Batch)batch).addMessage(sequence, in, requiredBytes);
-                if(batch.isComplete()) {
-                    if(logger.isTraceEnabled()) {
+                try {
+                    ((V2Batch) batch).addMessage(sequence, in, requiredBytes);
+                } catch (Throwable th) {
+                    // batch has to release its internal buffer before bubbling up the exception
+                    batch.release();
+
+                    // re throw the same error after released the internal buffer
+                    throw th;
+                }
+
+                if (batch.isComplete()) {
+                    if (logger.isTraceEnabled()) {
                         logger.trace("Sending batch size: " + this.batch.size() + ", windowSize: " + batch.getBatchSize() + " , seq: " + sequence);
                     }
                     out.add(batch);

src/main/java/org/logstash/beats/FlowLimiterHandler.java

@@ -0,0 +1,56 @@
+package org.logstash.beats;
+
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelHandler.Sharable;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+/**
+ * Configure the channel where it's installed to operate the reads in pull mode,
+ * disabling the autoread and explicitly invoking the read operation.
+ * The flow control to keep the outgoing buffer under control is done
+ * avoiding to read in new bytes if the outgoing direction became not writable, this
+ * excert back pressure to the TCP layer and ultimately to the upstream system.
+ * */
+@Sharable
+public final class FlowLimiterHandler extends ChannelInboundHandlerAdapter {
+
+    private final static Logger logger = LogManager.getLogger(FlowLimiterHandler.class);
+
+    @Override
+    public void channelRegistered(final ChannelHandlerContext ctx) throws Exception {
+        ctx.channel().config().setAutoRead(false);
+        super.channelRegistered(ctx);
+    }
+
+    @Override
+    public void channelActive(final ChannelHandlerContext ctx) throws Exception {
+        super.channelActive(ctx);
+        if (isAutoreadDisabled(ctx.channel()) && ctx.channel().isWritable()) {
+            ctx.channel().read();
+        }
+    }
+
+    @Override
+    public void channelReadComplete(final ChannelHandlerContext ctx) throws Exception {
+        super.channelReadComplete(ctx);
+        if (isAutoreadDisabled(ctx.channel()) && ctx.channel().isWritable()) {
+            ctx.channel().read();
+        }
+    }
+
+    private boolean isAutoreadDisabled(Channel channel) {
+        return !channel.config().isAutoRead();
+    }
+
+    @Override
+    public void channelWritabilityChanged(ChannelHandlerContext ctx) throws Exception {
+        ctx.channel().read();
+        super.channelWritabilityChanged(ctx);
+
+        logger.debug("Writability on channel {} changed to {}", ctx.channel(), ctx.channel().isWritable());
+    }
+
+}

src/main/java/org/logstash/beats/OOMConnectionCloser.java

@@ -0,0 +1,88 @@
+package org.logstash.beats;
+
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.PooledByteBufAllocator;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.internal.PlatformDependent;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+public class OOMConnectionCloser extends ChannelInboundHandlerAdapter {
+
+    private final PooledByteBufAllocator allocator;
+
+    static class DirectMemoryUsage {
+        final long used;
+        final long pinned;
+        private final PooledByteBufAllocator allocator;
+        final short ratio;
+
+        private DirectMemoryUsage(long used, long pinned, PooledByteBufAllocator allocator) {
+            this.used = used;
+            this.pinned = pinned;
+            this.allocator = allocator;
+            this.ratio = (short) Math.round(((double) pinned / used) * 100);
+        }
+
+        static DirectMemoryUsage capture(PooledByteBufAllocator allocator) {
+            long usedDirectMemory = allocator.metric().usedDirectMemory();
+            long pinnedDirectMemory = allocator.pinnedDirectMemory();
+            return new DirectMemoryUsage(usedDirectMemory, pinnedDirectMemory, allocator);
+        }
+
+        boolean isCloseToOOM() {
+            long maxDirectMemory = PlatformDependent.maxDirectMemory();
+            int chunkSize = allocator.metric().chunkSize();
+            return ((maxDirectMemory - used) <= chunkSize) && ratio > 75;
+        }
+    }
+
+    private final static Logger logger = LogManager.getLogger(OOMConnectionCloser.class);
+
+    public static final Pattern DIRECT_MEMORY_ERROR = Pattern.compile("^Cannot reserve \\d* bytes of direct buffer memory.*$");
+
+    OOMConnectionCloser() {
+        allocator = (PooledByteBufAllocator) ByteBufAllocator.DEFAULT;
+    }
+
+    @Override
+    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+        DirectMemoryUsage direct = DirectMemoryUsage.capture(allocator);
+        logger.info("Direct memory status, used: {}, pinned: {}, ratio: {}", direct.used, direct.pinned, direct.ratio);
+        if (direct.isCloseToOOM()) {
+            logger.warn("Closing connection {} because running out of memory, used: {}, pinned: {}, ratio {}", ctx.channel(), direct.used, direct.pinned, direct.ratio);
+            ReferenceCountUtil.release(msg); // to free the memory used by the buffer
+            ctx.flush();
+            ctx.close();
+        } else {
+            super.channelRead(ctx, msg);
+        }
+    }
+
+    @Override
+    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
+        if (isDirectMemoryOOM(cause)) {
+            DirectMemoryUsage direct = DirectMemoryUsage.capture(allocator);
+            logger.info("Direct memory status, used: {}, pinned: {}, ratio: {}", direct.used, direct.pinned, direct.ratio);
+            logger.warn("Dropping connection {} due to lack of available Direct Memory. Please lower the number of concurrent connections or reduce the batch size. " +
+                    "Alternatively, raise -XX:MaxDirectMemorySize option in the JVM running Logstash", ctx.channel());
+            ctx.flush();
+            ctx.close();
+        } else {
+            super.exceptionCaught(ctx, cause);
+        }
+    }
+
+    private boolean isDirectMemoryOOM(Throwable th) {
+        if (!(th instanceof OutOfMemoryError)) {
+            return false;
+        }
+        Matcher m = DIRECT_MEMORY_ERROR.matcher(th.getMessage());
+        return m.matches();
+    }
+}

src/main/java/org/logstash/beats/Runner.java

@@ -17,7 +17,7 @@ static public void main(String[] args) throws Exception {
         // Check for leaks.
         // ResourceLeakDetector.setLevel(ResourceLeakDetector.Level.PARANOID);
 
-        Server server = new Server("0.0.0.0", DEFAULT_PORT, 15, Runtime.getRuntime().availableProcessors());
+        Server server = new Server("0.0.0.0", DEFAULT_PORT, 15, Runtime.getRuntime().availableProcessors(), true);
 
             if(args.length > 0 && args[0].equals("ssl")) {
             logger.debug("Using SSL");

src/main/java/org/logstash/beats/Server.java

@@ -1,13 +1,18 @@
 package org.logstash.beats;
 
 import io.netty.bootstrap.ServerBootstrap;
-import io.netty.channel.*;
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInitializer;
+import io.netty.channel.ChannelOption;
+import io.netty.channel.ChannelPipeline;
 import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.channel.socket.SocketChannel;
 import io.netty.channel.socket.nio.NioServerSocketChannel;
 import io.netty.handler.timeout.IdleStateHandler;
 import io.netty.util.concurrent.DefaultEventExecutorGroup;
 import io.netty.util.concurrent.EventExecutorGroup;
+import io.netty.util.internal.PlatformDependent;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.logstash.netty.SslHandlerProvider;
@@ -18,18 +23,34 @@ public class Server {
     private final int port;
     private final String host;
     private final int beatsHeandlerThreadCount;
+    private final boolean protectDirectMemory;
     private NioEventLoopGroup workGroup;
     private IMessageListener messageListener = new MessageListener();
     private SslHandlerProvider sslHandlerProvider;
     private BeatsInitializer beatsInitializer;
 
     private final int clientInactivityTimeoutSeconds;
 
-    public Server(String host, int p, int clientInactivityTimeoutSeconds, int threadCount) {
+    public Server(String host, int p, int clientInactivityTimeoutSeconds, int threadCount, boolean protectDirectMemory) {
         this.host = host;
         port = p;
         this.clientInactivityTimeoutSeconds = clientInactivityTimeoutSeconds;
         beatsHeandlerThreadCount = threadCount;
+        this.protectDirectMemory = protectDirectMemory;
+
+        validateMinimumDirectMemory();
+    }
+
+    /**
+     * Validate if the configured available direct memory is enough for safe processing, else throws a ConfigurationException
+     * */
+    private void validateMinimumDirectMemory() {
+        long maxDirectMemoryAllocatable = PlatformDependent.maxDirectMemory();
+        if (maxDirectMemoryAllocatable < 256 * 1024 * 1024) {
+            long roundedMegabytes = Math.round((double) maxDirectMemoryAllocatable / 1024 / 1024);
+            throw new IllegalArgumentException("Max direct memory should be at least 256MB but was " + roundedMegabytes + "MB, " +
+                    "please check your MaxDirectMemorySize and io.netty.maxDirectMemory settings");
+        }
     }
 
     public void setSslHandlerProvider(SslHandlerProvider sslHandlerProvider){
@@ -126,6 +147,9 @@ private class BeatsInitializer extends ChannelInitializer<SocketChannel> {
 
         public void initChannel(SocketChannel socket){
             ChannelPipeline pipeline = socket.pipeline();
+            if (protectDirectMemory) {
+                pipeline.addLast(new OOMConnectionCloser());
+            }
 
             if (isSslEnabled()) {
                 pipeline.addLast(SSL_HANDLER, sslHandlerProvider.sslHandlerForChannel(socket));
@@ -134,7 +158,12 @@ public void initChannel(SocketChannel socket){
                              new IdleStateHandler(localClientInactivityTimeoutSeconds, IDLESTATE_WRITER_IDLE_TIME_SECONDS, localClientInactivityTimeoutSeconds));
             pipeline.addLast(BEATS_ACKER, new AckEncoder());
             pipeline.addLast(CONNECTION_HANDLER, new ConnectionHandler());
-            pipeline.addLast(beatsHandlerExecutorGroup, new BeatsParser(), new BeatsHandler(localMessageListener));
+            if (protectDirectMemory) {
+                pipeline.addLast(new FlowLimiterHandler());
+                pipeline.addLast(new ThunderingGuardHandler());
+            }
+            pipeline.addLast(beatsHandlerExecutorGroup, new BeatsParser());
+            pipeline.addLast(beatsHandlerExecutorGroup, new BeatsHandler(localMessageListener));
         }