Kafka 5692 elect preferred

clients/src/main/java/org/apache/kafka/clients/admin/AdminClient.java

@@ -17,6 +17,7 @@
 
 package org.apache.kafka.clients.admin;
 
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.TopicPartitionReplica;
 import org.apache.kafka.common.acl.AclBinding;
 import org.apache.kafka.common.acl.AclBindingFilter;
@@ -508,4 +509,55 @@ public CreatePartitionsResult createPartitions(Map<String, NewPartitions> newPar
     public abstract CreatePartitionsResult createPartitions(Map<String, NewPartitions> newPartitions,
                                                             CreatePartitionsOptions options);
 
+    /*
+     * Elect the preferred replica of the given {@code partitions} as leader, or
+     * elect the preferred replica for all partitions as leader if the argument to {@code partitions} is null.
+     *
+     * This is a convenience method for {@link #electPreferredLeaders(Collection, ElectPreferredLeadersOptions)} with default options.
+     * See the overload for more details.
+     *
+     * @param partitions      The partitions for which the the preferred leader should be elected.
+     * @return                The ElectPreferredLeadersResult.
+     */
+    public ElectPreferredLeadersResult electPreferredLeaders(Collection<TopicPartition> partitions) {
+        return electPreferredLeaders(partitions, new ElectPreferredLeadersOptions());
+    }
+
+    /**
+     * Elect the preferred replica of the given {@code partitions} as leader, or
+     * elect the preferred replica for all partitions as leader if the argument to {@code partitions} is null.
+     *
+     * This operation is not transactional so it may succeed for some partitions while fail for others.
+     *
+     * It may take several seconds after this method returns
+     * success for all the brokers to become aware that the partitions have new leaders.
+     * During this time, {@link AdminClient#describeTopics(Collection)}
+     * may not return information about the new partitions leaders.
+     *
+     * This operation is supported by brokers with version 1.1.0 or higher.
+     *
+     * <p>The {@code KafkaFuture}s available from instances of this class may be completed
+     * exceptionally due to:</p>
+     * <ul>
+     *   <li>{@link org.apache.kafka.common.errors.ClusterAuthorizationException}
+     *   if the authenticated user didn't have alter access to the cluster.</li>
+     *   <li>{@link org.apache.kafka.common.errors.UnknownTopicOrPartitionException}
+     *   if the topic or partition did not exist within the cluster.</li>
+     *   <li>{@link org.apache.kafka.common.errors.InvalidTopicException}
+     *   if the topic was already queued for deletion.</li>
+     *   <li>{@link org.apache.kafka.common.errors.NotControllerException}
+     *   if the request was sent to a broker that was not the controller for the cluster.</li>
+     *   <li>{@link org.apache.kafka.common.errors.TimeoutException}
+     *   if the request timed out before the election was complete.</li>
+     *   <li>{@link org.apache.kafka.common.errors.LeaderNotAvailableException}
+     *   if the preferred leader was not alive or not in the ISR.</li>
+     * </ul>
+     *
+     * @param partitions      The partitions for which the the preferred leader should be elected.
+     * @param options         The options to use when electing the preferred leaders.
+     * @return                The ElectPreferredLeadersResult.
+     */
+    public abstract ElectPreferredLeadersResult electPreferredLeaders(Collection<TopicPartition> partitions,
+                                                                      ElectPreferredLeadersOptions options);
+
 }

clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersOptions.java

@@ -0,0 +1,31 @@
+/*
+ * 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.kafka.clients.admin;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.util.Collection;
+
+/**
+ * Options for {@link AdminClient#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)}.
+ *
+ * The API of this class is evolving, see {@link AdminClient} for details.
+ */
+@InterfaceStability.Evolving
+public class ElectPreferredLeadersOptions extends AbstractOptions<ElectPreferredLeadersOptions> {
+}

clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersResult.java

@@ -0,0 +1,121 @@
+/*
+ * 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.kafka.clients.admin;
+
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.errors.ApiException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+import org.apache.kafka.common.requests.ApiError;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * The result of {@link AdminClient#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)}
+ *
+ * The API of this class is evolving, see {@link AdminClient} for details.
+ */
+@InterfaceStability.Evolving
+public class ElectPreferredLeadersResult {
+
+    private final KafkaFutureImpl<Map<TopicPartition, ApiError>> electionFuture;
+
+    ElectPreferredLeadersResult(KafkaFutureImpl<Map<TopicPartition, ApiError>> electionFuture) {
+        this.electionFuture = electionFuture;
+    }
+
+    /**
+     * Get the result of the election for the given {@code partition}.
+     * If there was not an election triggered for the given {@code partition}, the
+     * returned future will complete with an error.
+     */
+    public KafkaFuture<Void> partitionResult(final TopicPartition partition) {
+        final KafkaFutureImpl<Void> result = new KafkaFutureImpl<>();
+        electionFuture.thenApply(new KafkaFuture.Function<Map<TopicPartition, ApiError>, Void>() {
+            @Override
+            public Void apply(Map<TopicPartition, ApiError> topicPartitions) {
+                if (!topicPartitions.containsKey(partition)) {
+                    result.completeExceptionally(new UnknownTopicOrPartitionException(
+                        "Preferred leader election for partition \"" + partition +
+                            "\" was not attempted"));
+                } else {
+                    ApiException exception = topicPartitions.get(partition).exception();
+                    if (exception == null) {
+                        result.complete(null);
+                    } else {
+                        result.completeExceptionally(exception);
+                    }
+                }
+                return null;
+            }
+        });
+        return result;
+    }
+
+    /**
+     * <p>Get a future for the topic partitions for which a leader election
+     * was attempted. A partition will be present in this result if
+     * an election was attempted even if the election was not successful.</p>
+     *
+     * <p>This method is provided to discover the partitions attempted when
+     * {@link AdminClient#electPreferredLeaders(Collection)} is called
+     * with a null {@code partitions} argument.</p>
+     */
+    public KafkaFuture<Set<TopicPartition>> partitions() {
+        final KafkaFutureImpl<Set<TopicPartition>> result = new KafkaFutureImpl<>();
+        electionFuture.thenApply(new KafkaFuture.Function<Map<TopicPartition, ApiError>, Void>() {
+            @Override
+            public Void apply(Map<TopicPartition, ApiError> topicPartitions) {
+                for (ApiError apiError : topicPartitions.values()) {
+                    if (apiError.isFailure()) {
+                        result.completeExceptionally(apiError.exception());
+                        return null;
+                    }
+                }
+                result.complete(topicPartitions.keySet());
+                return null;
+            }
+        });
+        return result;
+    }
+
+    /**
+     * Return a future which succeeds if all the topic elections succeed.
+     */
+    public KafkaFuture<Void> all() {
+        final KafkaFutureImpl<Void> result = new KafkaFutureImpl<>();
+        electionFuture.thenApply(new KafkaFuture.Function<Map<TopicPartition, ApiError>, Void>() {
+            @Override
+            public Void apply(Map<TopicPartition, ApiError> topicPartitions) {
+                for (ApiError apiError : topicPartitions.values()) {
+                    if (apiError.isFailure()) {
+                        result.completeExceptionally(apiError.exception());
+                        return null;
+                    }
+                }
+                result.complete(null);
+                return null;
+            }
+        });
+        return result;
+    }
+}

clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java

@@ -86,6 +86,8 @@
 import org.apache.kafka.common.requests.DescribeConfigsResponse;
 import org.apache.kafka.common.requests.DescribeLogDirsRequest;
 import org.apache.kafka.common.requests.DescribeLogDirsResponse;
+import org.apache.kafka.common.requests.ElectPreferredLeadersRequest;
+import org.apache.kafka.common.requests.ElectPreferredLeadersResponse;
 import org.apache.kafka.common.requests.MetadataRequest;
 import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.common.requests.Resource;
@@ -1873,4 +1875,30 @@ void handleFailure(Throwable throwable) {
         return new CreatePartitionsResult(new HashMap<String, KafkaFuture<Void>>(futures));
     }
 
+    @Override
+    public ElectPreferredLeadersResult electPreferredLeaders(final Collection<TopicPartition> partitions,
+                                                             ElectPreferredLeadersOptions options) {
+        final KafkaFutureImpl<Map<TopicPartition, ApiError>> electionFuture = new KafkaFutureImpl<>();
+        final long now = time.milliseconds();
+        runnable.call(new Call("electPreferredLeaders", calcDeadlineMs(now, options.timeoutMs()),
+                new ControllerNodeProvider()) {
+
+            @Override
+            public AbstractRequest.Builder createRequest(int timeoutMs) {
+                return new ElectPreferredLeadersRequest.Builder(partitions, timeoutMs);
+            }
+
+            @Override
+            public void handleResponse(AbstractResponse abstractResponse) {
+                ElectPreferredLeadersResponse response = (ElectPreferredLeadersResponse) abstractResponse;
+                electionFuture.complete(response.errors());
+            }
+
+            @Override
+            void handleFailure(Throwable throwable) {
+                electionFuture.completeExceptionally(throwable);
+            }
+        }, now);
+        return new ElectPreferredLeadersResult(electionFuture);
+    }
 }

clients/src/main/java/org/apache/kafka/common/KafkaFuture.java

@@ -84,6 +84,15 @@ public static <U> KafkaFuture<U> completedFuture(U value) {
         return future;
     }
 
+    /**
+     * Return a new future that has completed exceptionally with the given exception.
+     */
+    public static <T> KafkaFuture<T> exceptionalFuture(Throwable exception) {
+        KafkaFutureImpl<T> future = new KafkaFutureImpl<>();
+        future.completeExceptionally(exception);
+        return future;
+    }
+
     /** 
      * Returns a new KafkaFuture that is completed when all the given futures have completed.  If
      * any future throws an exception, the returned future returns it.  If multiple futures throw
@@ -104,6 +113,8 @@ public static KafkaFuture<Void> allOf(KafkaFuture<?>... futures) {
      */
     public abstract <R> KafkaFuture<R> thenApply(Function<T, R> function);
 
+    public abstract <R> KafkaFuture<R> thenCompose(Function<T, KafkaFuture<R>> function);
+
     protected abstract void addWaiter(BiConsumer<? super T, ? super Throwable> action);
 
     /**

clients/src/main/java/org/apache/kafka/common/internals/KafkaFutureImpl.java

@@ -70,6 +70,41 @@ public void accept(A a, Throwable exception) {
         }
     }
 
+    private static class Composer<A, B> extends BiConsumer<A, Throwable> {
+        private final Function<A, KafkaFuture<B>> function;
+        private final KafkaFutureImpl<B> future;
+
+        Composer(Function<A, KafkaFuture<B>> function, KafkaFutureImpl<B> future) {
+            this.function = function;
+            this.future = future;
+        }
+
+        // Called when "this" completes
+        @Override
+        public void accept(A a, Throwable exception) {
+            if (exception != null) {
+                future.completeExceptionally(exception);
+            } else {
+                try {
+                    // TODO typecast eek!
+                    KafkaFutureImpl<B> mapped = (KafkaFutureImpl) function.apply(a);
+                    mapped.addWaiter(new BiConsumer<B, Throwable>() {
+                        @Override
+                        public void accept(B b, Throwable throwable) {
+                            if (throwable != null) {
+                                future.completeExceptionally(throwable);
+                            } else {
+                                future.complete(b);
+                            }
+                        }
+                    });
+                } catch (Throwable t) {
+                    future.completeExceptionally(t);
+                }
+            }
+        }
+    }
+
     private static class SingleWaiter<R> extends BiConsumer<R, Throwable> {
         private R value = null;
         private Throwable exception = null;
@@ -146,6 +181,13 @@ public <R> KafkaFuture<R> thenApply(Function<T, R> function) {
         return future;
     }
 
+    @Override
+    public <R> KafkaFuture<R> thenCompose(Function<T, KafkaFuture<R>> function) {
+        KafkaFutureImpl<R> future = new KafkaFutureImpl<R>();
+        addWaiter(new Composer<>(function, future));
+        return future;
+    }
+
     @Override
     protected synchronized void addWaiter(BiConsumer<? super T, ? super Throwable> action) {
         if (exception != null) {

clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java

@@ -53,6 +53,8 @@
 import org.apache.kafka.common.requests.DescribeGroupsResponse;
 import org.apache.kafka.common.requests.DescribeLogDirsRequest;
 import org.apache.kafka.common.requests.DescribeLogDirsResponse;
+import org.apache.kafka.common.requests.ElectPreferredLeadersRequest;
+import org.apache.kafka.common.requests.ElectPreferredLeadersResponse;
 import org.apache.kafka.common.requests.EndTxnRequest;
 import org.apache.kafka.common.requests.EndTxnResponse;
 import org.apache.kafka.common.requests.FetchRequest;
@@ -171,7 +173,9 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
     SASL_AUTHENTICATE(36, "SaslAuthenticate", SaslAuthenticateRequest.schemaVersions(),
             SaslAuthenticateResponse.schemaVersions()),
     CREATE_PARTITIONS(37, "CreatePartitions", CreatePartitionsRequest.schemaVersions(),
-            CreatePartitionsResponse.schemaVersions());
+            CreatePartitionsResponse.schemaVersions()),
+    ELECT_PREFERRED_LEADERS(38, "ElectPreferredLeaders", ElectPreferredLeadersRequest.schemaVersions(),
+            ElectPreferredLeadersResponse.schemaVersions());
 
     private static final ApiKeys[] ID_TO_TYPE;
     private static final int MIN_API_KEY = 0;

clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java

@@ -214,6 +214,8 @@ public static AbstractRequest parseRequest(ApiKeys apiKey, short apiVersion, Str
                 return new SaslAuthenticateRequest(struct, apiVersion);
             case CREATE_PARTITIONS:
                 return new CreatePartitionsRequest(struct, apiVersion);
+            case ELECT_PREFERRED_LEADERS:
+                return new ElectPreferredLeadersRequest(struct, apiVersion);
             default:
                 throw new AssertionError(String.format("ApiKey %s is not currently handled in `parseRequest`, the " +
                         "code should be updated to do so.", apiKey));

clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java

@@ -146,6 +146,8 @@ public static AbstractResponse parseResponse(ApiKeys apiKey, Struct struct) {
                 return new SaslAuthenticateResponse(struct);
             case CREATE_PARTITIONS:
                 return new CreatePartitionsResponse(struct);
+            case ELECT_PREFERRED_LEADERS:
+                return new ElectPreferredLeadersResponse(struct);
             default:
                 throw new AssertionError(String.format("ApiKey %s is not currently handled in `parseResponse`, the " +
                         "code should be updated to do so.", apiKey));