KAFKA-7026: Sticky Assignor Partition Assignment Improvement (KIP-341)

clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java

@@ -566,7 +566,7 @@ public void handle(JoinGroupResponse joinResponse, RequestFuture<ByteBuffer> fut
                 // and send another join group request in next cycle.
                 synchronized (AbstractCoordinator.this) {
                     AbstractCoordinator.this.generation = new Generation(OffsetCommitRequest.DEFAULT_GENERATION_ID,
-                        joinResponse.data().memberId(), null);
+                            joinResponse.data().memberId(), null);
                     AbstractCoordinator.this.rejoinNeeded = true;
                     AbstractCoordinator.this.state = MemberState.UNJOINED;
                 }
@@ -654,7 +654,7 @@ private RequestFuture<Void> sendFindCoordinatorRequest(Node node) {
         FindCoordinatorRequest.Builder requestBuilder =
                 new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, this.groupId);
         return client.send(node, requestBuilder)
-                     .compose(new FindCoordinatorResponseHandler());
+                .compose(new FindCoordinatorResponseHandler());
     }
 
     private class FindCoordinatorResponseHandler extends RequestFutureAdapter<ClientResponse, Void> {
@@ -940,8 +940,8 @@ public GroupCoordinatorMetrics(Metrics metrics, String metricGrpPrefix) {
 
             this.heartbeatLatency = metrics.sensor("heartbeat-latency");
             this.heartbeatLatency.add(metrics.metricName("heartbeat-response-time-max",
-                this.metricGrpName,
-                "The max time taken to receive a response to a heartbeat request"), new Max());
+                    this.metricGrpName,
+                    "The max time taken to receive a response to a heartbeat request"), new Max());
             this.heartbeatLatency.add(createMeter(metrics, metricGrpName, "heartbeat", "heartbeats"));
 
             this.joinLatency = metrics.sensor("join-latency");
@@ -1148,8 +1148,8 @@ public boolean equals(final Object o) {
             if (o == null || getClass() != o.getClass()) return false;
             final Generation that = (Generation) o;
             return generationId == that.generationId &&
-                Objects.equals(memberId, that.memberId) &&
-                Objects.equals(protocol, that.protocol);
+                    Objects.equals(memberId, that.memberId) &&
+                    Objects.equals(protocol, that.protocol);
         }
 
         @Override

clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java

@@ -262,7 +262,7 @@ protected void onJoinComplete(int generation,
         maybeUpdateJoinedSubscription(assignedPartitions);
 
         // give the assignor a chance to update internal state based on the received assignment
-        assignor.onAssignment(assignment);
+        assignor.onAssignment(assignment, generation);
 
         // reschedule the auto commit starting from now
         if (autoCommitEnabled)

clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java

@@ -45,6 +45,7 @@
  *   TopicPartitions => [Topic Partitions]
  *     Topic         => String
  *     Partitions    => [int32]
+ *   UserData        => Bytes
  * </pre>
  *
  * The current implementation assumes that future versions will not break compatibility. When

clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java

@@ -58,13 +58,21 @@ public interface PartitionAssignor {
      */
     Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions);
 
-
     /**
      * Callback which is invoked when a group member receives its assignment from the leader.
      * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)}
      */
     void onAssignment(Assignment assignment);
 
+    /**
+     * Callback which is invoked when a group member receives its assignment from the leader.
+     * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)}
+     * @param generation The consumer group generation associated with this partition assignment (optional)
+     */
+    default void onAssignment(Assignment assignment, int generation) {
+        onAssignment(assignment);
+    }
+
 
     /**
      * Unique name for this assignor (e.g. "range" or "roundrobin" or "sticky")