Implement rebalance mechanism in Pulsar Functions

pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Worker.java

@@ -29,6 +29,7 @@
 import org.apache.pulsar.functions.worker.rest.api.WorkerImpl;
 
 import javax.ws.rs.GET;
+import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
 import javax.ws.rs.core.MediaType;
@@ -114,4 +115,18 @@ public Map<String, Collection<String>> getAssignments() {
     public List<ConnectorDefinition> getConnectorsList() throws IOException {
         return worker.getListOfConnectors(clientAppId());
     }
+
+    @PUT
+    @ApiOperation(
+            value = "Triggers a rebalance of functions to workers"
+    )
+    @ApiResponses(value = {
+            @ApiResponse(code = 403, message = "The requester doesn't have admin permissions"),
+            @ApiResponse(code = 400, message = "Invalid request"),
+            @ApiResponse(code = 408, message = "Request timeout")
+    })
+    @Path("/rebalance")
+    public void rebalance() {
+        worker.rebalance(uri.getRequestUri(), clientAppId());
+    }
 }

pulsar-functions/runtime/src/main/java/org/apache/pulsar/functions/worker/WorkerConfig.java

@@ -70,6 +70,8 @@ public class WorkerConfig implements Serializable, PulsarConfiguration {
     @Category
     private static final String CATEGORY_FUNC_RUNTIME_MNG = "Function Runtime Management";
     @Category
+    private static final String CATEGORY_FUNC_SCHEDULE_MNG = "Function Scheduling Management";
+    @Category
     private static final String CATEGORY_SECURITY = "Common Security Settings (applied for both worker and client)";
     @Category
     private static final String CATEGORY_WORKER_SECURITY = "Worker Security Settings";
@@ -219,33 +221,33 @@ public class WorkerConfig implements Serializable, PulsarConfiguration {
     )
     private String stateStorageServiceUrl;
     @FieldContext(
-        category = CATEGORY_FUNC_METADATA_MNG,
+        category = CATEGORY_FUNC_RUNTIME_MNG,
         doc = "The pulsar topic used for storing function assignment informations"
     )
     private String functionAssignmentTopicName;
     @FieldContext(
-        category = CATEGORY_FUNC_METADATA_MNG,
+        category = CATEGORY_FUNC_SCHEDULE_MNG,
         doc = "The scheduler class used by assigning functions to workers"
     )
     private String schedulerClassName;
     @FieldContext(
-        category = CATEGORY_FUNC_METADATA_MNG,
+        category = CATEGORY_FUNC_RUNTIME_MNG,
         doc = "The frequency of failure checks, in milliseconds"
     )
     private long failureCheckFreqMs;
     @FieldContext(
-        category = CATEGORY_FUNC_METADATA_MNG,
+        category = CATEGORY_FUNC_RUNTIME_MNG,
         doc = "The reschedule timeout of function assignment, in milliseconds"
     )
     private long rescheduleTimeoutMs;
     @FieldContext(
-        category = CATEGORY_FUNC_METADATA_MNG,
+        category = CATEGORY_FUNC_RUNTIME_MNG,
         doc = "The max number of retries for initial broker reconnects when function metadata manager"
             + " tries to create producer on metadata topics"
     )
     private int initialBrokerReconnectMaxRetries;
     @FieldContext(
-        category = CATEGORY_FUNC_METADATA_MNG,
+        category = CATEGORY_FUNC_RUNTIME_MNG,
         doc = "The max number of retries for writing assignment to assignment topic"
     )
     private int assignmentWriteMaxRetries;

pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/SchedulerManager.java

@@ -184,7 +184,7 @@ public synchronized void initialize() {
         }
     }
 
-    public Future<?> schedule() {
+    private Future<?> scheduleInternal(Runnable runnable, String errMsg) {
         if (!leaderService.isLeader()) {
             return CompletableFuture.completedFuture(null);
         }
@@ -201,9 +201,9 @@ public Future<?> schedule() {
                     boolean isLeader = leaderService.isLeader();
                     if (isLeader) {
                         try {
-                            invokeScheduler();
+                            runnable.run();
                         } catch (Throwable th) {
-                            log.error("Encountered error when invoking scheduler", th);
+                            log.error("Encountered error when invoking scheduler", errMsg);
                             errorNotifier.triggerError(th);
                         }
                     }
@@ -218,22 +218,12 @@ public Future<?> schedule() {
         }
     }
 
-    private void scheduleCompaction(ScheduledExecutorService executor, long scheduleFrequencySec) {
-        if (executor != null) {
-            executor.scheduleWithFixedDelay(() -> {
-                if (leaderService.isLeader() && isCompactionNeeded.get()) {
-                    compactAssignmentTopic();
-                    isCompactionNeeded.set(false);
-                }
-            }, scheduleFrequencySec, scheduleFrequencySec, TimeUnit.SECONDS);
+    public Future<?> schedule() {
+        return scheduleInternal(() -> invokeScheduler(), "Encountered error when invoking scheduler");
+    }
 
-            executor.scheduleWithFixedDelay(() -> {
-                if (leaderService.isLeader() && metadataTopicLastMessage.compareTo(functionMetaDataManager.getLastMessageSeen()) != 0) {
-                    metadataTopicLastMessage = functionMetaDataManager.getLastMessageSeen();
-                    compactFunctionMetadataTopic();
-                }
-            }, scheduleFrequencySec, scheduleFrequencySec, TimeUnit.SECONDS);
-        }
+    public Future<?> rebalance() {
+        return scheduleInternal(() -> invokeRebalance(), "Encountered error when invoking rebalance");
     }
 
     @VisibleForTesting
@@ -331,7 +321,62 @@ void invokeScheduler() {
             // update message id associated with current view of assignments map
             lastMessageProduced = messageId;
         }
-
+    }
+
+    private void invokeRebalance() {
+
+        Set<String> currentMembership = membershipManager.getCurrentMembership()
+                .stream().map(workerInfo -> workerInfo.getWorkerId()).collect(Collectors.toSet());
+
+        Map<String, Map<String, Assignment>> workerIdToAssignments = functionRuntimeManager.getCurrentAssignments();
+
+        // filter out assignments of workers that are not currently in the active membership
+        List<Assignment> currentAssignments = workerIdToAssignments
+                .entrySet()
+                .stream()
+                .filter(workerIdToAssignmentEntry -> {
+                    String workerId = workerIdToAssignmentEntry.getKey();
+                    // remove assignments to workers that don't exist / died for now.
+                    // wait for failure detector to unassign them in the future for re-scheduling
+                    if (!currentMembership.contains(workerId)) {
+                        return false;
+                    }
+
+                    return true;
+                })
+                .flatMap(stringMapEntry -> stringMapEntry.getValue().values().stream())
+                .collect(Collectors.toList());
+
+
+        List<Assignment> rebalancedAssignments = scheduler.rebalance(currentAssignments, currentMembership);
+
+        for (Assignment assignment : rebalancedAssignments) {
+            MessageId messageId = publishNewAssignment(assignment, false);
+            // Directly update in memory assignment cache since I am leader
+            log.info("Rebalance - new assignment: {}", assignment);
+            functionRuntimeManager.processAssignment(assignment);
+            // update message id associated with current view of assignments map
+            lastMessageProduced = messageId;
+        }
+        log.info("Total number of new assignments computed for rebalance: {}", rebalancedAssignments.size());
+    }
+
+    private void scheduleCompaction(ScheduledExecutorService executor, long scheduleFrequencySec) {
+        if (executor != null) {
+            executor.scheduleWithFixedDelay(() -> {
+                if (leaderService.isLeader() && isCompactionNeeded.get()) {
+                    compactAssignmentTopic();
+                    isCompactionNeeded.set(false);
+                }
+            }, scheduleFrequencySec, scheduleFrequencySec, TimeUnit.SECONDS);
+
+            executor.scheduleWithFixedDelay(() -> {
+                if (leaderService.isLeader() && metadataTopicLastMessage.compareTo(functionMetaDataManager.getLastMessageSeen()) != 0) {
+                    metadataTopicLastMessage = functionMetaDataManager.getLastMessageSeen();
+                    compactFunctionMetadataTopic();
+                }
+            }, scheduleFrequencySec, scheduleFrequencySec, TimeUnit.SECONDS);
+        }
     }
 
     private void compactAssignmentTopic() {

pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/WorkerImpl.java

@@ -32,13 +32,18 @@
 import org.apache.pulsar.functions.worker.WorkerService;
 import org.apache.pulsar.functions.worker.WorkerUtils;
 
+import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.core.Response;
 import javax.ws.rs.core.Response.Status;
+import javax.ws.rs.core.UriBuilder;
 import java.io.IOException;
+import java.net.URI;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Future;
 import java.util.function.Supplier;
 
 import static com.google.common.base.Preconditions.checkNotNull;
@@ -48,6 +53,8 @@
 public class WorkerImpl {
 
     private final Supplier<WorkerService> workerServiceSupplier;
+    private Future<?> currentRebalanceFuture;
+
 
     public WorkerImpl(Supplier<WorkerService> workerServiceSupplier) {
         this.workerServiceSupplier = workerServiceSupplier;
@@ -198,4 +205,27 @@ public List<ConnectorDefinition> getListOfConnectors(String clientRole) {
 
         return this.worker().getConnectorsManager().getConnectors();
     }
+
+    public void rebalance(final URI uri, final String clientRole) {
+        if (!isWorkerServiceAvailable()) {
+            throwUnavailableException();
+        }
+
+        if (worker().getWorkerConfig().isAuthorizationEnabled() && !isSuperUser(clientRole)) {
+            log.error("Client [{}] is not authorized rebalance cluster", clientRole);
+            throw new RestException(Status.UNAUTHORIZED, "client is not authorize to perform operation");
+        }
+
+        if (worker().getLeaderService().isLeader()) {
+            if (currentRebalanceFuture == null || currentRebalanceFuture.isDone()) {
+                currentRebalanceFuture = this.worker().getSchedulerManager().rebalance();
+            } else {
+                throw new RestException(Status.BAD_REQUEST, "Rebalance already in progress");
+            }
+        } else {
+            WorkerInfo workerInfo = worker().getMembershipManager().getLeader();
+            URI redirect = UriBuilder.fromUri(uri).host(workerInfo.getWorkerHostname()).port(workerInfo.getPort()).build();
+            throw new WebApplicationException(Response.temporaryRedirect(redirect).build());
+        }
+    }
 }

pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/v2/WorkerApiV2Resource.java

@@ -39,10 +39,12 @@
 import javax.servlet.http.HttpServletRequest;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.GET;
+import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
+import javax.ws.rs.core.UriInfo;
 
 @Slf4j
 @Path("/worker")
@@ -59,6 +61,8 @@ public class WorkerApiV2Resource implements Supplier<WorkerService> {
     protected ServletContext servletContext;
     @Context
     protected HttpServletRequest httpRequest;
+    @Context
+    protected UriInfo uri;
 
     public WorkerApiV2Resource() {
         this.worker = new WorkerImpl(this);
@@ -138,4 +142,18 @@ public Map<String, Collection<String>> getAssignments() {
     public List<ConnectorDefinition> getConnectorsList() throws IOException {
         return worker.getListOfConnectors(clientAppId());
     }
+
+    @PUT
+    @ApiOperation(
+            value = "Triggers a rebalance of functions to workers"
+    )
+    @ApiResponses(value = {
+            @ApiResponse(code = 403, message = "The requester doesn't have admin permissions"),
+            @ApiResponse(code = 400, message = "Invalid request"),
+            @ApiResponse(code = 408, message = "Request timeout")
+    })
+    @Path("/rebalance")
+    public void rebalance() {
+        worker.rebalance(uri.getRequestUri(), clientAppId());
+    }
 }

pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/IScheduler.java

@@ -18,9 +18,11 @@
  */
 package org.apache.pulsar.functions.worker.scheduler;
 
+import org.apache.pulsar.functions.proto.Function;
 import org.apache.pulsar.functions.proto.Function.Assignment;
 import org.apache.pulsar.functions.proto.Function.Instance;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 
@@ -34,8 +36,24 @@ public interface IScheduler {
      * @param currentAssignments
      *            current assignments
      * @param workers
+     *            current list of active workers
      * @return
+     *            A list of new assignments
      */
     List<Assignment> schedule(List<Instance> unassignedFunctionInstances, List<Assignment> currentAssignments,
             Set<String> workers);
+
+    /**
+     * Rebalances function instances scheduled to workers
+     *
+     * @param currentAssignments
+     *            current assignments
+     * @param workers
+     *            current list of active workers
+     * @return
+     *            A list of new assignments
+     */
+    default List<Function.Assignment> rebalance(List<Function.Assignment> currentAssignments, Set<String> workers){
+        return Collections.emptyList();
+    }
 }