Status

Current state: Draft Under Discussion

Discussion thread: here 

JIRA: KAFKA-15045

...

Back in KIP-441, we introduced a new task assignor which prioritized stateful task availability over stickiness, but kept the old assignment logic which prioritized stickiness called the StickyTaskAssignor. As a safeguard we added a backdoor internal config for the task assignor, the idea being that we could recommend this to users in case of critical bugs in the new assignor. However over time it has become clear that there are valid use cases for wanting the old StickyTaskAssignor over the new HighAvailabilityTaskAssignor, primarily relating to the possibility of extreme task shuffling induced by the HAAssignor in less stable environments. As the overall assignment logic is becoming more and more sophisticated, now is a good time to finally move the internal task assignor config to a full-fledged public API.

Public Interfaces

The original change here will be introducing a new public config for the task assignor class. Note that the thread-level assignment will remain an internal implementation detail (see "Rejected Alternatives" for further thoughts and reasoning)

There are also reasonable scenarios where users may want to plug in their own task assignor. For example, a user may want to configure their own static assignment to work around an issue in the available assignors. Or they may want to implement their own assignment logic that considers metrics they've collected, or task migration cost related to their specific processors.

Finally, there are good reasons for a user to want to extend or modify the behaviour of the Kafka Streams partition assignor beyond just changing the task assignment. For example, a user may want to implement their own initialization logic that initializes resources (much the same way the Streams Partition Assignor initializes internal topics).

With these motivations in mind, we are proposing to add a new group of configurable interfaces for plugging custom behaviour into the Streams Partition Assignor. This configuration will supplant the existing internal task assignor config. In this KIP, we will limit the scope of these configs to supplying a custom task assignor. However, this gives us a pattern based on which to, in future KIPs, add further configs which a user can set to plug in custom behavior.

Public Interfaces

We will introduce a new config that supplies an instance of TaskAssignor  (discussed below). In the future, additional plugins can use the same partition.assignor  prefix:

/** {@code partition.assignor.task.assignor.class} */
public

/** {@code task.assignor.class} */
public static final String TASK_ASSIGNOR_CLASS_CONFIG = "task.assignor.class";
private static final String PARTITION_ASSIGNOR_TASK_ASSIGNOR_CLASS_DOCCONFIG = "A task assignor class or class name implementing the <@link TaskAssignor> interface"
        + " that will be used to assign tasks to client nodes during a rebalance. Note that tasks will be distributed equally among the"
        + "StreamThreads within each client with maximal stickiness ie returning as many tasks as possible to their previous owner on that"
        + "client node. Defaults to the <@link HighAvailabilityTaskAssignor> class.";

We will also remove the old internal config (which we can do without deprecation as this was an internal config and thus by definition not part of the public API, also discussed in more detail in "Rejected Alternatives")

public static class InternalConfig {
        // This will be removed
        public static final String INTERNAL_TASK_ASSIGNOR_CLASS = "internal.task.assignor.class";
}

...

partition.assignor.task.assignor.class";
private static final String PARTITION_ASSIGNOR_TASK_ASSIGNOR_CLASS_DOC = "A task assignor class or class name implementing the <@link TaskAssignor> interface". Defaults to the <@link HighAvailabilityTaskAssignor> class.";

We will also remove the old internal config (which we can do without deprecation as this was an internal config and thus by definition not part of the public API, also discussed in more detail in "Rejected Alternatives")

public static class InternalConfig {
        // This will be removed
        public static final String INTERNAL_TASK_ASSIGNOR_CLASS = "internal.task.assignor.class";
}

Note that the thread-level assignment will remain an un-configurable internal implementation detail of the partition assignor (see "Rejected Alternatives" for further thoughts and reasoning).

To enable users to actually plug something in by implementing taskAssignor , we will need to move the TaskAssignor interface from o.a.k.streams.processor.internals.assignment to a public package, along with some of the supporting classes such as the assignment configs container class and ClientState which both appear in the TaskAssignor#assign method (although those will be heavily refactored, discussed below). All these new public APIs will be placed in a new non-internal package that mirrors their old internal package, specifically org.apache.kafka.streams.processor.assignment.

Both the input parameter and return value will be encapsulated in wrapper classes for the sake of forwards compatibility. This will let us avoid the cycle of adding, deprecating, and removing new #assign overloads if/when we want to evolve the assignor in the future, for example to pass in additional metadata or enable the assignor to output new kinds of information or instructions to the StreamsPartitionAssignor. The analogous ConsumerPartitionAssignor works similarly, returning a single GroupAssignment object that wraps the collection of individual consumer assignments for the same reason.

package org.apache.kafka.streams.processor.assignment

...

Both the input parameter and return value will be encapsulated in wrapper classes for the sake of forwards compatibility. This will let us avoid the cycle of adding, deprecating, and removing new #assign overloads if/when we want to evolve the assignor in the future, for example to pass in additional metadata or enable the assignor to output new kinds of information or instructions to the StreamsPartitionAssignor. The analogous ConsumerPartitionAssignor works similarly, returning a single GroupAssignment object that wraps the collection of individual consumer assignments for the same reason.

package org.apache.kafka.streams.processor.assignment;

public interface TaskAssignor extends Configurable {

  /**
   * @param applicationMetadata the metadata for this Streams application
   ;

public interface TaskAssignor extends Configurable {

  /**
   * @param applicationMetadata the metadata for this Streams application
   * @return the assignment of active and standby tasks to Streams client nodes
   */
  TaskAssignment assign(final ApplicationMetadata applicationMetadata);

  /**
   * Wrapper class for the final assignment of active and standbys tasks to individual Streams 
   * client nodes
   */
  class TaskAssignment {
    private final Collection<NodeAssignment> nodeAssignments;

	/**
     * @return the assignment of active and standby tasks to Streamsnodes
 client nodes
   */
   TaskAssignment assign(finalpublic ApplicationMetadata applicationMetadata)Collection<NodeAssignment> assignment();

    /**
   * Wrapper class* for@return the final assignmentnumber of activeStreams andclient standbysnodes tasksto towhich individualtasks Streamswere assigned
   * client nodes
   */
  class TaskAssignment {
public int numNodes();

  private final Collection<NodeAssignment> nodeAssignments;

	/**
     * @return a theString assignmentrepresentation of tasks to nodes
     */
    public Collection<NodeAssignment> assignment();

    /**the returned assignment, in processId order
     * @return the number of Streams client nodes to which tasks were assigned/
     */@Override
    public intString numNodestoString();

    /**
     * @return a String representation of the returned assignment, in processId order
     */
    @Override
    public String toString();
  }
}
}

Another reason for introducing the new TaskAssignment and ApplicationMetadata classes is to clean up the way assignment is performed today, as the current API is really not fit for public consumption. Currently, the TaskAssignor is provided a set of ClientState objects representing each client node. The ClientState is however not just the input to the assignor, but also its output – the assignment of tasks to nodes is performed by mutating the ClientStates passed in. The return value of the #assign method is a simple boolean indicating to the StreamsPartitionAssignor whether it should request a followup probing rebalance, a feature associated only with the HighAvailabilityTaskAssignor.

...

1. To provide a clean separation of input/output by splitting the ClientState into a readan input-only NodeState metadata class and a writean output-only NodeAssignment return value class
2. To decouple the followup rebalance request from the probing rebalance feature and give the assignor more direct control over the followup rebalance schedule, by allowing it to indicate which node(s) should trigger a rejoin and when to request the subsequent rebalance

...

package org.apache.kafka.streams.processor.assignment;

/**
 * A simple writewrapper classaround forUUID thethat assignorabstracts toa fillProcess inID
public theclass desiredProcessID placement{
 of active and standbyprivate tasksfinal on Streams client nodes 
 */
public interface NodeAssignmentUUID id;

    public ProcessID(final UUID id) {
   UUID processId();

    void assignActive(final TaskId task);

  void assignStandby(final TaskId task); this.id = id;
    }

  void removeActive(final TaskIdpublic taskid();

 {
      void removeStandby(final TaskIdreturn task)id;

   Set<TaskId> activeAssignment();}

   Set<TaskId> activeStatefulAssignmentint hashCode(); {
  
   Set<TaskId> activeStatelessAssignment();

  Set<TaskId>return standbyAssignmentid.hashCode();

  /**
  }

 * Request a followupboolean rebalanceequals(final toProcessID beother) triggered{
 by  one  of  the if consumers(other on== thisnull node after the|| getClass() != other.getClass()) {
   * given interval has elapsed. The {@link StreamThread} selected forreturn thisfalse;
 will be chosen at random.
   * <p>}
   * NOTE: A best effort will be made to enforce a rebalance according to the requested schedule,
   * but there is no guarantee that another rebalance will not occur before this time has elapsed.
   * Similarly, there is no guarantee the followup rebalance will occur, and must be re-requested
   * if, for example, the requesting consumer crashes or drops out of the group. Such an event
   * is, however, guaranteed to trigger a new rebalance itself, at which point the assignor
   * can re-evaluate whether to request an additional rebalance or not.
   *
   * @param followupRebalanceDelay how long this node should wait before initiating a new rebalance
   */
  void requestFollowupRebalance(final Duration followupRebalanceDelay);

return id.equals(other.id);
    }
}

/**
 * A simple interface for the assignor to return the desired placement of active and standby tasks on Streams client nodes 
 */
public interface NodeAssignment {
  ProcessID processId();

  Set<TaskId> activeAssignment();

  Set<TaskId> activeStatefulAssignment();
  
  Set<TaskId> activeStatelessAssignment();

  Set<TaskId> standbyAssignment();

  /**
   * @return the actual deadline in objective time, using ms since the epoch, after which the
   * followup rebalance will be attempted. Equivalent to {@code 'now + followupRebalanceDelay'}
   */
  long followupRebalanceDeadline();
 }

and

package org.apache.kafka.streams.processor.assignment;

/**
 * A read-only metadata class representing the current state of each Streams client node with at least one StreamThread participating in this rebalance
 */
public interface NodeState {
  /**
   * @return the actualprocessId deadlineof inthe objectiveapplication time,instance usingrunning mson sincethis thenode
 epoch, after which the*/
  ProcessID processId();

  /**
   * followup rebalance will be attempted. Equivalent to {@code 'now + followupRebalanceDelay'} Returns the number of StreamThreads on this node, which is equal to the number of main consumers
   */
 and represents long followupRebalanceDeadline();
 }

and

package org.apache.kafka.streams.processor.assignment;

/**
 * A read-only metadata class representing the current state of each Streams client node with at least one StreamThread participating in this rebalance
 */
public interface NodeState {
    
}

The NodeState will be wrapped up along with the other inputs to the assignor (such as the configuration and set of tasks to be assigned, as well as various utilities that may be useful) in the final new interface, the ApplicationMetadata:

package org.apache.kafka.streams.processor.assignment;

/**
 * A read-only metadata class representing the current state of each Streams client node with at least one StreamThread participating in this rebalance
 */
public interface ApplicationMetadata {
    its overall capacity.
   * <p>
   * NOTE: this is actually the "minimum capacity" of a node, or the minimum number of assigned
   * active tasks below which the node will have been over-provisioned and unable to give every
   * available StreamThread an active task assignment
   *
   * @return the number of consumers on this node
   */
  int numStreamThreads();

  /**
   * @return the set of consumer client ids for all StreamThreads on the given node
   */
  SortedSet<String> consumers();

  /**
     * @return athe mapset fromof theall {@codeactive processId}tasks toowned {@linkby NodeState}consumers foron allthis Streamsnode clientsince nodesthe inprevious this apprebalance
     */
    Map<UUID, NodeState> nodeStatesSortedSet<TaskId> previousActiveTasks();

    /**
     * Makes@return athe remoteset callof toall fetchstandby changelogtasks topicowned endby offsetsconsumers and,on ifthis successful,node usessince the resultsprevious to computerebalance
   */
  * task lags for each {@link NodeState}.
     *
     * @return whetherSortedSet<TaskId> previousStandbyTasks();

  /**
   * Returns the total lag across all logged stores in the task. Equal to the end offset fetchsum andif lag computation was successfulthis client
   * did */
not have any state boolean computeTaskLags();

 for this task on disk.
   /**
     * @return aend simpleoffset containersum class- withoffset thesum
 Streams configs relevant* to assignment
     */
   Task.LATEST_OFFSET AssignmentConfigs assignmentConfigs();

    /**
     * @return the set of all tasks in this topology which must be assigned to a nodeif this was previously an active running task on this client
     */
  long lagFor(final Set<TaskId>TaskId allTasks(task);

    /**
     *
 @return the previous tasks *assigned @returnto thethis setconsumer ofordered stateful and changelogged tasks by lag, filtered for any tasks that don't exist in this topologyassignment
     */
  SortedSet<TaskId> prevTasksByLag(final Set<TaskId>String statefulTasks(consumer);

    /**
   * Returns a *collection Assigncontaining standbyall tasks(and toonly) nodesstateful accordingtasks toin the defaulttopology logic.
by {@link TaskId},
   * <p>
mapped to its "offset lag * If rack-aware client tags are configured, the rack-aware standby task assignor will be used
     *
     * @param nodeAssignments the current assignment of tasks to nodes
  sum". This is computed as the difference between the changelog end offset
   * and the current offset, summed across all logged state stores in the task.
   *
   * @return a map from all stateful tasks to their lag sum
   */
  Map<TaskId, Long> void defaultStandbyTaskAssignment(final Map<UUID, NodeAssignment> nodeAssignmentsstatefulTasksToLagSums();

    /**
   * The * Optimize the active task assignment for rack-awareness{@link HostInfo} of this node, if set via the
   *  *
  {@link org.apache.kafka.streams.StreamsConfig#APPLICATION_SERVER_CONFIG application.server} config
   *
 @param nodeAssignments the* current@return assignmentthe ofhost tasksinfo tofor nodes
this node if configured, else * @param tasks the set of tasks to reassign if possible. Must already be assigned to a node
     */
    void optimizeRackAwareActiveTasks(final Map<UUID, NodeAssignment> nodeAssignments, final SortedSet<TaskId> tasks);

    /**
     * Optimize the standby task assignment for rack-awareness
     *
     * @param nodeAssignments the current assignment of tasks to nodes
     */
    void optimizeRackAwareStandbyTasks(final Map<UUID, NodeAssignment> nodeAssignments); 
}{@code Optional.empty()}
   */
  Optional<HostInfo> hostInfo();

  /**
   * The client tags for this client node, if set any have been via configs using the
   * {@link org.apache.kafka.streams.StreamsConfig#clientTagPrefix}
   * <p>
   * Can be used however you want, or passed in to enable the rack-aware standby task assignor.
   *
   * @return all the client tags found in this node's {@link org.apache.kafka.streams.StreamsConfig}
   */
  Map<String, String> clientTags();
 }

The NodeState will be wrapped up along with the other inputs to the assignor (such as the configuration and set of tasks to be assigned, as well as various utilities that may be useful) in the final new interface, the ApplicationMetadata. The methods on the ApplicationMetadata are The first four methods on the ApplicationMetadata are basically just the current inputs to the #assign method. The last three methods are new APIs but pre-existing functionality, essentially presenting a clean way for users to take advantage of the current optimizations and algorithms that are utilized by the built-in assignors, so that users don't have to re-implement complex features such as rack-awareness. The #defaultStandbyTaskAssignment API will just delegate to the appropriate standby task assignor (either basic default or client tag based standby rack awareness, depending on the existence of client tags in the configuration). Similarly, the #optimizeRackAware{Active/Standby}Tasks API will just delegate to the new RackAwareTaskAssignor that is being added in KIP-925.

Last, we have the AssignmentConfigs, which are (and would remain) just a basic container class, although we will migrate from public fields to standard getters for each of the configs passed into the assignor. Going forward, when a KIP is proposed to introduce a new config intended for the assignor, it should include the appropriate getter(s) in this class as part of the accepted proposal.

...

...

:

package org.apache.kafka.streams.processor.assignment;

/**
 * A read-only metadata class representing the current state of each Streams client node with at least one StreamThread participating in this rebalance
 */
public interface ApplicationMetadata {
    /**
     * @return a map from the {@code processId} to {@link NodeState} for all Streams client nodes in this app
     */
    Map<ProcessID, NodeState> nodeStates();

    /**
     * Makes a remote call to fetch changelog topic end offsets and, if successful, uses the results to compute
     * task lags for each {@link NodeState}.
     *
     * @return whether the end offset fetch and lag computation was successful
     */
    boolean computeTaskLags();

    /**
     * @return a simple container class with the Streams configs relevant to assignment
     */
    AssignmentConfigs assignmentConfigs();

    /**
     * @return the set of all tasks in this topology which must be assigned to a node
     */
    Set<TaskId> allTasks();

    /**
     *
     * @return the set of stateful and changelogged tasks in this topology
     */
    Set<TaskId> statefulTasks(); 
}

We'll also move some of the existing assignment functionality into a utils class that can be called by implementors of TaskAssignor:

package org.apache.kafka.streams.processor.assignment;

/**
 * A set of utilities to help implement task assignment
 */
public final class TaskAssignmentUtils {
    /**
     * Assign standby tasks to nodes according to the default logic.
     * <p>
     * If rack-aware client tags are configured, the rack-aware standby task assignor will be used
     *
     * @param nodeAssignments the current assignment of tasks to nodes
     */
    public static void defaultStandbyTaskAssignment(final ApplicationMetadata applicationMetadata, final Map<ProcessID, NodeAssignment> nodeAssignments) {...}

    /**
     * Optimize the active task assignment for rack-awareness
     *
     * @param nodeAssignments the current assignment of tasks to nodes
     * @param tasks the set of tasks to reassign if possible. Must already be assigned to a node
     */
    public static void optimizeRackAwareActiveTasks(final ApplicationMetadata applicationMetadata, final Map<ProcessID, NodeAssignment> nodeAssignments, final SortedSet<TaskId> tasks) {...}

    /**
     * Optimize the standby task assignment for rack-awareness
     *
     * @param nodeAssignments the current assignment of tasks to nodes
     */
    public static void optimizeRackAwareStandbyTasks(final ApplicationMetadata applicationMetadata, final Map<ProcessID, NodeAssignment> nodeAssignments) {...}
}

TaskAssignmentUtils  provides new APIs but pre-existing functionality, essentially presenting a clean way for users to take advantage of the current optimizations and algorithms that are utilized by the built-in assignors, so that users don't have to re-implement complex features such as rack-awareness. The #defaultStandbyTaskAssignment API will just delegate to the appropriate standby task assignor (either basic default or client tag based standby rack awareness, depending on the existence of client tags in the configuration). Similarly, the #optimizeRackAware{Active/Standby}Tasks API will just delegate to the new RackAwareTaskAssignor that is being added in KIP-925.

Last, we have the AssignmentConfigs, which are (and would remain) just a basic container class, although we will migrate from public fields to standard getters for each of the configs passed into the assignor. Going forward, when a KIP is proposed to introduce a new config intended for the assignor, it should include the appropriate getter(s) in this class as part of the accepted proposal.

package org.apache.kafka.streams.processor.assignment;

public class AssignmentConfigs {
    public long acceptableRecoveryLag();
    public int maxWarmupReplicas();
    public int numStandbyReplicas();
    public long probingRebalanceIntervalMs();
    public List<String> rackAwareAssignmentTags();
    public int trafficCost();
    public int nonOverlapCost();
}

Finally, as part of this change, we're moving some of the behavior that can fail into the task assignor. In particular, we're moving the bits that compute lags for stateful tasks into the implementation of ApplicationMetadata.computeTaskLags . This means we need some way to communicate to the streams partition assignor that it should retain the same assignment and schedule a follow-up rebalance. To do this, we will add the exception type StreamsAssignorRetryableException . If the TaskAssignor  throws this exception, StreamsPartitionAssignor  catches it, does fallback assignment, and schedules a follow-up rebalance.

Proposed Changes

No actual changes to functionality, mainly moving an internal config to the public part of StreamsConfig and bringing along a few currently-internal classes into the public API as some new interfaces and a new public assignment package. Code-wise the largest change is really the breaking up of the ClientState into the new NodeState and NodeAssignment interfaces, but that will be handled transparently to the user for all existing built-in-assignors, which will continue to work the same as before. 

...

1. One obvious question here is whether we want to still deprecate the old internal config anyway, out of compassion for any who may already be using it despite it not being considered public. Personally I think this would be reasonable but don't feel strongly one way or another.
2. Another possibility that was considered and ultimately decided against was whether to encompass the thread-level assignment in this KIP, and bring that into the public API or and make it pluggable as well. We determined that this did not seem necessary to do as part of the initial KIP, especially considering the large scope we have already reached. However it's worth noting that a followup KIP that builds on the new public API(s) introduced here would become much more feasible should someone wish to customize the thread-level logic at some point in the future. If/when that question is brought up, we'll have to address a few other concerns we had that contributed to our decision to exclude this for now, such as validating the thread assignment for correctness according to the cooperative rebalancing protocol, or niche optimizations like transient standbys to avoid losing in-memory state, and some other subtle logic that currently resides in the last leg of the StreamsPartitionAssignor's algorithm that tackles the distribution of node tasks to threads

...