Controller

common/src/main/java/com/flipkart/varadhi/Constants.java

@@ -11,6 +11,7 @@ public class Constants {
     public static final String USER_ID_HEADER = "x_user_id";
 
     public static final TopicCapacityPolicy DefaultTopicCapacity = new TopicCapacityPolicy(100, 400, 2);
+    public static String SYSTEM_IDENTITY = "System";
 
     public static class PathParams {
         public static final String PATH_PARAM_ORG = "org";

consumer/src/main/java/com/flipkart/varadhi/consumer/ConsumerApiMgr.java

@@ -16,16 +16,20 @@
 @Slf4j
 public class ConsumerApiMgr implements ConsumerApi {
     private final ConsumersManager consumersManager;
+    private final ConsumerInfo consumerInfo;
 
-    public ConsumerApiMgr(ConsumersManager consumersManager) {
+
+    public ConsumerApiMgr(ConsumersManager consumersManager, ConsumerInfo consumerInfo) {
         this.consumersManager = consumersManager;
+        this.consumerInfo = consumerInfo;
     }
 
     @Override
     public CompletableFuture<Void> start(ShardOperation.StartData operation) {
         log.info("Consumer: Starting shard {}", operation);
         VaradhiSubscription subscription = operation.getSubscription();
         SubscriptionUnitShard shard = operation.getShard();
+        consumerInfo.recordShardAssignment(subscription.getName(), shard.getShardId(), shard.getCapacityRequest());
         StorageSubscription<StorageTopic> mainSub = shard.getMainSubscription().getSubscriptionToConsume();
         ConsumptionFailurePolicy failurePolicy =
                 new ConsumptionFailurePolicy(subscription.getRetryPolicy(), shard.getRetrySubscription(),
@@ -47,6 +51,8 @@ public CompletableFuture<Void> start(ShardOperation.StartData operation) {
     @Override
     public CompletableFuture<Void> stop(ShardOperation.StopData operation) {
         VaradhiSubscription subscription = operation.getSubscription();
+        SubscriptionUnitShard shard = operation.getShard();
+        consumerInfo.purgeShardAssignment(subscription.getName(), shard.getShardId(), shard.getCapacityRequest());
         return consumersManager.stopSubscription(
                 subscription.getName(),
                 operation.getShardId()
@@ -60,7 +66,6 @@ public CompletableFuture<ShardStatus> getShardStatus(String subscriptionId, int
 
     @Override
     public CompletableFuture<ConsumerInfo> getConsumerInfo() {
-        //TODO::Return assignments as well.
-        return CompletableFuture.completedFuture(consumersManager.getInfo());
+        return CompletableFuture.completedFuture(consumerInfo);
     }
 }

consumer/src/main/java/com/flipkart/varadhi/consumer/ConsumersManager.java

@@ -1,7 +1,6 @@
 package com.flipkart.varadhi.consumer;
 
 import com.flipkart.varadhi.entities.*;
-import com.flipkart.varadhi.entities.cluster.ConsumerInfo;
 
 import java.util.concurrent.CompletableFuture;
 
@@ -37,6 +36,4 @@ CompletableFuture<Void> startSubscription(
     ConsumerState getConsumerState(String subscription, int shardId);
 
     // TODO likely need status on the starting / stopping as well; as the above status is for a running consumer..
-
-    ConsumerInfo getInfo();
 }

consumer/src/main/java/com/flipkart/varadhi/consumer/impl/ConsumersManagerImpl.java

@@ -4,17 +4,10 @@
 import com.flipkart.varadhi.consumer.ConsumersManager;
 import com.flipkart.varadhi.consumer.ConsumptionFailurePolicy;
 import com.flipkart.varadhi.entities.*;
-import com.flipkart.varadhi.entities.cluster.ConsumerInfo;
-import com.flipkart.varadhi.entities.TopicPartitions;
 
 import java.util.concurrent.CompletableFuture;
 
 public class ConsumersManagerImpl implements ConsumersManager {
-    private final ConsumerInfo consumerInfo;
-
-    public ConsumersManagerImpl(ConsumerInfo consumerInfo) {
-        this.consumerInfo = consumerInfo;
-    }
 
     @Override
     public CompletableFuture<Void> startSubscription(
@@ -44,9 +37,4 @@ public void resumeSubscription(String subscription, int shardId) {
     public ConsumerState getConsumerState(String subscription, int shardId) {
         return null;
     }
-
-    @Override
-    public ConsumerInfo getInfo() {
-        return consumerInfo;
-    }
 }

controller/src/main/java/com/flipkart/varadhi/controller/ControllerApiMgr.java

@@ -1,9 +1,11 @@
 package com.flipkart.varadhi.controller;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 
+import com.flipkart.varadhi.controller.config.ControllerConfig;
 import com.flipkart.varadhi.entities.cluster.*;
 import com.flipkart.varadhi.core.cluster.*;
 import com.flipkart.varadhi.entities.SubscriptionShards;
@@ -15,6 +17,8 @@
 import io.micrometer.core.instrument.MeterRegistry;
 import lombok.extern.slf4j.Slf4j;
 
+import static com.flipkart.varadhi.Constants.SYSTEM_IDENTITY;
+
 @Slf4j
 public class ControllerApiMgr implements ControllerApi {
     private final ShardAssigner shardAssigner;
@@ -23,28 +27,29 @@ public class ControllerApiMgr implements ControllerApi {
     private final OperationMgr operationMgr;
 
     public ControllerApiMgr(
-            ConsumerClientFactory consumerClientFactory, MetaStoreProvider metaStoreProvider,
+            ControllerConfig config, ConsumerClientFactory consumerClientFactory, MetaStoreProvider metaStoreProvider,
             MeterRegistry meterRegistry
     ) {
         this.consumerClientFactory = consumerClientFactory;
         this.shardAssigner = new ShardAssigner(metaStoreProvider.getAssignmentStore(), meterRegistry);
         this.metaStore = metaStoreProvider.getMetaStore();
-        this.operationMgr = new OperationMgr(metaStoreProvider.getOpStore());
+        this.operationMgr = new OperationMgr(config, metaStoreProvider.getOpStore());
     }
 
     public CompletableFuture<Void> addConsumerNodes(List<ConsumerNode> clusterConsumers) {
         return CompletableFuture.allOf(clusterConsumers.stream()
                         .map(cc -> getConsumerInfo(cc.getConsumerId()).thenAccept(
-                                cc::updateWithConsumerInfo)).toArray(CompletableFuture[]::new))
-                .thenAccept(v -> shardAssigner.addConsumerNodes(clusterConsumers));
+                                cc::initFromConsumerInfo)).toArray(CompletableFuture[]::new))
+                .thenCompose(v -> shardAssigner.addConsumerNodes(clusterConsumers));
     }
 
     @Override
     public CompletableFuture<SubscriptionStatus> getSubscriptionStatus(String subscriptionId, String requestedBy) {
         VaradhiSubscription subscription = metaStore.getSubscription(subscriptionId);
         return getSubscriptionStatus(subscription).exceptionally(t -> {
             throw new IllegalStateException(
-                    String.format("Failure in getting subscription status, try again after sometime. %s",
+                    String.format(
+                            "Failure in getting subscription status, try again after sometime. %s",
                             t.getMessage()
                     ));
         });
@@ -84,32 +89,31 @@ private SubscriptionStatus getSubscriptionStatusFromShardStatus(
     public CompletableFuture<SubscriptionOperation> startSubscription(
             String subscriptionId, String requestedBy
     ) {
-        //TODO:: Fix it -assignment failure is not failing the start op. Task failure in the operation mgr queue.
         VaradhiSubscription subscription = metaStore.getSubscription(subscriptionId);
         return getSubscriptionStatus(subscription).exceptionally(t -> {
             // If not temporary, then alternate needs to be provided to allow recovery from this.
             throw new IllegalStateException(
-                    String.format("Failure in getting subscription status, try again after sometime. %s",
+                    String.format(
+                            "Failure in getting subscription status, try again after sometime. %s",
                             t.getMessage()
                     ));
         }).thenApply(ss -> {
             if (ss.getState() == SubscriptionState.RUNNING || ss.getState() == SubscriptionState.STARTING) {
                 throw new InvalidOperationForResourceException("Subscription is already running or starting.");
             }
             log.info("Starting the Subscription: {}", subscriptionId);
-            // operationMgr is not expected to create a subOp and throw, so failure is not handled here.
-            // TODO:: fix this w.r.to failure in getOrCreateShardAssignment or its chain
-            return operationMgr.requestSubStart(
-                    subscriptionId, requestedBy, subOp -> getOrCreateShardAssignment(subscription).thenCompose(
-                            assignments -> startShards(subOp, subscription, assignments)));
+            SubscriptionOperation operation = SubscriptionOperation.startOp(subscriptionId, requestedBy);
+            operationMgr.createAndEnqueue(operation, subOp -> getOrCreateShardAssignment(subscription).thenCompose(
+                    assignments -> startShards((SubscriptionOperation) subOp, subscription, assignments)));
+            return operation;
         });
     }
 
     private CompletableFuture<List<Assignment>> getOrCreateShardAssignment(VaradhiSubscription subscription) {
         List<Assignment> assignedShards = shardAssigner.getSubscriptionAssignment(subscription.getName());
         if (assignedShards.isEmpty()) {
             List<SubscriptionUnitShard> unAssigned = getSubscriptionShards(subscription.getShards());
-            return shardAssigner.assignShard(unAssigned, subscription);
+            return shardAssigner.assignShard(unAssigned, subscription, new HashSet<>());
         } else {
             log.info(
                     "{} Shards for Subscription {} are already assigned.", assignedShards.size(),
@@ -145,13 +149,16 @@ private CompletableFuture<Void> startShard(
             // IsAssigned is started|starting|errored.
             // Stopping isn't considered as assigned, as start/stop shouldn't be running in parallel.
             if (!shardStatus.isAssigned()) {
-                ShardOperation shardOp = operationMgr.requestShardStart(subOpId, shard, subscription);
-                return consumer.start((ShardOperation.StartData) shardOp.getOpData()).whenComplete((v, t) -> {
-                    if (t != null) {
-                        markShardOpFailed(shardOp, t);
-                    } else {
-                        log.info("Scheduled shard start({}).", shardOp);
-                    }
+                ShardOperation startOp = ShardOperation.startOp(subOpId, shard, subscription);
+                return operationMgr.createAndExecute(startOp, op -> {
+                    ShardOperation shardOp = (ShardOperation) op;
+                    return consumer.start((ShardOperation.StartData) shardOp.getOpData()).whenComplete((v, t) -> {
+                        if (t != null) {
+                            markShardOpFailed(shardOp, t);
+                        } else {
+                            log.info("Scheduled shard start({}).", shardOp);
+                        }
+                    });
                 });
             } else {
                 log.info("Subscription:{} Shard:{} is already started. Skipping.", subId, shardId);
@@ -167,17 +174,18 @@ public CompletableFuture<SubscriptionOperation> stopSubscription(
         VaradhiSubscription subscription = metaStore.getSubscription(subscriptionId);
         return getSubscriptionStatus(subscription).exceptionally(t -> {
             throw new IllegalStateException(
-                    String.format("Failure in getting subscription status, try again after sometime. %s",
+                    String.format(
+                            "Failure in getting subscription status, try again after sometime. %s",
                             t.getMessage()
                     ));
         }).thenApply(ss -> {
             if (ss.getState() == SubscriptionState.STOPPED) {
                 throw new InvalidOperationForResourceException("Subscription is already stopped.");
             }
             log.info("Stopping the Subscription: {}", subscriptionId);
-            // operationMgr is not expected to create a subOp and throw, so failure is not handled here.
-            // TODO:: fix this w.r.to failure in stopShards
-            return operationMgr.requestSubStop(subscriptionId, requestedBy, subOp -> stopShards(subOp, subscription));
+            SubscriptionOperation operation = SubscriptionOperation.stopOp(subscriptionId, requestedBy);
+            operationMgr.createAndEnqueue(operation, subOp -> stopShards((SubscriptionOperation) subOp, subscription));
+            return operation;
         });
     }
 
@@ -187,7 +195,7 @@ private CompletableFuture<Void> stopShards(SubscriptionOperation subOp, VaradhiS
         SubscriptionShards shards = subscription.getShards();
         List<Assignment> assignments = shardAssigner.getSubscriptionAssignment(subId);
         log.info(
-                "Found {} assigned Shards for the Subscription:{} with Shards:{}.", assignments.size(),
+                "Found {} assigned Shards for the Subscription:{} with total {} Shards.", assignments.size(),
                 subId, shards.getShardCount()
         );
         CompletableFuture<Void> future = CompletableFuture.allOf(assignments.stream()
@@ -196,7 +204,7 @@ private CompletableFuture<Void> stopShards(SubscriptionOperation subOp, VaradhiS
                 )).toArray(CompletableFuture[]::new)).exceptionally(t -> {
             markSubOpFailed(subOp, t);
             return null;
-        }).thenCompose(v -> shardAssigner.unAssignShard(assignments, subscription));
+        }).thenCompose(v -> shardAssigner.unAssignShard(assignments, subscription, true));
         log.info("Scheduled Stop on {} shards for SubOp({}).", shards.getShardCount(), subOp.getData());
         return future;
     }
@@ -209,13 +217,16 @@ private CompletableFuture<Void> stopShard(
         ConsumerApi consumer = getAssignedConsumer(assignment);
         return consumer.getShardStatus(subId, shardId).thenCompose(shardStatus -> {
             if (!shardStatus.isAssigned()) {
-                ShardOperation shardOp = operationMgr.requestShardStop(subOpId, shard, subscription);
-                return consumer.stop((ShardOperation.StopData) shardOp.getOpData()).whenComplete((v, t) -> {
-                    if (t != null) {
-                        markShardOpFailed(shardOp, t);
-                    } else {
-                        log.info("Scheduled shard stop({}).", shardOp);
-                    }
+                ShardOperation stopOp = ShardOperation.stopOp(subOpId, shard, subscription);
+                return operationMgr.createAndExecute(stopOp, op -> {
+                    ShardOperation shardOp = (ShardOperation) op;
+                    return consumer.stop((ShardOperation.StopData) shardOp.getOpData()).whenComplete((v, t) -> {
+                        if (t != null) {
+                            markShardOpFailed(shardOp, t);
+                        } else {
+                            log.info("Scheduled shard stop({}).", shardOp);
+                        }
+                    });
                 });
             } else {
                 log.info("Subscription:{} Shard:{} is already Stopped. Skipping.", subId, shardId);
@@ -262,14 +273,31 @@ public CompletableFuture<Void> update(ShardOperation.OpData opData) {
     }
 
     public CompletableFuture<Void> consumerNodeLeft(String consumerNodeId) {
-        shardAssigner.consumerNodeLeft(consumerNodeId);
-        return CompletableFuture.completedFuture(null);
+        log.info("ConsumerNode {} left the cluster.", consumerNodeId);
+        return shardAssigner.consumerNodeLeft(consumerNodeId).thenAccept((v) -> {
+            List<Assignment> assignments = shardAssigner.getConsumerNodeAssignment(consumerNodeId);
+            assignments.forEach(assignment -> {
+                log.info("Assignment {} needs to be re-assigned", assignment);
+                SubscriptionOperation operation = SubscriptionOperation.reAssignShardOp(assignment, SYSTEM_IDENTITY);
+                operationMgr.createAndEnqueue(operation, subOp -> reAssignShard((SubscriptionOperation) subOp));
+            });
+        });
+    }
+
+    private CompletableFuture<Void> reAssignShard(SubscriptionOperation subOp) {
+        //TODO:: failure recovery of re-assign needs to be taken care of.
+        SubscriptionOperation.ReassignShardData data = (SubscriptionOperation.ReassignShardData) subOp.getData();
+        Assignment currentAssignment = data.getAssignment();
+        VaradhiSubscription subscription = metaStore.getSubscription(currentAssignment.getSubscriptionId());
+        SubscriptionUnitShard shard = subscription.getShards().getShard(currentAssignment.getShardId());
+        return shardAssigner.reAssignShard(currentAssignment, subscription, false).thenCompose(a ->
+                startShard(subOp.getId(), a, shard, subscription));
     }
 
     public CompletableFuture<Void> consumerNodeJoined(ConsumerNode consumerNode) {
-        return getConsumerInfo(consumerNode.getConsumerId()).thenAccept(ci -> {
-            consumerNode.updateWithConsumerInfo(ci);
-            shardAssigner.consumerNodeJoined(consumerNode);
+        return getConsumerInfo(consumerNode.getConsumerId()).thenCompose(ci -> {
+            consumerNode.initFromConsumerInfo(ci);
+            return shardAssigner.consumerNodeJoined(consumerNode);
         });
     }
 

controller/src/main/java/com/flipkart/varadhi/controller/OpExecutor.java

@@ -0,0 +1,9 @@
+package com.flipkart.varadhi.controller;
+
+import com.flipkart.varadhi.entities.cluster.Operation;
+
+import java.util.concurrent.CompletableFuture;
+
+public interface OpExecutor<T extends Operation> {
+    CompletableFuture<Void> execute(T operation);
+}