Merge pull request #1310 from smallrye/bug/uni-memoize-protocol

smallrye · Jun 26, 2023 · 7883a82 · 7883a82
2 parents 4c4fc0f + 03757ee
commit 7883a82
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/implementation/src/main/java/io/smallrye/mutiny/operators/uni/UniMemoizeOp.java b/implementation/src/main/java/io/smallrye/mutiny/operators/uni/UniMemoizeOp.java
@@ -54,17 +54,19 @@ public void subscribe(UniSubscriber<? super I> subscriber) {
                 case INIT:
                     state = State.WAITING_FOR_UPSTREAM;
                     awaiters.add(subscriber);
+                    subscriber.onSubscribe(new MemoizedSubscription(subscriber));
                     currentContext = subscriber.context();
                     upstream().subscribe().withSubscriber(this);
                     break;
                 case WAITING_FOR_UPSTREAM:
                     awaiters.add(subscriber);
+                    subscriber.onSubscribe(new MemoizedSubscription(subscriber));
                     break;
                 case CACHING:
+                    subscriber.onSubscribe(new MemoizedSubscription(subscriber));
                     forwardTo(subscriber);
                     break;
             }
-            subscriber.onSubscribe(new MemoizedSubscription(subscriber));
         } finally {
             internalLock.unlock();
         }

diff --git a/implementation/src/test/java/io/smallrye/mutiny/groups/UniMemoizeTest.java b/implementation/src/test/java/io/smallrye/mutiny/groups/UniMemoizeTest.java
@@ -534,4 +534,19 @@ public void reproducer_1303() throws ExecutionException, InterruptedException {
         }
     }
 
+    @Test
+    public void checkProtocolCorrectness() {
+        var log = new ArrayList<String>();
+        var uni = Uni.createFrom().item(() -> 58).memoize().indefinitely()
+                .onSubscription().invoke(() -> log.add("sub"))
+                .onItem().invoke(n -> log.add(String.valueOf(n)));
+
+        Integer res = uni.await().atMost(Duration.ofSeconds(5));
+        assertThat(res).isEqualTo(58);
+        assertThat(log).containsExactly("sub", "58");
+
+        uni.await().atMost(Duration.ofSeconds(5));
+        assertThat(res).isEqualTo(58);
+        assertThat(log).containsExactly("sub", "58", "sub", "58");
+    }
 }