KAFKA-5692: Change PreferredReplicaLeaderElectionCommand to use Admin…

checkstyle/import-control.xml

@@ -116,6 +116,7 @@
 
     <subpackage name="protocol">
       <allow pkg="org.apache.kafka.common.errors" />
+      <allow pkg="org.apache.kafka.common.message" />
       <allow pkg="org.apache.kafka.common.protocol.types" />
       <allow pkg="org.apache.kafka.common.record" />
       <allow pkg="org.apache.kafka.common.requests" />
@@ -140,6 +141,7 @@
     <subpackage name="requests">
       <allow pkg="org.apache.kafka.common.acl" />
       <allow pkg="org.apache.kafka.common.protocol" />
+      <allow pkg="org.apache.kafka.common.message" />
       <allow pkg="org.apache.kafka.common.network" />
       <allow pkg="org.apache.kafka.common.requests" />
       <allow pkg="org.apache.kafka.common.resource" />

clients/src/main/java/org/apache/kafka/clients/NetworkClient.java

@@ -676,7 +676,8 @@ public Node leastLoadedNode(long now) {
 
     public static AbstractResponse parseResponse(ByteBuffer responseBuffer, RequestHeader requestHeader) {
         Struct responseStruct = parseStructMaybeUpdateThrottleTimeMetrics(responseBuffer, requestHeader, null, 0);
-        return AbstractResponse.parseResponse(requestHeader.apiKey(), responseStruct);
+        return AbstractResponse.parseResponse(requestHeader.apiKey(), responseStruct,
+                requestHeader.apiVersion());
     }
 
     private static Struct parseStructMaybeUpdateThrottleTimeMetrics(ByteBuffer responseBuffer, RequestHeader requestHeader,
@@ -811,7 +812,8 @@ private void handleCompletedReceives(List<ClientResponse> responses, long now) {
                     req.header.apiKey(), req.header.correlationId(), responseStruct);
             }
             // If the received response includes a throttle delay, throttle the connection.
-            AbstractResponse body = AbstractResponse.parseResponse(req.header.apiKey(), responseStruct);
+            AbstractResponse body = AbstractResponse.
+                    parseResponse(req.header.apiKey(), responseStruct, req.header.apiVersion());
             maybeThrottle(body, req.header.apiVersion(), req.destination, now);
             if (req.isInternalRequest && body instanceof MetadataResponse)
                 metadataUpdater.handleCompletedMetadataResponse(req.header, now, (MetadataResponse) body);

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

@@ -792,6 +792,58 @@ public DeleteConsumerGroupsResult deleteConsumerGroups(Collection<String> groupI
         return deleteConsumerGroups(groupIds, new DeleteConsumerGroupsOptions());
     }
 
+    /**
+     * Elect the preferred broker of the given {@code partitions} as leader, or
+     * elect the preferred broker 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 preferred leader should be elected.
+     * @return                The ElectPreferredLeadersResult.
+     */
+    public ElectPreferredLeadersResult electPreferredLeaders(Collection<TopicPartition> partitions) {
+        return electPreferredLeaders(partitions, new ElectPreferredLeadersOptions());
+    }
+
+    /**
+     * Elect the preferred broker of the given {@code partitions} as leader, or
+     * elect the preferred broker 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 in the cluster to become aware that the partitions have new leaders.
+     * During this time, {@link AdminClient#describeTopics(Collection)}
+     * may not return information about the partitions' new leaders.
+     *
+     * This operation is supported by brokers with version 2.2.0 or higher.
+     *
+     * <p>The following exceptions can be anticipated when calling {@code get()} on the futures obtained from
+     * the returned {@code ElectPreferredLeadersResult}:</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 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);
+
     /**
      * Get the metrics kept by the adminClient
      */

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,136 @@
+/*
+ * 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.protocol.Errors;
+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;
+    private final Set<TopicPartition> partitions;
+
+    ElectPreferredLeadersResult(KafkaFutureImpl<Map<TopicPartition, ApiError>> electionFuture, Set<TopicPartition> partitions) {
+        this.electionFuture = electionFuture;
+        this.partitions = partitions;
+    }
+
+    /**
+     * 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.whenComplete(new KafkaFuture.BiConsumer<Map<TopicPartition, ApiError>, Throwable>() {
+            @Override
+            public void accept(Map<TopicPartition, ApiError> topicPartitions, Throwable throwable) {
+                if (throwable != null) {
+                    result.completeExceptionally(throwable);
+                } else if (!topicPartitions.containsKey(partition)) {
+                    result.completeExceptionally(new UnknownTopicOrPartitionException(
+                            "Preferred leader election for partition \"" + partition +
+                                    "\" was not attempted"));
+                } else {
+                    if (partitions == null && topicPartitions.isEmpty()) {
+                        result.completeExceptionally(Errors.CLUSTER_AUTHORIZATION_FAILED.exception());
+                    }
+                    ApiException exception = topicPartitions.get(partition).exception();
+                    if (exception == null) {
+                        result.complete(null);
+                    } else {
+                        result.completeExceptionally(exception);
+                    }
+                }
+            }
+        });
+        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() {
+        if (partitions != null) {
+            return KafkaFutureImpl.completedFuture(this.partitions);
+        } else {
+            final KafkaFutureImpl<Set<TopicPartition>> result = new KafkaFutureImpl<>();
+            electionFuture.whenComplete(new KafkaFuture.BiConsumer<Map<TopicPartition, ApiError>, Throwable>() {
+                @Override
+                public void accept(Map<TopicPartition, ApiError> topicPartitions, Throwable throwable) {
+                    if (throwable != null) {
+                        result.completeExceptionally(throwable);
+                    } else if (topicPartitions.isEmpty()) {
+                        result.completeExceptionally(Errors.CLUSTER_AUTHORIZATION_FAILED.exception());
+                    } else {
+                        for (ApiError apiError : topicPartitions.values()) {
+                            if (apiError.isFailure()) {
+                                result.completeExceptionally(apiError.exception());
+                            }
+                        }
+                        result.complete(topicPartitions.keySet());
+                    }
+                }
+            });
+            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

@@ -105,6 +105,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.ExpireDelegationTokenRequest;
 import org.apache.kafka.common.requests.ExpireDelegationTokenResponse;
 import org.apache.kafka.common.requests.FindCoordinatorRequest;
@@ -2777,4 +2779,35 @@ void handleFailure(Throwable throwable) {
     public Map<MetricName, ? extends Metric> metrics() {
         return Collections.unmodifiableMap(this.metrics.metrics());
     }
+
+    @Override
+    public ElectPreferredLeadersResult electPreferredLeaders(final Collection<TopicPartition> partitions,
+                                                             ElectPreferredLeadersOptions options) {
+        final Set<TopicPartition> partitionSet = partitions != null ? new HashSet<>(partitions) : null;
+        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(
+                        ElectPreferredLeadersRequest.toRequestData(partitions, timeoutMs));
+            }
+
+            @Override
+            public void handleResponse(AbstractResponse abstractResponse) {
+                ElectPreferredLeadersResponse response = (ElectPreferredLeadersResponse) abstractResponse;
+                electionFuture.complete(
+                        ElectPreferredLeadersRequest.fromResponseData(response.data()));
+            }
+
+            @Override
+            void handleFailure(Throwable throwable) {
+                electionFuture.completeExceptionally(throwable);
+            }
+        }, now);
+        return new ElectPreferredLeadersResult(electionFuture, partitionSet);
+    }
+
 }

clients/src/main/java/org/apache/kafka/common/errors/PreferredLeaderNotAvailableException.java

@@ -0,0 +1,28 @@
+/*
+ * 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.common.errors;
+
+public class PreferredLeaderNotAvailableException extends InvalidMetadataException {
+
+    public PreferredLeaderNotAvailableException(String message) {
+        super(message);
+    }
+
+    public PreferredLeaderNotAvailableException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

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

@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.protocol;
 
+import org.apache.kafka.common.message.ElectPreferredLeadersRequestData;
+import org.apache.kafka.common.message.ElectPreferredLeadersResponseData;
 import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.SchemaException;
 import org.apache.kafka.common.protocol.types.Struct;
@@ -35,10 +37,10 @@
 import org.apache.kafka.common.requests.ControlledShutdownResponse;
 import org.apache.kafka.common.requests.CreateAclsRequest;
 import org.apache.kafka.common.requests.CreateAclsResponse;
-import org.apache.kafka.common.requests.CreatePartitionsRequest;
-import org.apache.kafka.common.requests.CreatePartitionsResponse;
 import org.apache.kafka.common.requests.CreateDelegationTokenRequest;
 import org.apache.kafka.common.requests.CreateDelegationTokenResponse;
+import org.apache.kafka.common.requests.CreatePartitionsRequest;
+import org.apache.kafka.common.requests.CreatePartitionsResponse;
 import org.apache.kafka.common.requests.CreateTopicsRequest;
 import org.apache.kafka.common.requests.CreateTopicsResponse;
 import org.apache.kafka.common.requests.DeleteAclsRequest;
@@ -53,12 +55,12 @@
 import org.apache.kafka.common.requests.DescribeAclsResponse;
 import org.apache.kafka.common.requests.DescribeConfigsRequest;
 import org.apache.kafka.common.requests.DescribeConfigsResponse;
+import org.apache.kafka.common.requests.DescribeDelegationTokenRequest;
+import org.apache.kafka.common.requests.DescribeDelegationTokenResponse;
 import org.apache.kafka.common.requests.DescribeGroupsRequest;
 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.DescribeDelegationTokenRequest;
-import org.apache.kafka.common.requests.DescribeDelegationTokenResponse;
 import org.apache.kafka.common.requests.EndTxnRequest;
 import org.apache.kafka.common.requests.EndTxnResponse;
 import org.apache.kafka.common.requests.ExpireDelegationTokenRequest;
@@ -186,7 +188,9 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
     RENEW_DELEGATION_TOKEN(39, "RenewDelegationToken", RenewDelegationTokenRequest.schemaVersions(), RenewDelegationTokenResponse.schemaVersions()),
     EXPIRE_DELEGATION_TOKEN(40, "ExpireDelegationToken", ExpireDelegationTokenRequest.schemaVersions(), ExpireDelegationTokenResponse.schemaVersions()),
     DESCRIBE_DELEGATION_TOKEN(41, "DescribeDelegationToken", DescribeDelegationTokenRequest.schemaVersions(), DescribeDelegationTokenResponse.schemaVersions()),
-    DELETE_GROUPS(42, "DeleteGroups", DeleteGroupsRequest.schemaVersions(), DeleteGroupsResponse.schemaVersions());
+    DELETE_GROUPS(42, "DeleteGroups", DeleteGroupsRequest.schemaVersions(), DeleteGroupsResponse.schemaVersions()),
+    ELECT_PREFERRED_LEADERS(43, "ElectPreferredLeaders", ElectPreferredLeadersRequestData.SCHEMAS,
+            ElectPreferredLeadersResponseData.SCHEMAS);
 
     private static final ApiKeys[] ID_TO_TYPE;
     private static final int MIN_API_KEY = 0;