Fix SetupPayload being released too early

Author: rdegnan

gradle.properties

@@ -1,2 +1,2 @@
 group=io.netifi.proteus
-version=1.5.3
+version=1.5.4

proteus-client/src/main/java/io/netifi/proteus/DefaultProteusBrokerService.java

@@ -17,14 +17,12 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.Unpooled;
-import io.netty.util.ReferenceCountUtil;
 import io.opentracing.Tracer;
 import io.rsocket.Payload;
 import io.rsocket.RSocket;
 import io.rsocket.rpc.rsocket.RequestHandlingRSocket;
 import io.rsocket.transport.ClientTransport;
 import io.rsocket.util.ByteBufPayload;
-import io.rsocket.util.DefaultPayload;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
 import java.time.Duration;
@@ -137,15 +135,10 @@ static Payload getSetupPayload(
       String group,
       long accessKey,
       ByteBuf accessToken) {
-    ByteBuf metadata = null;
-    try {
-      metadata =
-          DestinationSetupFlyweight.encode(
-              alloc, computedFromDestination, group, accessKey, accessToken);
-      return DefaultPayload.create(Unpooled.EMPTY_BUFFER, metadata);
-    } finally {
-      ReferenceCountUtil.safeRelease(metadata);
-    }
+    ByteBuf metadata =
+        DestinationSetupFlyweight.encode(
+            alloc, computedFromDestination, group, accessKey, accessToken);
+    return ByteBufPayload.create(Unpooled.EMPTY_BUFFER, metadata);
   }
 
   private synchronized void reconcileSuppliers(Set<Broker> incomingBrokers) {

proteus-client/src/main/java/io/netifi/proteus/rsocket/WeightedReconnectingRSocket.java

@@ -183,7 +183,7 @@ synchronized void resetStatistics() {
     errorPercentage.reset(1.0);
   }
 
-  private RSocketFactory.ClientRSocketFactory getClientFactory(String destination) {
+  private RSocketFactory.ClientRSocketFactory getClientFactory(Payload setupPayload) {
     RSocketFactory.ClientRSocketFactory connect =
         RSocketFactory.connect().frameDecoder(Frame::retain);
 
@@ -203,7 +203,7 @@ private RSocketFactory.ClientRSocketFactory getClientFactory(String destination)
     }
 
     return connect
-        .setupPayload(setupPayloadSupplier.apply(destination))
+        .setupPayload(setupPayload)
         .keepAliveAckTimeout(Duration.ofSeconds(0))
         .keepAliveTickPeriod(Duration.ofSeconds(0));
   }
@@ -228,9 +228,10 @@ void connect() {
                   WeightedClientTransportSupplier weighedClientTransportSupplier =
                       transportSupplier.get();
                   String destination = destinationNameFactory.get();
+                  Payload setupPayload = setupPayloadSupplier.apply(destination);
 
                   long start = System.nanoTime();
-                  return getClientFactory(destination)
+                  return getClientFactory(setupPayload)
                       .errorConsumer(
                           throwable ->
                               logger.error(
@@ -276,7 +277,8 @@ void connect() {
                                     })
                                 .subscribe();
                             setRSocket(rSocket);
-                          });
+                          })
+                      .doFinally(signalType -> setupPayload.release());
                 }))
         .doOnError(t -> logger.error("error trying to broker", t))
         .retry()

proteus-client/src/test/java/io/netifi/proteus/ProteusBrokerServiceTest.java

@@ -18,6 +18,7 @@ public void setupPayloadLeakTest() {
     for (int i = 0; i < 100000; i++) {
       Payload payload =
           DefaultProteusBrokerService.getSetupPayload(alloc, "foo", "bar", 123L, token);
+      payload.release();
     }
     Assert.assertEquals(0, directBuffersCount(alloc));
     Assert.assertEquals(0, heapBuffersCount(alloc));