JVMCBC-1654 Replace recursive calls with waitForBucketConfig

(backported from commit 907ca15)
Reviewed-on: https://review.couchbase.org/c/couchbase-jvm-clients/+/228885

Motivation
----------
Some code paths were making recursive calls to
defer/retry until bucket topology is available.

These paths did not enforce timeouts, and would
retry forever if the bucket did not exist.

To make matters worse, each invocation added another
nested layer of CompletableFuture, leading to unconstrained
resource usage.

Modifications
-------------
Restructure the code to use waitForBucketConfig() instead
of scheduling recursively calls.

Simplify waitForBucketConfig/Topology(), and fix an issue
where it was failing to convert RetryExhaustedException
to our UnambiguousTimeoutException because
Exceptions.isRetryExhausted() was not recognizing
our repackaged version of RetryExhaustedException.

Change-Id: Iedbac8b75b7ec4ea428ccd22095b5dff438f37ee
Reviewed-on: https://review.couchbase.org/c/couchbase-jvm-clients/+/230380
Tested-by: Build Bot <[email protected]>
Reviewed-by: David Nault <[email protected]>

Author: dnault

core-io/src/main/java/com/couchbase/client/core/kv/RangeScanOrchestrator.java

@@ -20,9 +20,6 @@
 import com.couchbase.client.core.Reactor;
 import com.couchbase.client.core.annotation.Stability;
 import com.couchbase.client.core.api.kv.CoreKvParamValidators;
-import com.couchbase.client.core.config.BucketCapabilities;
-import com.couchbase.client.core.config.BucketConfig;
-import com.couchbase.client.core.config.CouchbaseBucketConfig;
 import com.couchbase.client.core.error.AuthenticationFailureException;
 import com.couchbase.client.core.error.CollectionNotFoundException;
 import com.couchbase.client.core.error.CouchbaseException;
@@ -38,10 +35,12 @@
 import com.couchbase.client.core.io.CollectionIdentifier;
 import com.couchbase.client.core.msg.CancellationReason;
 import com.couchbase.client.core.msg.ResponseStatus;
-import com.couchbase.client.core.msg.kv.MutationToken;
 import com.couchbase.client.core.msg.kv.RangeScanCancelRequest;
 import com.couchbase.client.core.msg.kv.RangeScanContinueRequest;
 import com.couchbase.client.core.msg.kv.RangeScanCreateRequest;
+import com.couchbase.client.core.topology.BucketCapability;
+import com.couchbase.client.core.topology.ClusterTopologyWithBucket;
+import com.couchbase.client.core.topology.CouchbaseBucketTopology;
 import reactor.core.Exceptions;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -51,12 +50,12 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 
+import static com.couchbase.client.core.util.BucketConfigUtil.waitForBucketTopology;
 import static com.couchbase.client.core.util.Validators.notNull;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
@@ -72,26 +71,9 @@ public class RangeScanOrchestrator {
   public static final int RANGE_SCAN_DEFAULT_BATCH_BYTE_LIMIT = 15000;
   public static final int RANGE_SCAN_DEFAULT_BATCH_ITEM_LIMIT = 50;
 
-  /**
-   * Holds the reference to core.
-   */
   private final Core core;
-
-  /**
-   * Holds the pointer to the keyspace that should be used.
-   */
   private final CollectionIdentifier collectionIdentifier;
 
-  /**
-   * Stores the current configuration if already set.
-   */
-  private volatile BucketConfig currentBucketConfig;
-
-  /**
-   * Checks if run against older clusters the feature might not be supported.
-   */
-  private volatile boolean capabilityEnabled = false;
-
   /**
    * Creates a new {@link RangeScanOrchestrator} which can be shared across calls.
    *
@@ -101,86 +83,77 @@ public class RangeScanOrchestrator {
   public RangeScanOrchestrator(final Core core, final CollectionIdentifier collectionIdentifier) {
     this.core = notNull(core, "Core");
     this.collectionIdentifier = notNull(collectionIdentifier, "CollectionIdentifier");
+  }
 
-    core.configurationProvider().configs().subscribe(cc -> {
-      BucketConfig bucketConfig = cc.bucketConfig(collectionIdentifier.bucket());
-      if (bucketConfig != null) {
-        currentBucketConfig = bucketConfig;
-        capabilityEnabled = bucketConfig.bucketCapabilities().contains(BucketCapabilities.RANGE_SCAN);
-      }
-    });
+  private Mono<CouchbaseBucketTopology> currentBucketTopology(CoreScanOptions options) {
+    Duration scanTimeout = options.commonOptions().timeout().orElse(core.environment().timeoutConfig().kvScanTimeout());
+    Duration timeout = min(scanTimeout, core.environment().timeoutConfig().connectTimeout()); // default scan timeout is 75 seconds; don't want to wait that long!
+
+    return waitForBucketTopology(core, collectionIdentifier.bucket(), timeout)
+      .map(ClusterTopologyWithBucket::bucket)
+      .cast(CouchbaseBucketTopology.class)
+      .onErrorMap(ClassCastException.class, t -> new FeatureNotAvailableException("Only Couchbase buckets are supported with KV Range Scan", t))
+
+      .filter(bucket -> bucket.hasCapability(BucketCapability.RANGE_SCAN))
+      .switchIfEmpty(Mono.error(FeatureNotAvailableException::rangeScan));
   }
 
-  /**
-   * Performs a range scan between a start and an end term (reactive).
-   *
-   * @param rangeScan
-   * @param options
-   * @return a {@link Flux} of returned items, or a failed flux during errors.
-   */
+  private static <T extends Comparable<T>> T min(T a, T b) {
+    return a.compareTo(b) < 0 ? a : b;
+  }
 
-    public Flux<CoreRangeScanItem> rangeScan(CoreRangeScan rangeScan, CoreScanOptions options) {
+  public Flux<CoreRangeScanItem> rangeScan(CoreRangeScan rangeScan, CoreScanOptions options) {
     return Flux.defer(() -> {
       CoreKvParamValidators.validateScanParams(rangeScan, options);
 
-      if (currentBucketConfig == null) {
-        // We might not have a config yet if bootstrap is still in progress, wait 100ms
-        // and then try again. In a steady state this should not happen.
-        return Mono
-          .delay(Duration.ofMillis(100), core.context().environment().scheduler())
-          .flatMapMany(ign -> rangeScan(rangeScan, options));
-      } else if (!(currentBucketConfig instanceof CouchbaseBucketConfig)) {
-        return Flux.error(new IllegalStateException("Only Couchbase buckets are supported with KV Range Scan"));
-      }
-      Map<Short, MutationToken> consistencyMap=options.consistencyMap();
-      return streamForPartitions((partition, start) -> {
-        byte[] actualStartTerm = start == null ? rangeScan.from().id().getBytes(UTF_8) : start;
-        return RangeScanCreateRequest.forRangeScan(actualStartTerm, rangeScan, options, partition, core.context(),
-            collectionIdentifier, consistencyMap);
-      }, options);
+      return currentBucketTopology(options)
+        .flatMapMany(topology -> streamForPartitions(
+          topology,
+          options,
+          (partition, start) -> {
+            byte[] actualStartTerm = start == null ? rangeScan.from().id().getBytes(UTF_8) : start;
+            return RangeScanCreateRequest.forRangeScan(
+              actualStartTerm,
+              rangeScan,
+              options,
+              partition,
+              core.context(),
+              collectionIdentifier,
+              options.consistencyMap()
+            );
+          }));
     });
   }
 
-  /**
-   * Performs a sampling scan (reactive).
-   *
-   * @param samplingScan
-   * @param options
-   * @return a {@link Flux} of returned items, or a failed flux during errors.
-   */
   public Flux<CoreRangeScanItem> samplingScan(CoreSamplingScan samplingScan, CoreScanOptions options) {
-    return Flux
-      .defer(() -> {
-        CoreKvParamValidators.validateScanParams(samplingScan, options);
-        return Mono.just(0);
-      })
-      .thenMany(Flux.defer(() -> {
-        if (currentBucketConfig == null) {
-          // We might not have a config yet if bootstrap is still in progress, wait 100ms
-          // and then try again. In a steady state this should not happen.
-          return Mono
-            .delay(Duration.ofMillis(100), core.context().environment().scheduler())
-            .flatMapMany(ign -> samplingScan(samplingScan, options));
-        } else if (!(currentBucketConfig instanceof CouchbaseBucketConfig)) {
-          return Flux.error(new IllegalStateException("Only Couchbase buckets are supported with KV Range Scan"));
-        }
-        Map<Short, MutationToken> consistencyMap=options.consistencyMap();
-        return streamForPartitions((partition, ignored) -> RangeScanCreateRequest.forSamplingScan(samplingScan,
-              options, partition, core.context(), collectionIdentifier, consistencyMap), options);
-        }).take(samplingScan.limit()));
-  }
-
-  @SuppressWarnings("unchecked")
-  private Flux<CoreRangeScanItem> streamForPartitions(final BiFunction<Short, byte[], RangeScanCreateRequest> createSupplier,
-                                                      final CoreScanOptions options) {
+    return Flux.defer(() -> {
+      CoreKvParamValidators.validateScanParams(samplingScan, options);
 
-    if (!capabilityEnabled) {
-      return Flux.error(FeatureNotAvailableException.rangeScan());
-    }
+      return currentBucketTopology(options)
+        .flatMapMany(topology -> streamForPartitions(
+          topology,
+          options,
+          (partition, ignored) -> RangeScanCreateRequest.forSamplingScan(
+            samplingScan,
+            options,
+            partition,
+            core.context(),
+            collectionIdentifier,
+            options.consistencyMap()
+          )
+        ))
+        .take(samplingScan.limit());
+    });
+  }
 
+  private Flux<CoreRangeScanItem> streamForPartitions(
+    CouchbaseBucketTopology topology,
+    CoreScanOptions options,
+    BiFunction<Short, byte[], RangeScanCreateRequest> createSupplier
+  ) {
     final AtomicLong itemsStreamed = new AtomicLong();
 
-    int numPartitions = ((CouchbaseBucketConfig) currentBucketConfig).numberOfPartitions();
+    int numPartitions = topology.numberOfPartitions();
 
     List<Flux<CoreRangeScanItem>> partitionStreams = new ArrayList<>(numPartitions);
     for (short i = 0; i < numPartitions; i++) {