Fix groupBy initial request off by one (#2450)

This commit fixes the following misalignment:

 1. drainLoop in GroupByMain does s.request(e) when e groups are
 produced to the downstream. Each group at the moment has at least 1
 element enqueued into it.
 2. When UnicastGroupedFlux consumes elements for the first time, it
 requests the first element and does main.s.request(e) where e covers
 the first element
 3. That said that GroupByMain fulfilled the demand for the first
 element and then UnicastGroupedFlux does the same for the second time.
 It leads that for each group we have 1 extra demand which then ends up
 with overflow.
 4. To avoid that, this PR adds isFirstRequest check which allows
 removing that redundant demand for the first element on the first
 requestN

Signed-off-by: Oleh Dokuka [email protected]

Author: OlegDokuka

reactor-core/src/main/java/reactor/core/publisher/FluxGroupBy.java

@@ -508,6 +508,7 @@ public K key() {
 		volatile boolean outputFused;
 
 		int produced;
+		boolean isFirstRequest = true;
 
 		UnicastGroupedFlux(K key,
 				Queue<V> queue,
@@ -565,7 +566,16 @@ void drainRegular(Subscriber<? super V> a) {
 				if (e != 0) {
 					GroupByMain<?, K, V> main = parent;
 					if (main != null) {
-						main.s.request(e);
+						if (this.isFirstRequest) {
+							this.isFirstRequest = false;
+							long toRequest = e - 1;
+
+							if (toRequest > 0) {
+								main.s.request(toRequest);
+							}
+						} else {
+							main.s.request(e);
+						}
 					}
 					if (r != Long.MAX_VALUE) {
 						REQUESTED.addAndGet(this, -e);
@@ -744,7 +754,16 @@ void tryReplenish() {
 				produced = 0;
 				GroupByMain<?, K, V> main = parent;
 				if (main != null) {
-					main.s.request(p);
+					if (this.isFirstRequest) {
+						this.isFirstRequest = false;
+						p--;
+
+						if (p > 0) {
+							main.s.request(p);
+						}
+					} else {
+						main.s.request(p);
+					}
 				}
 			}
 		}

reactor-core/src/test/java/reactor/core/publisher/FluxGroupByTest.java

@@ -374,6 +374,37 @@ public void twoGroupsLongAsyncMergeHidden() {
 		  .assertComplete();
 	}
 
+	@Test
+	public void twoGroupsLongAsyncMergeHidden2() {
+		ForkJoinPool forkJoinPool = new ForkJoinPool();
+
+		for (int j = 0; j < 100; j++) {
+			AssertSubscriber<Long> ts = AssertSubscriber.create();
+			AtomicLong dropped = new AtomicLong();
+
+			Hooks.onNextDropped(__ -> dropped.incrementAndGet());
+			try {
+				final int total = 100_000;
+				Flux.range(0, total)
+				    .groupBy(i -> (i / 2d) * 2d, 42)
+				    .flatMap(it -> it.take(1)
+				                     .hide(), 2)
+				    .publishOn(Schedulers.fromExecutorService(forkJoinPool), 2)
+				    .count()
+				    .subscribe(ts);
+
+				ts.await(Duration.ofSeconds(50));
+
+				ts.assertValues(total - dropped.get())
+				  .assertNoError()
+				  .assertComplete();
+
+			} finally {
+				Hooks.resetOnNextDropped();
+			}
+		}
+	}
+
 	@Test
 	public void twoGroupsConsumeWithSubscribe() {
 		ForkJoinPool forkJoinPool = new ForkJoinPool();