tcp: suggested changes

Author: AgraVator

netty/src/main/java/io/grpc/netty/NettyClientHandler.java

@@ -486,6 +486,7 @@ private void onRstStreamRead(int streamId, long errorCode) {
 
   @Override
   public void close(ChannelHandlerContext ctx, ChannelPromise promise) throws Exception {
+    tcpMetrics.recordTcpInfo(ctx.channel());
     logger.fine("Network channel being closed by the application.");
     if (ctx.channel().isActive()) { // Ignore notification that the socket was closed
       lifecycleManager.notifyShutdown(

netty/src/main/java/io/grpc/netty/NettyServer.java

@@ -21,6 +21,7 @@
 import static io.netty.channel.ChannelOption.ALLOCATOR;
 import static io.netty.channel.ChannelOption.SO_KEEPALIVE;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -68,6 +69,7 @@
 import java.util.concurrent.Callable;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import javax.annotation.Nullable;
 
 /**
  * Netty-based server implementation.
@@ -138,7 +140,8 @@ class NettyServer implements InternalServer, InternalWithLogId {
       long maxConnectionAgeInNanos, long maxConnectionAgeGraceInNanos,
       boolean permitKeepAliveWithoutCalls, long permitKeepAliveTimeInNanos,
       int maxRstCount, long maxRstPeriodNanos,
-      Attributes eagAttributes, InternalChannelz channelz) {
+      Attributes eagAttributes, InternalChannelz channelz,
+      @Nullable MetricRecorder metricRecorder) {
     this.addresses = checkNotNull(addresses, "addresses");
     this.channelFactory = checkNotNull(channelFactory, "channelFactory");
     checkNotNull(channelOptions, "channelOptions");
@@ -174,6 +177,13 @@ class NettyServer implements InternalServer, InternalWithLogId {
     this.channelz = Preconditions.checkNotNull(channelz);
     this.logId = InternalLogId.allocate(getClass(), addresses.isEmpty() ? "No address" :
         String.valueOf(addresses));
+    this.metricRecorder = metricRecorder;
+  }
+
+  @VisibleForTesting
+  @Nullable
+  MetricRecorder getMetricRecorder() {
+    return metricRecorder;
   }
 
   @Override

netty/src/main/java/io/grpc/netty/NettyServerBuilder.java

@@ -32,6 +32,7 @@
 import io.grpc.ExperimentalApi;
 import io.grpc.ForwardingServerBuilder;
 import io.grpc.Internal;
+import io.grpc.MetricRecorder;
 import io.grpc.ServerBuilder;
 import io.grpc.ServerCredentials;
 import io.grpc.ServerStreamTracer;
@@ -60,6 +61,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
 import javax.net.ssl.SSLException;
 
 /**
@@ -116,6 +118,7 @@ public final class NettyServerBuilder extends ForwardingServerBuilder<NettyServe
   private int maxRstCount;
   private long maxRstPeriodNanos;
   private Attributes eagAttributes = Attributes.EMPTY;
+  @Nullable private MetricRecorder metricRecorder;
 
   /**
    * Creates a server builder that will bind to the given port.
@@ -703,6 +706,20 @@ void eagAttributes(Attributes eagAttributes) {
     this.eagAttributes = checkNotNull(eagAttributes, "eagAttributes");
   }
 
+  /**
+   * Sets the {@link MetricRecorder} to use for recording metrics.
+   *
+   * @param metricRecorder the metric recorder to use
+   * @return this
+   * @since 1.81.0
+   */
+  @CanIgnoreReturnValue
+  public NettyServerBuilder setMetricRecorder(
+      @Nullable MetricRecorder metricRecorder) {
+    this.metricRecorder = metricRecorder;
+    return this;
+  }
+
   NettyServer buildTransportServers(
       List<? extends ServerStreamTracer.Factory> streamTracerFactories) {
     assertEventLoopsAndChannelType();
@@ -737,7 +754,8 @@ NettyServer buildTransportServers(
         maxRstCount,
         maxRstPeriodNanos,
         eagAttributes,
-        this.serverImplBuilder.getChannelz());
+        this.serverImplBuilder.getChannelz(),
+        metricRecorder);
   }
 
   @VisibleForTesting

netty/src/main/java/io/grpc/netty/TcpMetrics.java

@@ -23,11 +23,18 @@
 import io.grpc.MetricInstrumentRegistry;
 import io.grpc.MetricRecorder;
 import io.netty.channel.Channel;
+import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
 
+/**
+ * Utility for collecting TCP metrics from Netty channels.
+ */
 final class TcpMetrics {
 
   private static final Metrics DEFAULT_METRICS;
@@ -51,9 +58,9 @@ static Metrics getDefaultMetrics() {
   static final class Metrics {
     final LongCounterMetricInstrument connectionsCreated;
     final LongUpDownCounterMetricInstrument connectionCount;
-    final LongCounterMetricInstrument packetsRetransmitted;
-    final LongCounterMetricInstrument recurringRetransmits;
-    final DoubleHistogramMetricInstrument minRtt;
+    @Nullable final LongCounterMetricInstrument packetsRetransmitted;
+    @Nullable final LongCounterMetricInstrument recurringRetransmits;
+    @Nullable final DoubleHistogramMetricInstrument minRtt;
 
     Metrics(MetricInstrumentRegistry registry, boolean epollAvailable) {
       List<String> requiredLabels = Collections.singletonList("grpc.target");
@@ -171,12 +178,32 @@ private static DoubleHistogramMetricInstrument safelyRegisterDoubleHistogram(
     }
   }
 
+  private static final class ChannelReflectionAccessor {
+    final Method tcpInfoMethod;
+    final Constructor<?> tcpInfoConstructor;
+    final Method totalRetransMethod;
+    final Method retransmitsMethod;
+    final Method rttMethod;
+
+    ChannelReflectionAccessor(
+        Method tcpInfoMethod, Constructor<?> tcpInfoConstructor, Method totalRetransMethod,
+        Method retransmitsMethod, Method rttMethod) {
+      this.tcpInfoMethod = tcpInfoMethod;
+      this.tcpInfoConstructor = tcpInfoConstructor;
+      this.totalRetransMethod = totalRetransMethod;
+      this.retransmitsMethod = retransmitsMethod;
+      this.rttMethod = rttMethod;
+    }
+  }
+
   static final class Tracker {
     private final MetricRecorder metricRecorder;
     private final String target;
     private final Metrics metrics;
     private final String epollSocketChannelClassName;
     private final String epollTcpInfoClassName;
+    private volatile ChannelReflectionAccessor channelReflectionAccessor;
+    private long lastTotalRetrans;
 
     Tracker(MetricRecorder metricRecorder, String target) {
       this(metricRecorder, target, DEFAULT_METRICS);
@@ -197,6 +224,9 @@ static final class Tracker {
       this.epollTcpInfoClassName = epollTcpInfoClassName;
     }
 
+    private static final Map<String, ChannelReflectionAccessor> accessorCache =
+        new HashMap<>();
+
     private static final long RECORD_INTERVAL_MILLIS;
 
     static {
@@ -261,43 +291,65 @@ void channelInactive(Channel channel) {
         metricRecorder.addLongUpDownCounter(metrics.connectionCount, -1,
             Collections.singletonList(target), labelValues);
         // Final collection on close
-        recordTcpInfo(channel);
+        recordTcpInfo(channel, true);
       }
     }
 
-    private void recordTcpInfo(Channel channel) {
+    void recordTcpInfo(Channel channel) {
+      recordTcpInfo(channel, false);
+    }
+
+    void recordTcpInfo(Channel channel, boolean isClosed) {
       if (metricRecorder == null || target == null) {
         return;
       }
       java.util.List<String> labelValues = getLabelValues(channel);
       try {
-        if (channel.getClass().getName().equals(epollSocketChannelClassName)) {
-          Class<?> tcpInfoClass = Class.forName(epollTcpInfoClassName);
-          Method tcpInfoMethod = channel.getClass().getMethod("tcpInfo", tcpInfoClass);
-          Object info = tcpInfoClass.getDeclaredConstructor().newInstance();
-          tcpInfoMethod.invoke(channel, info);
-
-          Method totalRetransMethod = tcpInfoClass.getMethod("totalRetrans");
-          Method retransmitsMethod = tcpInfoClass.getMethod("retransmits");
-          Method rttMethod = tcpInfoClass.getMethod("rtt");
-
-          long totalRetrans = (Long) totalRetransMethod.invoke(info);
-          int retransmits = (Integer) retransmitsMethod.invoke(info);
-          long rtt = (Long) rttMethod.invoke(info);
-
-          if (metrics.packetsRetransmitted != null) {
-            metricRecorder.addLongCounter(metrics.packetsRetransmitted, totalRetrans,
-                Collections.singletonList(target), labelValues);
+        if (channelReflectionAccessor == null) {
+          if (!channel.getClass().getName().equals(epollSocketChannelClassName)) {
+            return;
           }
-          if (metrics.recurringRetransmits != null) {
-            metricRecorder.addLongCounter(metrics.recurringRetransmits, retransmits,
-                Collections.singletonList(target), labelValues);
+          synchronized (accessorCache) {
+            channelReflectionAccessor = accessorCache.get(epollTcpInfoClassName);
+            if (channelReflectionAccessor == null) {
+              Class<?> tcpInfoClass = Class.forName(epollTcpInfoClassName);
+              Method tcpInfoMethod = channel.getClass().getMethod("tcpInfo", tcpInfoClass);
+              Constructor<?> tcpInfoConstructor = tcpInfoClass.getDeclaredConstructor();
+              Method totalRetransMethod = tcpInfoClass.getMethod("totalRetrans");
+              Method retransmitsMethod = tcpInfoClass.getMethod("retransmits");
+              Method rttMethod = tcpInfoClass.getMethod("rtt");
+
+              channelReflectionAccessor = new ChannelReflectionAccessor(
+                  tcpInfoMethod, tcpInfoConstructor, totalRetransMethod, retransmitsMethod,
+                  rttMethod);
+              accessorCache.put(epollTcpInfoClassName, channelReflectionAccessor);
+            }
           }
-          if (metrics.minRtt != null) {
-            metricRecorder.recordDoubleHistogram(metrics.minRtt,
-                rtt / 1000000.0, // Convert microseconds to seconds
+        }
+
+        Object info = channelReflectionAccessor.tcpInfoConstructor.newInstance();
+        channelReflectionAccessor.tcpInfoMethod.invoke(channel, info);
+
+        long totalRetrans = (Long) channelReflectionAccessor.totalRetransMethod.invoke(info);
+        int retransmits = (Integer) channelReflectionAccessor.retransmitsMethod.invoke(info);
+        long rtt = (Long) channelReflectionAccessor.rttMethod.invoke(info);
+
+        if (metrics.packetsRetransmitted != null) {
+          long delta = totalRetrans - lastTotalRetrans;
+          if (delta > 0) {
+            metricRecorder.addLongCounter(metrics.packetsRetransmitted, delta,
                 Collections.singletonList(target), labelValues);
           }
+          lastTotalRetrans = totalRetrans;
+        }
+        if (isClosed && metrics.recurringRetransmits != null) {
+          metricRecorder.addLongCounter(metrics.recurringRetransmits, retransmits,
+              Collections.singletonList(target), labelValues);
+        }
+        if (metrics.minRtt != null) {
+          metricRecorder.recordDoubleHistogram(metrics.minRtt,
+              rtt / 1000000.0, // Convert microseconds to seconds
+              Collections.singletonList(target), labelValues);
         }
       } catch (Throwable t) {
         // Epoll not available or error getting tcp_info, just ignore.

netty/src/test/java/io/grpc/netty/NettyClientTransportTest.java

@@ -1198,7 +1198,8 @@ private void startServer(int maxStreamsPerConnection, int maxHeaderListSize,
             MAX_RST_COUNT_DISABLED,
             0,
             Attributes.EMPTY,
-            channelz);
+            channelz,
+            null);
     server.start(serverListener);
     address = TestUtils.testServerAddress((InetSocketAddress) server.getListenSocketAddress());
     authority = GrpcUtil.authorityFromHostAndPort(address.getHostString(), address.getPort());

netty/src/test/java/io/grpc/netty/NettyServerBuilderTest.java

@@ -189,4 +189,15 @@ public void useNioTransport_shouldNotThrow() {
 
     builder.assertEventLoopsAndChannelType();
   }
+
+  @Test
+  public void setMetricRecorder_propagatedToServer() throws Exception {
+    io.grpc.MetricRecorder recorder = mock(io.grpc.MetricRecorder.class);
+    builder.setMetricRecorder(recorder);
+
+    NettyServer server = builder.buildTransportServers(
+        ImmutableList.<ServerStreamTracer.Factory>of());
+
+    assertThat(server.getMetricRecorder()).isSameInstanceAs(recorder);
+  }
 }

netty/src/test/java/io/grpc/netty/NettyServerTest.java

@@ -161,7 +161,8 @@ class NoHandlerProtocolNegotiator implements ProtocolNegotiator {
             0,
             0, // ignore
             Attributes.EMPTY,
-            channelz);
+            channelz,
+            null);
     final SettableFuture<Void> serverShutdownCalled = SettableFuture.create();
     ns.start(new ServerListener() {
       @Override
@@ -218,7 +219,8 @@ public void multiPortStartStopGet() throws Exception {
             0,
             0, // ignore
             Attributes.EMPTY,
-            channelz);
+            channelz,
+            null);
     final SettableFuture<Void> shutdownCompleted = SettableFuture.create();
     ns.start(new ServerListener() {
       @Override
@@ -298,7 +300,8 @@ public void multiPortConnections() throws Exception {
             0,
             0, // ignore
             Attributes.EMPTY,
-            channelz);
+            channelz,
+            null);
     final SettableFuture<Void> shutdownCompleted = SettableFuture.create();
     ns.start(new ServerListener() {
       @Override
@@ -366,7 +369,8 @@ public void getPort_notStarted() {
             0,
             0, // ignore
             Attributes.EMPTY,
-            channelz);
+            channelz,
+            null);
 
     assertThat(ns.getListenSocketAddress()).isEqualTo(addr);
     assertThat(ns.getListenSocketAddresses()).isEqualTo(addresses);
@@ -447,7 +451,8 @@ class TestProtocolNegotiator implements ProtocolNegotiator {
             0,
             0, // ignore
             eagAttributes,
-            channelz);
+            channelz,
+            null);
     ns.start(new ServerListener() {
       @Override
       public ServerTransportListener transportCreated(ServerTransport transport) {
@@ -501,7 +506,8 @@ public void channelzListenSocket() throws Exception {
             0,
             0, // ignore
             Attributes.EMPTY,
-            channelz);
+            channelz,
+            null);
     final SettableFuture<Void> shutdownCompleted = SettableFuture.create();
     ns.start(new ServerListener() {
       @Override
@@ -649,7 +655,8 @@ private NettyServer getServer(List<SocketAddress> addr, EventLoopGroup ev) {
         0,
         0, // ignore
         Attributes.EMPTY,
-        channelz);
+        channelz,
+        null);
   }
 
   private static class NoopServerTransportListener implements ServerTransportListener {