insights/proto/kafka_assigner/v1/service.proto

syntax = "proto3";
package kafka_assigner.v1;

// KafkaAssigner is the coordination API for Kafka partition consumers.
//
// Consumers call the Register RPC to receive assignment commands.
// All coordination state lives in etcd; this gRPC API is a relay layer
// that removes the etcd dependency from consumer code.
//
// Every assigner instance serves gRPC. Consumers can connect to any
// instance. Each instance watches etcd and relays commands to its
// locally-connected consumers.
service KafkaAssigner {
  // Register registers a consumer and opens a server-side stream of
  // assignment commands. The first message is always a Snapshot
  // containing the consumer's current partition assignments.
  //
  // The stream remains open for the lifetime of the consumer.
  // If the stream drops, the consumer's etcd lease expires and
  // its partitions are reassigned.
  rpc Register(RegisterRequest) returns (stream AssignmentCommand);

  // PartitionReady signals that the consumer has finished warming
  // a partition and is ready to take ownership.
  rpc PartitionReady(PartitionReadyRequest) returns (PartitionReadyResponse);

  // PartitionReleased signals that the consumer has released a
  // partition after a handoff completion.
  rpc PartitionReleased(PartitionReleasedRequest) returns (PartitionReleasedResponse);

  // Deregister signals that a consumer is shutting down gracefully.
  //
  // The coordinator transitions the consumer to Draining status,
  // which excludes it from new partition assignments and triggers
  // handoffs for its current partitions.
  //
  // The response tells the consumer whether to wait for all Release
  // events to complete (WAIT_FOR_DRAIN) or to exit immediately
  // (SHUTDOWN_NOW). For example, during a StatefulSet rollout the
  // same pod name will come back, so waiting is unnecessary.
  rpc Deregister(DeregisterRequest) returns (DeregisterResponse);
}

// ── Register ─────────────────────────────────────────────

message RegisterRequest {
  // Unique consumer name (e.g., pod name or hostname).
  string consumer_name = 1;
  // Kafka topic this consumer wants partitions for.
  // The assigner ensures a TopicConfig exists in etcd for this topic,
  // fetching the partition count from Kafka if needed.
  string topic = 2;
}

// AssignmentCommand is a server-pushed command on the Register stream.
message AssignmentCommand {
  oneof command {
    // Batch assignment update: partitions added and/or removed.
    // The first message on connect is an assignment with the consumer's
    // current partitions in `assigned` and an empty `unassigned`.
    AssignmentUpdate assignment = 1;

    // Warm a partition in preparation for taking ownership.
    // Consumer should call PartitionReady when warming completes.
    WarmPartition warm = 2;

    // Release a partition. The handoff is complete and the consumer
    // should stop processing this partition.
    ReleasePartition release = 3;
  }
}

message AssignmentUpdate {
  repeated TopicPartition assigned = 1;
  repeated TopicPartition unassigned = 2;
}

message WarmPartition {
  TopicPartition partition = 1;
  // The current owner of the partition (for diagnostics).
  string current_owner = 2;
}

message ReleasePartition {
  TopicPartition partition = 1;
  // The new owner taking over (for diagnostics).
  string new_owner = 2;
}

message TopicPartition {
  string topic = 1;
  uint32 partition = 2;
}

// ── PartitionReady ───────────────────────────────────────

message PartitionReadyRequest {
  string consumer_name = 1;
  TopicPartition partition = 2;
}

message PartitionReadyResponse {}

// ── PartitionReleased ────────────────────────────────────

message PartitionReleasedRequest {
  string consumer_name = 1;
  TopicPartition partition = 2;
}

message PartitionReleasedResponse {}

// ── Deregister ──────────────────────────────────────────

message DeregisterRequest {
  // The consumer name that is shutting down.
  string consumer_name = 1;
}

message DeregisterResponse {
  DeregisterAction action = 1;
}

enum DeregisterAction {
  // Unspecified / default. Treated the same as SHUTDOWN_NOW.
  DEREGISTER_ACTION_UNSPECIFIED = 0;
  // Exit immediately. Used when the coordinator expects the pod to
  // come back under the same name (e.g., StatefulSet rollout) or
  // when K8s awareness is unavailable.
  DEREGISTER_ACTION_SHUTDOWN_NOW = 1;
  // Keep the gRPC stream open and wait for all Release events to
  // complete before exiting. Used when partitions must be drained
  // to other consumers (e.g., Deployment rollout or downscale).
  DEREGISTER_ACTION_WAIT_FOR_DRAIN = 2;
}