rsocket#1077 type safe discard element consumer

luczus · Dec 1, 2022 · 8a30105 · 8a30105
1 parent e31038f
commit 8a30105
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 18 deletions.
diff --git a/rsocket-core/src/main/java/io/rsocket/core/DefaultRSocketClient.java b/rsocket-core/src/main/java/io/rsocket/core/DefaultRSocketClient.java
@@ -22,6 +22,7 @@
 import io.rsocket.frame.FrameType;
 import java.util.AbstractMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
@@ -46,16 +47,20 @@
  */
 class DefaultRSocketClient extends ResolvingOperator<RSocket>
     implements CoreSubscriber<RSocket>, CorePublisher<RSocket>, RSocketClient {
-  static final Consumer<ReferenceCounted> DISCARD_ELEMENTS_CONSUMER =
-      referenceCounted -> {
-        if (referenceCounted.refCnt() > 0) {
-          try {
-            referenceCounted.release();
-          } catch (IllegalReferenceCountException e) {
-            // ignored
-          }
-        }
-      };
+  static final Consumer<?> DISCARD_ELEMENTS_CONSUMER =
+          item -> Optional.ofNullable(item)
+                  .filter(ReferenceCounted.class::isInstance)
+                  .map(ReferenceCounted.class::cast)
+                  .ifPresent(referenceCounted ->
+                  {
+                    if (referenceCounted.refCnt() > 0) {
+                      try {
+                        referenceCounted.release();
+                      } catch (IllegalReferenceCountException e) {
+                        // ignored
+                      }
+                    }
+                  });
 
   static final Object ON_DISCARD_KEY;
 

diff --git a/rsocket-core/src/main/java/io/rsocket/core/SendUtils.java b/rsocket-core/src/main/java/io/rsocket/core/SendUtils.java
@@ -33,20 +33,25 @@
 import io.rsocket.frame.RequestFireAndForgetFrameCodec;
 import io.rsocket.frame.RequestResponseFrameCodec;
 import io.rsocket.frame.RequestStreamFrameCodec;
+
+import java.util.Optional;
 import java.util.function.Consumer;
 import reactor.core.publisher.Operators;
 import reactor.util.context.Context;
 
 final class SendUtils {
   private static final Consumer<?> DROPPED_ELEMENTS_CONSUMER =
-      data -> {
-        try {
-          ReferenceCounted referenceCounted = (ReferenceCounted) data;
-          referenceCounted.release();
-        } catch (Throwable e) {
-          // ignored
-        }
-      };
+          data -> Optional.ofNullable(data)
+                  .filter(ReferenceCounted.class::isInstance)
+                  .map(ReferenceCounted.class::cast)
+                  .ifPresent(referenceCounted ->
+                  {
+                    try {
+                      referenceCounted.release();
+                    } catch (Throwable e) {
+                      // ignored
+                    }
+                  });
 
   static final Context DISCARD_CONTEXT = Operators.enableOnDiscard(null, DROPPED_ELEMENTS_CONSUMER);
 

diff --git a/rsocket-core/src/test/java/io/rsocket/core/DefaultRSocketClientTests.java b/rsocket-core/src/test/java/io/rsocket/core/DefaultRSocketClientTests.java
@@ -38,6 +38,7 @@
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.assertj.core.api.Assertions;
@@ -62,6 +63,8 @@
 import reactor.util.context.ContextView;
 import reactor.util.retry.Retry;
 
+import static org.mockito.Mockito.*;
+
 public class DefaultRSocketClientTests {
 
   ClientSocketRule rule;
@@ -637,6 +640,25 @@ public void shouldStartOriginalSourceOnceIfRacing() {
     }
   }
 
+
+  @Test
+  void discardElementsConsumerShouldAcceptOtherTypesThanReferenceCounted() {
+    Consumer discardElementsConsumer = DefaultRSocketClient.DISCARD_ELEMENTS_CONSUMER;
+    discardElementsConsumer.accept(new Object());
+  }
+
+  @Test
+  void droppedElementsConsumerReleaseReference() {
+    ReferenceCounted referenceCounted =  mock(ReferenceCounted.class);
+    when(referenceCounted.release()).thenReturn(true);
+    when(referenceCounted.refCnt()).thenReturn(1);
+
+    Consumer discardElementsConsumer = DefaultRSocketClient.DISCARD_ELEMENTS_CONSUMER;
+    discardElementsConsumer.accept(referenceCounted);
+
+    verify(referenceCounted).release();
+  }
+
   public static class ClientSocketRule extends AbstractSocketRule<RSocket> {
 
     protected RSocketClient client;

diff --git a/rsocket-core/src/test/java/io/rsocket/core/SendUtilsTest.java b/rsocket-core/src/test/java/io/rsocket/core/SendUtilsTest.java
@@ -0,0 +1,32 @@
+package io.rsocket.core;
+
+import io.netty.util.ReferenceCounted;
+import org.junit.jupiter.api.Test;
+
+import java.util.function.Consumer;
+
+import static org.mockito.Mockito.*;
+
+class SendUtilsTest {
+
+    @Test
+    void droppedElementsConsumerShouldAcceptOtherTypesThanReferenceCounted() {
+        Consumer value = extractDroppedElementConsumer();
+        value.accept(new Object());
+    }
+
+    @Test
+    void droppedElementsConsumerReleaseReference() {
+        ReferenceCounted referenceCounted =  mock(ReferenceCounted.class);
+        when(referenceCounted.release()).thenReturn(true);
+
+        Consumer value = extractDroppedElementConsumer();
+        value.accept(referenceCounted);
+
+        verify(referenceCounted).release();
+    }
+
+    private static Consumer<?> extractDroppedElementConsumer() {
+        return (Consumer<?>) SendUtils.DISCARD_CONTEXT.stream().findAny().get().getValue();
+    }
+}