loadbalancer reconnects when downstream connection goes away (#547)

* updates to netty's duplex connection

* loadbalancer reconnects when downstream connection goes away

* fixed loadbalancer tests

Author: robertroeser

.gitignore

@@ -65,7 +65,7 @@ atlassian-ide-plugin.xml
 
 # NetBeans specific files/directories
 .nbattrs
-/bin
+**/bin/*
 
 #.gitignore in subdirectory
 .gitignore

rsocket-core/build.gradle

@@ -27,7 +27,6 @@ dependencies {
     api 'io.netty:netty-buffer'
     api 'io.projectreactor:reactor-core'
 
-    implementation 'org.jctools:jctools-core'
     implementation 'org.slf4j:slf4j-api'
 
     compileOnly 'com.google.code.findbugs:jsr305'

rsocket-core/src/main/java/io/rsocket/StreamIdSupplier.java

@@ -16,28 +16,31 @@
 
 package io.rsocket;
 
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+
 final class StreamIdSupplier {
 
-  private int streamId;
+  private static final AtomicIntegerFieldUpdater<StreamIdSupplier> STREAM_ID =
+      AtomicIntegerFieldUpdater.newUpdater(StreamIdSupplier.class, "streamId");
+  private volatile int streamId;
 
   private StreamIdSupplier(int streamId) {
     this.streamId = streamId;
   }
 
-  synchronized int nextStreamId() {
-    streamId += 2;
-    return streamId;
-  }
-
-  synchronized boolean isBeforeOrCurrent(int streamId) {
-    return this.streamId >= streamId && streamId > 0;
-  }
-
   static StreamIdSupplier clientSupplier() {
     return new StreamIdSupplier(-1);
   }
 
   static StreamIdSupplier serverSupplier() {
     return new StreamIdSupplier(0);
   }
+
+  int nextStreamId() {
+    return STREAM_ID.addAndGet(this, 2);
+  }
+
+  boolean isBeforeOrCurrent(int streamId) {
+    return this.streamId >= streamId && streamId > 0;
+  }
 }

rsocket-core/src/main/java/io/rsocket/fragmentation/FragmentationDuplexConnection.java

@@ -16,21 +16,24 @@
 
 package io.rsocket.fragmentation;
 
-import static io.rsocket.fragmentation.FrameReassembler.createFrameReassembler;
-import static io.rsocket.util.AbstractionLeakingFrameUtils.toAbstractionLeakingFrame;
-
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.PooledByteBufAllocator;
+import io.netty.util.collection.IntObjectHashMap;
+import io.netty.util.collection.LongObjectHashMap;
 import io.rsocket.DuplexConnection;
 import io.rsocket.Frame;
 import io.rsocket.util.AbstractionLeakingFrameUtils;
 import io.rsocket.util.NumberUtils;
-import java.util.Objects;
-import org.jctools.maps.NonBlockingHashMapLong;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import java.util.Collection;
+import java.util.Objects;
+
+import static io.rsocket.fragmentation.FrameReassembler.createFrameReassembler;
+import static io.rsocket.util.AbstractionLeakingFrameUtils.toAbstractionLeakingFrame;
+
 /**
  * A {@link DuplexConnection} implementation that fragments and reassembles {@link Frame}s.
  *
@@ -46,8 +49,7 @@ public final class FragmentationDuplexConnection implements DuplexConnection {
 
   private final FrameFragmenter frameFragmenter;
 
-  private final NonBlockingHashMapLong<FrameReassembler> frameReassemblers =
-      new NonBlockingHashMapLong<>();
+  private final IntObjectHashMap<FrameReassembler> frameReassemblers = new IntObjectHashMap<>();
 
   /**
    * Creates a new instance.
@@ -85,7 +87,16 @@ public FragmentationDuplexConnection(
 
     delegate
         .onClose()
-        .doFinally(signalType -> frameReassemblers.values().forEach(FrameReassembler::dispose))
+        .doFinally(
+            signalType -> {
+              Collection<FrameReassembler> values;
+              synchronized (this) {
+                values = frameReassemblers.values();
+              }
+              for (FrameReassembler reassembler : values) {
+                reassembler.dispose();
+              }
+            })
         .subscribe();
   }
 
@@ -134,9 +145,13 @@ private Flux<Frame> toFragmentedFrames(int streamId, io.rsocket.framing.Frame fr
   }
 
   private Mono<Frame> toReassembledFrames(int streamId, io.rsocket.framing.Frame fragment) {
-    FrameReassembler frameReassembler =
-        frameReassemblers.computeIfAbsent(
-            (long) streamId, i -> createFrameReassembler(byteBufAllocator));
+    FrameReassembler frameReassembler;
+
+    synchronized (this) {
+      frameReassembler =
+          frameReassemblers.computeIfAbsent(
+               streamId, i -> createFrameReassembler(byteBufAllocator));
+    }
 
     return Mono.justOrEmpty(frameReassembler.reassemble(fragment))
         .map(frame -> toAbstractionLeakingFrame(byteBufAllocator, streamId, frame));