fix(net): add outbound buffer limits and write timeouts

p2p/src/main/java/org/tron/p2p/connection/Channel.java

@@ -10,6 +10,7 @@
 import io.netty.handler.codec.protobuf.ProtobufVarint32LengthFieldPrepender;
 import io.netty.handler.timeout.ReadTimeoutException;
 import io.netty.handler.timeout.ReadTimeoutHandler;
+import io.netty.handler.timeout.WriteTimeoutException;
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
@@ -79,6 +80,7 @@ public class Channel {
   @Getter
   private long avgLatency;
   private long count;
+  private final WriteBufferHandler writeBufferHandler = new WriteBufferHandler(this);
 
   public void init(ChannelPipeline pipeline, String nodeId, boolean discoveryMode) {
     this.discoveryMode = discoveryMode;
@@ -89,6 +91,7 @@ public void init(ChannelPipeline pipeline, String nodeId, boolean discoveryMode)
     pipeline.addLast(TrafficStats.tcp);
     pipeline.addLast("protoPrepend", new ProtobufVarint32LengthFieldPrepender());
     pipeline.addLast("protoDecode", new P2pProtobufVarint32FrameDecoder(this));
+    pipeline.addLast("writeTimeoutHandler", writeBufferHandler);
     pipeline.addLast("messageHandler", messageHandler);
   }
 
@@ -102,6 +105,7 @@ public void processException(Throwable throwable) {
     }
     SocketAddress address = ctx.channel().remoteAddress();
     if (throwable instanceof ReadTimeoutException
+        || throwable instanceof WriteTimeoutException
         || throwable instanceof IOException
         || throwable instanceof CorruptedFrameException) {
       logger.warn("Close peer {}, reason: {}", address, throwable.getMessage());
@@ -177,7 +181,7 @@ public void send(byte[] data) {
       }
 
       ByteBuf byteBuf = Unpooled.wrappedBuffer(data);
-      ctx.writeAndFlush(byteBuf).addListener((ChannelFutureListener) future -> {
+      writeBufferHandler.writeAndFlush(ctx, byteBuf).addListener((ChannelFutureListener) future -> {
         if (!future.isSuccess() && !isDisconnect) {
           logger.warn("Send to {} failed, message-type:{}, cause:{}",
               ctx.channel().remoteAddress(), ByteArray.byte2int(type),

p2p/src/main/java/org/tron/p2p/connection/WriteBufferHandler.java

@@ -0,0 +1,139 @@
+package org.tron.p2p.connection;
+
+import com.google.protobuf.CodedOutputStream;
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelPromise;
+import io.netty.handler.timeout.WriteTimeoutHandler;
+import java.nio.channels.ClosedChannelException;
+import java.util.LinkedHashSet;
+import java.util.Set;
+import lombok.extern.slf4j.Slf4j;
+
+/** Bounds pending TCP writes and closes connections whose writes stall. */
+@Slf4j(topic = "net")
+final class WriteBufferHandler extends WriteTimeoutHandler {
+
+  private static final Budget GLOBAL_BUDGET = new Budget(200L * 1024 * 1024, 500L * 1024 * 1024);
+
+  private final Channel channel;
+  private final Budget budget;
+  // Guarded by budget, including sends originating outside the channel's event loop.
+  private long pendingBytes;
+  private boolean closing;
+
+  WriteBufferHandler(Channel channel) {
+    this(channel, GLOBAL_BUDGET);
+  }
+
+  WriteBufferHandler(Channel channel, Budget budget) {
+    super(10);
+    this.channel = channel;
+    this.budget = budget;
+  }
+
+  ChannelFuture writeAndFlush(ChannelHandlerContext ctx, ByteBuf data) {
+    // Reserve before Netty can enqueue a write task on the event loop. Include the frame prefix
+    // that protoPrepend adds later, after the payload has already been compressed by Channel.send.
+    long size = (long) data.readableBytes()
+        + CodedOutputStream.computeUInt32SizeNoTag(data.readableBytes());
+    ChannelPromise promise = ctx.newPromise();
+    promise.setUncancellable();
+    if (!budget.reserve(this, size)) {
+      data.release();
+      return promise.setFailure(new ClosedChannelException());
+    }
+    promise.addListener(future -> budget.release(this, size));
+    if (channel.isDisconnect()) {
+      data.release();
+      return promise.setFailure(new ClosedChannelException());
+    }
+    return ctx.writeAndFlush(data, promise);
+  }
+
+  private void disconnect() {
+    logger.warn("Close peer {}, pending outbound data limit exceeded",
+        channel.getInetSocketAddress());
+    channel.close();
+  }
+
+  /** Shared by TCP connections; the global limit triggers eviction rather than blocking I/O. */
+  static final class Budget {
+
+    private final long peerLimit;
+    private final long globalLimit;
+    private final Set<WriteBufferHandler> peers = new LinkedHashSet<>();
+    private long pendingBytes;
+    private WriteBufferHandler evicting;
+
+    Budget(long peerLimit, long globalLimit) {
+      this.peerLimit = peerLimit;
+      this.globalLimit = globalLimit;
+    }
+
+    private boolean reserve(WriteBufferHandler handler, long size) {
+      WriteBufferHandler victim;
+      synchronized (this) {
+        if (handler.closing || handler.channel.isDisconnect()) {
+          return false;
+        }
+        handler.pendingBytes += size;
+        pendingBytes += size;
+        peers.add(handler);
+        if (handler.pendingBytes > peerLimit) {
+          handler.closing = true;
+          victim = handler;
+          if (evicting == null) {
+            evicting = handler;
+          }
+        } else {
+          victim = selectVictim();
+        }
+      }
+      // Closing may immediately complete writes and re-enter release(). Never close under the lock.
+      if (victim != null) {
+        victim.disconnect();
+      }
+      return true;
+    }
+
+    private void release(WriteBufferHandler handler, long size) {
+      WriteBufferHandler victim;
+      synchronized (this) {
+        handler.pendingBytes -= size;
+        pendingBytes -= size;
+        if (handler.pendingBytes == 0) {
+          peers.remove(handler);
+          if (evicting == handler) {
+            evicting = null;
+          }
+        }
+        victim = selectVictim();
+      }
+      if (victim != null) {
+        victim.disconnect();
+      }
+    }
+
+    // Called with the budget lock held. Wait for the current victim's writes to be released
+    // before selecting another, even if multiple producers observe the global limit concurrently.
+    private WriteBufferHandler selectVictim() {
+      if (evicting != null || pendingBytes < globalLimit) {
+        return null;
+      }
+      WriteBufferHandler largest = null;
+      for (WriteBufferHandler candidate : peers) {
+        if (!candidate.closing && !candidate.channel.isDisconnect()
+            && (largest == null || candidate.pendingBytes > largest.pendingBytes)) {
+          largest = candidate;
+        }
+      }
+      if (largest != null) {
+        largest.closing = true;
+        evicting = largest;
+      }
+      return largest;
+    }
+  }
+}

p2p/src/test/java/org/tron/p2p/connection/ChannelCoreTest.java

@@ -2,12 +2,14 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.channel.ChannelOutboundBuffer;
 import io.netty.channel.embedded.EmbeddedChannel;
 import io.netty.handler.codec.CorruptedFrameException;
 import io.netty.handler.timeout.ReadTimeoutException;
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
+import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -107,6 +109,36 @@ public void sendWritesTheFramedMessage() throws Exception {
     netty.finishAndReleaseAll();
   }
 
+  @Test
+  public void stalledWriteClosesConnectionAfterTenSeconds() throws Exception {
+    EmbeddedChannel netty = new EmbeddedChannel() {
+      @Override
+      protected void doWrite(ChannelOutboundBuffer buffer) {
+        // A socket whose remote peer does not read: leave writes pending in Netty.
+      }
+    };
+    netty.freezeTime();
+    Channel channel = new Channel();
+    channel.init(netty.pipeline(), "", false);
+    attach(channel, netty);
+    Field ctx = Channel.class.getDeclaredField("ctx");
+    ctx.setAccessible(true);
+    ctx.set(channel, netty.pipeline().lastContext());
+    try {
+      channel.send(new PingMessage());
+      netty.advanceTimeBy(9, TimeUnit.SECONDS);
+      netty.runScheduledPendingTasks();
+      Assert.assertTrue(netty.isOpen());
+
+      netty.advanceTimeBy(1, TimeUnit.SECONDS);
+      netty.runScheduledPendingTasks();
+      Assert.assertFalse("a pending write must time out even without a business queue",
+          netty.isOpen());
+    } finally {
+      netty.finishAndReleaseAll();
+    }
+  }
+
   @Test
   public void sendIsSuppressedOnceDisconnected() throws Exception {
     EmbeddedChannel netty = new EmbeddedChannel(new ChannelInboundHandlerAdapter());