KAFKA-13836: Improve KRaft broker heartbeat logic

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

@@ -29,8 +29,8 @@ public class BrokerHeartbeatRequest extends AbstractRequest {
     public static class Builder extends AbstractRequest.Builder<BrokerHeartbeatRequest> {
         private final BrokerHeartbeatRequestData data;
 
-        public Builder(BrokerHeartbeatRequestData data) {
-            super(ApiKeys.BROKER_HEARTBEAT);
+        public Builder(BrokerHeartbeatRequestData data, short version) {
+            super(ApiKeys.BROKER_HEARTBEAT, version);
             this.data = data;
         }
 

clients/src/main/resources/common/message/BrokerHeartbeatRequest.json

@@ -18,14 +18,17 @@
   "type": "request",
   "listeners": ["controller"],
   "name": "BrokerHeartbeatRequest",
-  "validVersions": "0",
+  // Version 1 add CurrentPublishedOffset.
+  "validVersions": "0-1",
   "flexibleVersions": "0+",
   "fields": [
     { "name": "BrokerId", "type": "int32", "versions": "0+", "entityType": "brokerId",
       "about": "The broker ID." },
     { "name": "BrokerEpoch", "type": "int64", "versions": "0+", "default": "-1",
       "about": "The broker epoch." },
     { "name": "CurrentMetadataOffset", "type": "int64", "versions": "0+",
+      "about": "The highest metadata offset which the broker has seen." },
+    { "name": "CurrentPublishedOffset", "type": "int64", "versions": "1+",
       "about": "The highest metadata offset which the broker has reached." },
     { "name": "WantFence", "type": "bool", "versions": "0+",
       "about": "True if the broker wants to be fenced, false otherwise." },

clients/src/main/resources/common/message/BrokerHeartbeatResponse.json

@@ -17,7 +17,8 @@
   "apiKey": 63,
   "type": "response",
   "name": "BrokerHeartbeatResponse",
-  "validVersions": "0",
+  // Version 1 is the same as version 0.
+  "validVersions": "0-1",
   "flexibleVersions": "0+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",

clients/src/test/java/org/apache/kafka/common/requests/RequestResponseTest.java

@@ -3252,9 +3252,10 @@ private BrokerHeartbeatRequest createBrokerHeartbeatRequest(short v) {
                 .setBrokerId(1)
                 .setBrokerEpoch(1)
                 .setCurrentMetadataOffset(1)
+                .setCurrentPublishedOffset(1)
                 .setWantFence(false)
                 .setWantShutDown(false);
-        return new BrokerHeartbeatRequest.Builder(data).build(v);
+        return new BrokerHeartbeatRequest.Builder(data, ApiKeys.BROKER_HEARTBEAT.latestVersion()).build(v);
     }
 
     private BrokerHeartbeatResponse createBrokerHeartbeatResponse() {

core/src/main/scala/kafka/server/BrokerLifecycleManager.scala

@@ -30,6 +30,8 @@ import org.apache.kafka.metadata.{BrokerState, VersionRange}
 import org.apache.kafka.queue.EventQueue.DeadlineFunction
 import org.apache.kafka.common.utils.{ExponentialBackoff, LogContext, Time}
 import org.apache.kafka.queue.{EventQueue, KafkaEventQueue}
+import org.apache.kafka.server.common.MetadataVersion
+
 import scala.jdk.CollectionConverters._
 
 
@@ -51,9 +53,12 @@ import scala.jdk.CollectionConverters._
  * In some cases we expose a volatile variable which can be read from any thread, but only
  * written from the event queue thread.
  */
-class BrokerLifecycleManager(val config: KafkaConfig,
-                             val time: Time,
-                             val threadNamePrefix: Option[String]) extends Logging {
+class BrokerLifecycleManager(
+  val config: KafkaConfig,
+  val featureCache: FinalizedFeatureCache,
+  val time: Time,
+  val threadNamePrefix: Option[String]
+) extends Logging {
   val logContext = new LogContext(s"[BrokerLifecycleManager id=${config.nodeId}] ")
 
   this.logIdent = logContext.logPrefix()
@@ -126,8 +131,16 @@ class BrokerLifecycleManager(val config: KafkaConfig,
   private var _highestMetadataOffsetProvider: () => Long = _
 
   /**
-   * True only if we are ready to unfence the broker.  This variable can only be read or
-   * written from the event queue thread.
+   * A thread-safe callback function which gives this manager the current highest published
+   * offset.
+   * This variable can only be read or written from the event queue thread.
+   */
+  private var _highestPublishedOffsetProvider: () => Long = _
+
+  /**
+   * True only if we are ready to unfence the broker, which means the metadata listener has
+   * completed its first publish operation. This variable can only be read or written from
+   * the event queue thread.
    */
   private var readyToUnfence = false
 
@@ -186,12 +199,18 @@ class BrokerLifecycleManager(val config: KafkaConfig,
    * @param clusterId                     The cluster ID.
    */
   def start(highestMetadataOffsetProvider: () => Long,
+            highestPublishedOffsetProvider: () => Long,
             channelManager: BrokerToControllerChannelManager,
             clusterId: String,
             advertisedListeners: ListenerCollection,
             supportedFeatures: util.Map[String, VersionRange]): Unit = {
-    eventQueue.append(new StartupEvent(highestMetadataOffsetProvider,
-      channelManager, clusterId, advertisedListeners, supportedFeatures))
+    eventQueue.append(new StartupEvent(
+      highestMetadataOffsetProvider,
+      highestPublishedOffsetProvider,
+      channelManager,
+      clusterId,
+      advertisedListeners,
+      supportedFeatures))
   }
 
   def setReadyToUnfence(): CompletableFuture[Void] = {
@@ -253,13 +272,16 @@ class BrokerLifecycleManager(val config: KafkaConfig,
     }
   }
 
-  private class StartupEvent(highestMetadataOffsetProvider: () => Long,
-                     channelManager: BrokerToControllerChannelManager,
-                     clusterId: String,
-                     advertisedListeners: ListenerCollection,
-                     supportedFeatures: util.Map[String, VersionRange]) extends EventQueue.Event {
+  private class StartupEvent(
+    highestMetadataOffsetProvider: () => Long,
+    highestPublishedOffsetProvider: () => Long,
+    channelManager: BrokerToControllerChannelManager,
+    clusterId: String,
+    advertisedListeners: ListenerCollection,
+    supportedFeatures: util.Map[String, VersionRange]) extends EventQueue.Event {
     override def run(): Unit = {
       _highestMetadataOffsetProvider = highestMetadataOffsetProvider
+      _highestPublishedOffsetProvider = highestPublishedOffsetProvider
       _channelManager = channelManager
       _channelManager.start()
       _state = BrokerState.STARTING
@@ -322,6 +344,7 @@ class BrokerLifecycleManager(val config: KafkaConfig,
           _brokerEpoch = message.data().brokerEpoch()
           registered = true
           initialRegistrationSucceeded = true
+          _state = BrokerState.REGISTERED
           info(s"Successfully registered broker $nodeId with broker epoch ${_brokerEpoch}")
           scheduleNextCommunicationImmediately() // Immediately send a heartbeat
         } else {
@@ -340,16 +363,23 @@ class BrokerLifecycleManager(val config: KafkaConfig,
 
   private def sendBrokerHeartbeat(): Unit = {
     val metadataOffset = _highestMetadataOffsetProvider()
+    val publishedOffset = _highestPublishedOffsetProvider()
+    val version = if (featureCache.getMetadataVersion().exists(_ >= MetadataVersion.IBP_3_3_IV1.featureLevel())) {
+      1.toShort
+    } else {
+      0.toShort
+    }
     val data = new BrokerHeartbeatRequestData().
       setBrokerEpoch(_brokerEpoch).
       setBrokerId(nodeId).
       setCurrentMetadataOffset(metadataOffset).
+      setCurrentPublishedOffset(publishedOffset).
       setWantFence(!readyToUnfence).
       setWantShutDown(_state == BrokerState.PENDING_CONTROLLED_SHUTDOWN)
     if (isTraceEnabled) {
       trace(s"Sending broker heartbeat $data")
     }
-    _channelManager.sendRequest(new BrokerHeartbeatRequest.Builder(data),
+    _channelManager.sendRequest(new BrokerHeartbeatRequest.Builder(data, version),
       new BrokerHeartbeatResponseHandler())
   }
 
@@ -376,7 +406,7 @@ class BrokerLifecycleManager(val config: KafkaConfig,
         if (errorCode == Errors.NONE) {
           failedAttempts = 0
           _state match {
-            case BrokerState.STARTING =>
+            case BrokerState.REGISTERED =>
               if (message.data().isCaughtUp) {
                 info(s"The broker has caught up. Transitioning from STARTING to RECOVERY.")
                 _state = BrokerState.RECOVERY

core/src/main/scala/kafka/server/BrokerServer.scala

@@ -150,7 +150,7 @@ class BrokerServer(
 
   var metadataListener: BrokerMetadataListener = null
 
-  var metadataPublisher: BrokerMetadataPublisher = null
+  var _metadataPublisher: BrokerMetadataPublisher = null
 
   val brokerFeatures: BrokerFeatures = BrokerFeatures.createDefault()
 
@@ -177,14 +177,18 @@ class BrokerServer(
 
   def replicaManager: ReplicaManager = _replicaManager
 
+  def metadataPublisher: BrokerMetadataPublisher = _metadataPublisher
+
   override def startup(): Unit = {
     if (!maybeChangeStatus(SHUTDOWN, STARTING)) return
     try {
       info("Starting broker")
 
       config.dynamicConfig.initialize(zkClientOpt = None)
+      featureCache = new FinalizedFeatureCache(brokerFeatures)
 
       lifecycleManager = new BrokerLifecycleManager(config,
+        featureCache,
         time,
         threadNamePrefix)
 
@@ -226,8 +230,6 @@ class BrokerServer(
       clientToControllerChannelManager.start()
       forwardingManager = new ForwardingManagerImpl(clientToControllerChannelManager)
 
-      featureCache = new FinalizedFeatureCache(brokerFeatures)
-
       val apiVersionManager = ApiVersionManager(
         ListenerType.BROKER,
         config,
@@ -342,10 +344,22 @@ class BrokerServer(
           k -> VersionRange.of(v.min, v.max)
       }.asJava
 
-      lifecycleManager.start(() => metadataListener.highestMetadataOffset,
-        BrokerToControllerChannelManager(controllerNodeProvider, time, metrics, config,
-          "heartbeat", threadNamePrefix, config.brokerSessionTimeoutMs.toLong),
-        metaProps.clusterId, networkListeners, featuresRemapped)
+      lifecycleManager.start(
+        () => metadataListener.highestMetadataOffset,
+        () => metadataPublisher.publishedOffset,
+        BrokerToControllerChannelManager(
+          controllerNodeProvider,
+          time,
+          metrics,
+          config,
+          "heartbeat",
+          threadNamePrefix,
+          config.brokerSessionTimeoutMs.toLong
+        ),
+        metaProps.clusterId,
+        networkListeners,
+        featuresRemapped
+      )
 
       // Register a listener with the Raft layer to receive metadata event notifications
       raftManager.register(metadataListener)
@@ -422,7 +436,7 @@ class BrokerServer(
       lifecycleManager.initialCatchUpFuture.get()
 
       // Apply the metadata log changes that we've accumulated.
-      metadataPublisher = new BrokerMetadataPublisher(config,
+      _metadataPublisher = new BrokerMetadataPublisher(config,
         metadataCache,
         logManager,
         replicaManager,
@@ -438,7 +452,7 @@ class BrokerServer(
       // replicaManager, groupCoordinator, and txnCoordinator. The log manager may perform
       // a potentially lengthy recovery-from-unclean-shutdown operation here, if required.
       try {
-        metadataListener.startPublishing(metadataPublisher).get()
+        metadataListener.startPublishing(_metadataPublisher).get()
       } catch {
         case t: Throwable => throw new RuntimeException("Received a fatal error while " +
           "waiting for the broker to catch up with the current cluster metadata.", t)

core/src/main/scala/kafka/server/ControllerApis.scala

@@ -545,7 +545,9 @@ class ControllerApis(val requestChannel: RequestChannel,
     authHelper.authorizeClusterOperation(request, CLUSTER_ACTION)
     val context = new ControllerRequestContext(request.context.principal,
       requestTimeoutMsToDeadlineNs(time, config.brokerHeartbeatIntervalMs))
-    controller.processBrokerHeartbeat(context, heartbeatRequest.data).handle[Unit] { (reply, e) =>
+    controller.processBrokerHeartbeat(
+      context, heartbeatRequest.data, heartbeatRequest.version()
+    ).handle[Unit] { (reply, e) =>
       def createResponseCallback(requestThrottleMs: Int,
                                  reply: BrokerHeartbeatReply,
                                  e: Throwable): BrokerHeartbeatResponse = {

core/src/main/scala/kafka/server/FinalizedFeatureCache.scala

@@ -151,6 +151,10 @@ class FinalizedFeatureCache(private val brokerFeatures: BrokerFeatures) extends
       Collections.unmodifiableMap(newFeatures)), highestMetadataOffset))
   }
 
+  def getMetadataVersion(): Option[Short] = {
+    featuresAndEpoch.flatMap(featuresAndEpoch => Option(featuresAndEpoch.features.features().get(MetadataVersion.FEATURE_NAME))).map(_.max())
+  }
+
   /**
    * Causes the current thread to wait no more than timeoutMs for the specified condition to be met.
    * It is guaranteed that the provided condition will always be invoked only from within a

core/src/test/java/kafka/test/MockController.java

@@ -367,7 +367,8 @@ public CompletableFuture<Map<ConfigResource, ApiError>> legacyAlterConfigs(
     @Override
     public CompletableFuture<BrokerHeartbeatReply> processBrokerHeartbeat(
         ControllerRequestContext context,
-        BrokerHeartbeatRequestData request
+        BrokerHeartbeatRequestData request,
+        short version
     ) {
         throw new UnsupportedOperationException();
     }