CosmosDb: Remove blocking calls from CFP

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/ChangeFeedObserver.java

@@ -3,6 +3,7 @@
 package com.azure.data.cosmos.internal.changefeed;
 
 import com.azure.data.cosmos.CosmosItemProperties;
+import reactor.core.publisher.Mono;
 
 import java.util.List;
 
@@ -30,6 +31,7 @@ public interface ChangeFeedObserver {
      *
      * @param context the context specifying partition for this observer, etc.
      * @param docs the documents changed.
+     * @return a deferred operation of this call.
      */
-    void processChanges(ChangeFeedObserverContext context, List<CosmosItemProperties> docs);
+    Mono<Void> processChanges(ChangeFeedObserverContext context, List<CosmosItemProperties> docs);
 }

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/exceptions/ObserverException.java

@@ -11,9 +11,9 @@ public class ObserverException extends RuntimeException {
     /**
      * Initializes a new instance of the {@link ObserverException} class using the specified internal exception.
      *
-     * @param originalException {@link Exception} thrown by the user code.
+     * @param originalException {@link Throwable} thrown by the user code.
      */
-    public ObserverException(Exception originalException) {
-        super(DefaultMessage, originalException.getCause());
+    public ObserverException(Throwable originalException) {
+        super(DefaultMessage, originalException);
     }
 }

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/AutoCheckpointer.java

@@ -7,6 +7,7 @@
 import com.azure.data.cosmos.internal.changefeed.ChangeFeedObserverCloseReason;
 import com.azure.data.cosmos.internal.changefeed.ChangeFeedObserverContext;
 import com.azure.data.cosmos.internal.changefeed.CheckpointFrequency;
+import reactor.core.publisher.Mono;
 
 import java.time.Duration;
 import java.time.ZoneId;
@@ -42,15 +43,24 @@ public void close(ChangeFeedObserverContext context, ChangeFeedObserverCloseReas
     }
 
     @Override
-    public void processChanges(ChangeFeedObserverContext context, List<CosmosItemProperties> docs) {
-        this.observer.processChanges(context, docs);
-        this.processedDocCount ++;
-
-        if (this.isCheckpointNeeded()) {
-            context.checkpoint().block();
-            this.processedDocCount = 0;
-            this.lastCheckpointTime = ZonedDateTime.now(ZoneId.of("UTC"));
+    public Mono<Void> processChanges(ChangeFeedObserverContext context, List<CosmosItemProperties> docs) {
+        return this.observer.processChanges(context, docs)
+            .then(this.afterProcessChanges(context));
+    }
+
+    private Mono<Void> afterProcessChanges(ChangeFeedObserverContext context) {
+        final AutoCheckpointer self = this;
+
+        self.processedDocCount ++;
+
+        if (self.isCheckpointNeeded()) {
+            return context.checkpoint()
+                .doOnSuccess((Void) -> {
+                    self.processedDocCount = 0;
+                    self.lastCheckpointTime = ZonedDateTime.now(ZoneId.of("UTC"));
+                });
         }
+        return Mono.empty();
     }
 
     private boolean isCheckpointNeeded() {

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/BootstrapperImpl.java

@@ -21,6 +21,9 @@ class BootstrapperImpl implements Bootstrapper {
     private final Duration lockTime;
     private final Duration sleepTime;
 
+    private boolean isInitialized;
+    private boolean isLockAcquired;
+
     public BootstrapperImpl(PartitionSynchronizer synchronizer, LeaseStore leaseStore, Duration lockTime, Duration sleepTime)
     {
         if (synchronizer == null) throw new IllegalArgumentException("synchronizer");
@@ -32,45 +35,48 @@ public BootstrapperImpl(PartitionSynchronizer synchronizer, LeaseStore leaseStor
         this.leaseStore = leaseStore;
         this.lockTime = lockTime;
         this.sleepTime = sleepTime;
+
+        this.isInitialized = false;
     }
 
     @Override
     public Mono<Void> initialize() {
-        BootstrapperImpl self = this;
-
-        return Mono.fromRunnable( () -> {
-            while (true) {
-                boolean initialized = self.leaseStore.isInitialized().block();
+        final BootstrapperImpl self = this;
+        self.isInitialized = false;
 
-                if (initialized) break;
+        return Mono.just(self)
+            .flatMap( value -> self.leaseStore.isInitialized())
+            .flatMap(initialized -> {
+                self.isInitialized = initialized;
 
-                boolean isLockAcquired = self.leaseStore.acquireInitializationLock(self.lockTime).block();
+                if (initialized) {
+                    return Mono.empty();
+                } else {
+                    return self.leaseStore.acquireInitializationLock(self.lockTime)
+                        .flatMap(lockAcquired -> {
+                            self.isLockAcquired = lockAcquired;
 
-                try {
-                    if (!isLockAcquired) {
-                        logger.info("Another instance is initializing the store");
-                        try {
-                            Thread.sleep(self.sleepTime.toMillis());
-                        } catch (InterruptedException ex) {
-                            logger.warn("Unexpected exception caught", ex);
-                        }
-                        continue;
-                    }
-
-                    logger.info("Initializing the store");
-                    self.synchronizer.createMissingLeases().block();
-                    self.leaseStore.markInitialized().block();
-
-                } catch (RuntimeException ex) {
-                    break;
-                } finally {
-                    if (isLockAcquired) {
-                        self.leaseStore.releaseInitializationLock().block();
-                    }
+                            if (!self.isLockAcquired) {
+                                logger.info("Another instance is initializing the store");
+                                return Mono.just(isLockAcquired).delayElement(self.sleepTime);
+                            } else {
+                                return self.synchronizer.createMissingLeases()
+                                    .then(self.leaseStore.markInitialized());
+                            }
+                        })
+                        .onErrorResume(throwable -> {
+                            logger.warn("Unexpected exception caught", throwable);
+                            return Mono.just(self.isLockAcquired);
+                        })
+                        .flatMap(lockAcquired -> {
+                            if (self.isLockAcquired) {
+                                return self.leaseStore.releaseInitializationLock();
+                            }
+                            return Mono.just(lockAcquired);
+                        });
                 }
-            }
-
-            logger.info("The store is initialized");
-        });
+            })
+            .repeat( () -> !self.isInitialized)
+            .then();
     }
 }

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/ChangeFeedContextClientImpl.java

@@ -20,6 +20,8 @@
 import com.azure.data.cosmos.SqlQuerySpec;
 import com.azure.data.cosmos.internal.PartitionKeyRange;
 import com.azure.data.cosmos.internal.changefeed.ChangeFeedContextClient;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
@@ -33,6 +35,8 @@
  * Implementation for ChangeFeedDocumentClient.
  */
 public class ChangeFeedContextClientImpl implements ChangeFeedContextClient {
+    private final Logger logger = LoggerFactory.getLogger(ChangeFeedContextClientImpl.class);
+
     private final AsyncDocumentClient documentClient;
     private final CosmosContainer cosmosContainer;
     private Scheduler rxScheduler;

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/ChangeFeedProcessorBuilderImpl.java

@@ -81,7 +81,19 @@ public class ChangeFeedProcessorBuilderImpl implements ChangeFeedProcessor.Build
      */
     @Override
     public Mono<Void> start() {
-        return partitionManager.start();
+        if (this.partitionManager == null) {
+            ChangeFeedProcessorBuilderImpl self = this;
+            return this.initializeCollectionPropertiesForBuild()
+                .then(self.getLeaseStoreManager()
+                    .flatMap(leaseStoreManager -> self.buildPartitionManager(leaseStoreManager)))
+                .flatMap(partitionManager1 -> {
+                    self.partitionManager = partitionManager1;
+                    return self.partitionManager.start();
+                });
+
+        } else {
+            return partitionManager.start();
+        }
     }
 
     /**
@@ -294,13 +306,7 @@ public ChangeFeedProcessor build() {
             this.executorService = Executors.newCachedThreadPool();
         }
 
-        // TBD: Move this initialization code as part of the start() call.
-        return this.initializeCollectionPropertiesForBuild()
-            .then(self.getLeaseStoreManager().flatMap(leaseStoreManager -> self.buildPartitionManager(leaseStoreManager)))
-            .map(partitionManager1 -> {
-                self.partitionManager = partitionManager1;
-                return self;
-            }).block();
+        return this;
     }
 
     public ChangeFeedProcessorBuilderImpl() {

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/DefaultObserver.java

@@ -8,12 +8,13 @@
 import com.azure.data.cosmos.internal.changefeed.ChangeFeedObserverContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Mono;
 
 import java.util.List;
 import java.util.function.Consumer;
 
 class DefaultObserver implements ChangeFeedObserver {
-    private final  Logger log = LoggerFactory.getLogger(DefaultObserver.class);
+    private static final Logger log = LoggerFactory.getLogger(DefaultObserver.class);
     private Consumer<List<CosmosItemProperties>> consumer;
 
     public DefaultObserver(Consumer<List<CosmosItemProperties>> consumer) {
@@ -31,9 +32,11 @@ public void close(ChangeFeedObserverContext context, ChangeFeedObserverCloseReas
     }
 
     @Override
-    public void processChanges(ChangeFeedObserverContext context, List<CosmosItemProperties> docs) {
+    public Mono<Void> processChanges(ChangeFeedObserverContext context, List<CosmosItemProperties> docs) {
         log.info("Start processing from thread {}", Thread.currentThread().getId());
         consumer.accept(docs);
         log.info("Done processing from thread {}", Thread.currentThread().getId());
+
+        return Mono.empty();
     }
 }

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/DocumentServiceLeaseUpdaterImpl.java

@@ -96,53 +96,6 @@ public Mono<Lease> updateLease(Lease cachedLease, CosmosItem itemLink, CosmosIte
                 }
                 return false;
             });
-
-//        Lease lease = cachedLease;
-//
-//        for (int retryCount = RETRY_COUNT_ON_CONFLICT; retryCount > 0; retryCount--) {
-//            lease = updateLease.apply(lease);
-//
-//            if (lease == null) {
-//                return Mono.empty();
-//            }
-//
-//            lease.setTimestamp(ZonedDateTime.now(ZoneId.of("UTC")));
-//            CosmosItemProperties leaseDocument = this.tryReplaceLease(lease, itemLink).block();
-//
-//            if (leaseDocument != null) {
-//                return Mono.just(ServiceItemLease.fromDocument(leaseDocument));
-//            }
-//
-//            // Partition lease update conflict. Reading the current version of lease.
-//            CosmosItemProperties document = null;
-//            try {
-//                CosmosItemResponse response = this.client.readItem(itemLink, requestOptions)
-//                    .block();
-//                document = response.properties();
-//            } catch (RuntimeException re) {
-//                if (re.getCause() instanceof CosmosClientException) {
-//                    CosmosClientException ex = (CosmosClientException) re.getCause();
-//                    if (ex.statusCode() == HTTP_STATUS_CODE_NOT_FOUND) {
-//                        // Partition lease no longer exists
-//                        throw new LeaseLostException(lease);
-//                    }
-//                }
-//                throw  re;
-//            }
-//
-//            ServiceItemLease serverLease = ServiceItemLease.fromDocument(document);
-//            logger.info(
-//                "Partition {} update failed because the lease with token '{}' was updated by host '{}' with token '{}'. Will retry, {} retry(s) left.",
-//                lease.getLeaseToken(),
-//                lease.getConcurrencyToken(),
-//                serverLease.getOwner(),
-//                serverLease.getConcurrencyToken(),
-//                retryCount);
-//
-//            lease = serverLease;
-//        }
-//
-//        throw new LeaseLostException(lease);
     }
 
     private Mono<CosmosItemProperties> tryReplaceLease(Lease lease, CosmosItem itemLink) throws LeaseLostException {