From 45dfc53d402480fa1bef779a329ba014ca6f40a4 Mon Sep 17 00:00:00 2001
From: Eric Staffas <114935344+eric-picnic@users.noreply.github.com>
Date: Wed, 26 Oct 2022 08:28:35 +0200
Subject: [PATCH] Prefer `Flux#take(long, boolean)` over `Flux#take(long)` to
 limit upstream generation (#314)

---
 .../refasterrules/ReactorRules.java           | 31 +++++++++++++++++++
 .../refasterrules/ReactorRulesTestInput.java  |  4 +++
 .../refasterrules/ReactorRulesTestOutput.java |  4 +++
 3 files changed, 39 insertions(+)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
index 1b3bf9c3e9..d2eb81d06d 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
@@ -1,6 +1,7 @@
 package tech.picnic.errorprone.refasterrules;
 
 import static com.google.common.collect.MoreCollectors.toOptional;
+import static com.google.errorprone.BugPattern.SeverityLevel.WARNING;
 import static com.google.errorprone.refaster.ImportPolicy.STATIC_IMPORT_ALWAYS;
 import static java.util.function.Function.identity;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -28,7 +29,9 @@
 import reactor.test.StepVerifier;
 import reactor.test.publisher.PublisherProbe;
 import reactor.util.context.Context;
+import tech.picnic.errorprone.refaster.annotation.Description;
 import tech.picnic.errorprone.refaster.annotation.OnlineDocumentation;
+import tech.picnic.errorprone.refaster.annotation.Severity;
 import tech.picnic.errorprone.refaster.matchers.ThrowsCheckedException;
 
 /** Refaster rules related to Reactor expressions and statements. */
@@ -144,6 +147,34 @@ Mono<S> after(Mono<T> mono, S object) {
     }
   }
 
+  /**
+   * Prefer {@link Flux#take(long, boolean)} over {@link Flux#take(long)}.
+   *
+   * <p>In Reactor versions prior to 3.5.0, {@code Flux#take(long)} makes an unbounded request
+   * upstream, and is equivalent to {@code Flux#take(long, false)}. In 3.5.0, the behavior of {@code
+   * Flux#take(long)} will change to that of {@code Flux#take(long, true)}.
+   *
+   * <p>The intent with this Refaster rule is to get the new behavior before upgrading to Reactor
+   * 3.5.0.
+   */
+  // XXX: Drop this rule some time after upgrading to Reactor 3.6.0, or introduce a way to apply
+  // this rule only when an older version of Reactor is on the classpath.
+  // XXX: Once Reactor 3.6.0 is out, introduce a rule that rewrites code in the opposite direction.
+  @Description(
+      "Prior to Reactor 3.5.0, `take(n)` requests and unbounded number of elements upstream.")
+  @Severity(WARNING)
+  static final class FluxTake<T> {
+    @BeforeTemplate
+    Flux<T> before(Flux<T> flux, long n) {
+      return flux.take(n);
+    }
+
+    @AfterTemplate
+    Flux<T> after(Flux<T> flux, long n) {
+      return flux.take(n, /* limitRequest= */ true);
+    }
+  }
+
   /** Don't unnecessarily pass an empty publisher to {@link Mono#switchIfEmpty(Mono)}. */
   static final class MonoSwitchIfEmptyOfEmptyPublisher<T> {
     @BeforeTemplate
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
index 270dce8306..60c0902de4 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
@@ -58,6 +58,10 @@ Mono<String> testMonoThenReturn() {
     return Mono.empty().then(Mono.just("foo"));
   }
 
+  Flux<Integer> testFluxTake() {
+    return Flux.just(1, 2, 3).take(1);
+  }
+
   Mono<Integer> testMonoSwitchIfEmptyOfEmptyPublisher() {
     return Mono.just(1).switchIfEmpty(Mono.empty());
   }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
index 2c790c959b..e24d8157fe 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
@@ -60,6 +60,10 @@ Mono<String> testMonoThenReturn() {
     return Mono.empty().thenReturn("foo");
   }
 
+  Flux<Integer> testFluxTake() {
+    return Flux.just(1, 2, 3).take(1, true);
+  }
+
   Mono<Integer> testMonoSwitchIfEmptyOfEmptyPublisher() {
     return Mono.just(1);
   }