From 0ac832cb06a8ceb65e0baf249d2d15c0c85885cc Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 1 Jul 2020 14:09:50 -0700 Subject: [PATCH] Implement rebalance mechanism in Pulsar Functions --- .../apache/pulsar/broker/admin/v2/Worker.java | 15 +++ .../pulsar/functions/worker/WorkerConfig.java | 14 +- .../functions/worker/SchedulerManager.java | 83 +++++++++--- .../functions/worker/rest/api/WorkerImpl.java | 30 +++++ .../rest/api/v2/WorkerApiV2Resource.java | 18 +++ .../worker/scheduler/IScheduler.java | 18 +++ .../worker/scheduler/RoundRobinScheduler.java | 120 ++++++++++++++++++ .../scheduler/RoundRobinSchedulerTest.java | 78 ++++++++++++ 8 files changed, 351 insertions(+), 25 deletions(-) create mode 100644 pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinSchedulerTest.java diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Worker.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Worker.java index 332a43f540b66..d4638276e9a09 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Worker.java +++ b/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> getAssignments() { public List 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()); + } } diff --git a/pulsar-functions/runtime/src/main/java/org/apache/pulsar/functions/worker/WorkerConfig.java b/pulsar-functions/runtime/src/main/java/org/apache/pulsar/functions/worker/WorkerConfig.java index c056028fc6808..021c12ade90ef 100644 --- a/pulsar-functions/runtime/src/main/java/org/apache/pulsar/functions/worker/WorkerConfig.java +++ b/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; diff --git a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/SchedulerManager.java b/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/SchedulerManager.java index 0cd682a99746d..43413102e77c1 100644 --- a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/SchedulerManager.java +++ b/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 currentMembership = membershipManager.getCurrentMembership() + .stream().map(workerInfo -> workerInfo.getWorkerId()).collect(Collectors.toSet()); + + Map> workerIdToAssignments = functionRuntimeManager.getCurrentAssignments(); + + // filter out assignments of workers that are not currently in the active membership + List 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 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() { diff --git a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/WorkerImpl.java b/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/WorkerImpl.java index 88e74689f87be..f58f57bda1435 100644 --- a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/WorkerImpl.java +++ b/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 workerServiceSupplier; + private Future currentRebalanceFuture; + public WorkerImpl(Supplier workerServiceSupplier) { this.workerServiceSupplier = workerServiceSupplier; @@ -198,4 +205,27 @@ public List 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()); + } + } } diff --git a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/v2/WorkerApiV2Resource.java b/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/v2/WorkerApiV2Resource.java index b5217f2cad165..e233e80334a71 100644 --- a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/rest/api/v2/WorkerApiV2Resource.java +++ b/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 { protected ServletContext servletContext; @Context protected HttpServletRequest httpRequest; + @Context + protected UriInfo uri; public WorkerApiV2Resource() { this.worker = new WorkerImpl(this); @@ -138,4 +142,18 @@ public Map> getAssignments() { public List 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()); + } } diff --git a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/IScheduler.java b/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/IScheduler.java index e19ca71a44842..7c703259398db 100644 --- a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/IScheduler.java +++ b/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 schedule(List unassignedFunctionInstances, List currentAssignments, Set 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 rebalance(List currentAssignments, Set workers){ + return Collections.emptyList(); + } } diff --git a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinScheduler.java b/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinScheduler.java index 70e1d90b18da9..59b5137ed2ec4 100644 --- a/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinScheduler.java +++ b/pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinScheduler.java @@ -18,17 +18,25 @@ */ package org.apache.pulsar.functions.worker.scheduler; +import com.fasterxml.jackson.core.JsonProcessingException; +import lombok.Builder; +import lombok.Data; +import lombok.extern.slf4j.Slf4j; +import org.apache.pulsar.common.util.ObjectMapperFactory; import org.apache.pulsar.functions.proto.Function.Assignment; import org.apache.pulsar.functions.proto.Function.Instance; import com.google.common.collect.Lists; +import java.util.Comparator; import java.util.HashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Queue; import java.util.Set; +@Slf4j public class RoundRobinScheduler implements IScheduler { @Override @@ -70,4 +78,116 @@ private String findNextWorker(Map> workerIdToAssignment } return targetWorkerId; } + + @Override + public List rebalance(List currentAssignments, Set workers) { + + Map> workerToAssignmentMap = new HashMap<>(); + + workers.forEach(workerId -> workerToAssignmentMap.put(workerId, new LinkedList<>())); + + currentAssignments.forEach(assignment -> workerToAssignmentMap.computeIfAbsent(assignment.getWorkerId(), s -> new LinkedList<>()).add(assignment.getInstance())); + + List newAssignments = new LinkedList<>(); + + RebalanceStats rebalanceStats = new RebalanceStats(workerToAssignmentMap); + int iterations = 0; + while(true) { + iterations++; + + Map.Entry> mostAssignmentsWorker = findWorkerWithMostAssignments(workerToAssignmentMap); + + Map.Entry> leastAssignmentsWorker = findWorkerWithLeastAssignments(workerToAssignmentMap); + + if (mostAssignmentsWorker.getValue().size() == leastAssignmentsWorker.getValue().size() + || mostAssignmentsWorker.getValue().size() == leastAssignmentsWorker.getValue().size() + 1) { + break; + } + + String mostAssignmentsWorkerId = mostAssignmentsWorker.getKey(); + String leastAssignmentsWorkerId = leastAssignmentsWorker.getKey(); + + Queue src = workerToAssignmentMap.get(mostAssignmentsWorkerId); + Queue dest = workerToAssignmentMap.get(leastAssignmentsWorkerId); + + // update stats + rebalanceStats.decrementInstance(mostAssignmentsWorkerId); + rebalanceStats.incrementInstance(leastAssignmentsWorkerId); + + Instance instance = src.poll(); + Assignment newAssignment = Assignment.newBuilder() + .setInstance(instance) + .setWorkerId(leastAssignmentsWorkerId) + .build(); + newAssignments.add(newAssignment); + + dest.add(instance); + } + + log.info("Rebalance - iterations: {} stats: {}", iterations, rebalanceStats); + + return newAssignments; + } + + private Map.Entry> findWorkerWithLeastAssignments(Map> workerToAssignmentMap) { + return workerToAssignmentMap.entrySet().stream().min(Comparator.comparingInt(o -> o.getValue().size())).get(); + + } + + private Map.Entry> findWorkerWithMostAssignments(Map> workerToAssignmentMap) { + return workerToAssignmentMap.entrySet().stream().max(Comparator.comparingInt(o -> o.getValue().size())).get(); + } + + @Data + private static class RebalanceStats { + @Override + public String toString() { + try { + return ObjectMapperFactory.getThreadLocal().writerWithDefaultPrettyPrinter().writeValueAsString(workerStatsMap); + } catch (JsonProcessingException e) { + throw new RuntimeException(e); + } + } + + @Builder + @Data + private static class WorkerStats { + private int originalNumAssignments; + private int numAssignmentsAfterRebalance; + private int instancesAdded; + private int instancesRemoved; + } + + private Map workerStatsMap = new HashMap<>(); + + public RebalanceStats(Map> workerToAssignmentMap) { + for(Map.Entry> entry : workerToAssignmentMap.entrySet()) { + WorkerStats workerStats = WorkerStats.builder() + .originalNumAssignments(entry.getValue().size()) + .numAssignmentsAfterRebalance(entry.getValue().size()) + .build(); + workerStatsMap.put(entry.getKey(), workerStats); + } + } + + private void decrementInstance(String workerId) { + WorkerStats stats = workerStatsMap.get(workerId); + if (stats == null) { + throw new RuntimeException("Rebalance stats for worker " + workerId + " shouldn't be null"); + } + + stats.instancesRemoved++; + stats.numAssignmentsAfterRebalance--; + } + + private void incrementInstance(String workerId) { + WorkerStats stats = workerStatsMap.get(workerId); + if (stats == null) { + throw new RuntimeException("Rebalance stats for worker " + workerId + " shouldn't be null"); + } + + stats.instancesAdded++; + stats.numAssignmentsAfterRebalance++; + } + } } diff --git a/pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinSchedulerTest.java b/pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinSchedulerTest.java new file mode 100644 index 0000000000000..497587174debc --- /dev/null +++ b/pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/scheduler/RoundRobinSchedulerTest.java @@ -0,0 +1,78 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.functions.worker.scheduler; + +import lombok.extern.slf4j.Slf4j; +import org.apache.pulsar.functions.proto.Function; +import org.testng.Assert; +import org.testng.annotations.Test; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Set; + +@Slf4j +public class RoundRobinSchedulerTest { + + @Test + public void testRebalance() { + Function.FunctionMetaData function1 = Function.FunctionMetaData.newBuilder() + .setFunctionDetails(Function.FunctionDetails.newBuilder().setName("func-1") + .setNamespace("namespace-1").setTenant("tenant-1").setParallelism(1)).setVersion(0) + .build(); + + List assignments = new LinkedList<>(); + for (int i = 0; i < 10; i++) { + Function.Assignment assignment1 = Function.Assignment.newBuilder() + .setWorkerId("worker-1") + .setInstance(Function.Instance.newBuilder() + .setFunctionMetaData(function1).setInstanceId(i).build()) + .build(); + + assignments.add(assignment1); + } + + Set workers = new HashSet<>(); + for (int i = 0; i < 3; i++) { + workers.add("worker-" + i); + } + + RoundRobinScheduler roundRobinScheduler = new RoundRobinScheduler(); + + List newAssignments = roundRobinScheduler.rebalance(assignments, workers); + + Map workerAssignments = new HashMap<>(); + for (Function.Assignment assignment : newAssignments) { + Integer count = workerAssignments.get((assignment.getWorkerId())); + if (count == null) { + count = 0; + } + count++; + workerAssignments.put(assignment.getWorkerId(), count); + } + + Assert.assertEquals(workerAssignments.size(), 2); + for (Map.Entry entry : workerAssignments.entrySet()) { + Assert.assertEquals(entry.getValue().intValue(), 3); + } + } +}