fix(inkless:switch): Fix stale HW on switched leader promotion (#660)

* fix(inkless:switch): Fix stale HW on switched leader promotion

Reconcile already-switched leaders after makeLeader so a promoted
replica with a stale checkpointed HW advances to the committed seal.
Truncate any local classic tail above the seal, and fence leaders
whose local LEO is below the seal to avoid serving an incomplete prefix.
Keep follower truncation limited to the seal-commit delta.
Add a deterministic system test for the stale-HW promotion regression.

* Skip reconciliation of a partition if log is not found

Author: giuseppelillo

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

@@ -3299,22 +3299,73 @@ class ReplicaManager(val config: KafkaConfig,
   }
 
   /**
-   * Reconciles a partition whose classic-to-diskless seal was *just committed* in this delta
-   * (previous seal < 0, now >= 0); other transitions are left untouched. Records beyond the seal
-   * are uncommitted (the diskless log continues from the seal), so:
-   *  - LEO > seal: truncate down to the seal (a follower/deposed leader that held un-replicated
-   *    records past a seal a different broker committed at a lower HW);
-   *  - LEO < seal and leader: fence offline. Unreachable on the clean path (the driving leader
-   *    seals at its own HW == LEO; a cleanly-elected leader comes from the ISR with LEO >= seal),
-   *    so it implies a corrupt classic prefix -- don't serve an incomplete prefix below the seal.
+   * Reconcile a leader whose classic-to-diskless seal has already been committed.
    *
-   * Must run after makeLeader/makeFollower (so the log exists) and before any catch-up or
-   * consolidation fetcher starts (so they initialize from the reconciled LEO).
+   * A committed seal is the first diskless offset and the classic prefix [0, seal) must be fully
+   * present on any cleanly-elected leader. Once sealed, the local classic log cannot advance HW
+   * naturally, so a leader promoted with a stale checkpointed HW must restore HW to the seal.
+   *
+   *  - LEO > seal: truncate down to the seal. Records at or above the seal are uncommitted classic
+   *    tail; diskless owns that range.
+   *  - LEO == seal and HW < seal: advance HW to the seal so consumers can cross into diskless.
+   *  - LEO < seal: fence offline. The classic prefix is incomplete, so serving it would hide a
+   *    hole in acknowledged data.
+   *
+   * Must run after makeLeader (so the log exists) and before any consolidation fetcher starts.
+   */
+  private def maybeReconcileSwitchedLeader(tp: TopicPartition,
+                                           topicId: Uuid,
+                                           newRegistration: PartitionRegistration): Unit = {
+    val seal = newRegistration.classicToDisklessStartOffset
+    if (seal < 0) return
+
+    onlinePartition(tp).foreach { partition =>
+      partition.log match {
+        case Some(log) =>
+          try {
+            if (log.logEndOffset > seal) {
+              stateChangeLogger.info(s"Truncating switched partition $tp from LEO ${log.logEndOffset} " +
+                s"to classic-to-diskless start offset $seal")
+              // Seal is the classicToDisklessStartOffset, the first offset owned by diskless storage.
+              // truncateTo(seal) removes all entries with offset >= seal, leaving LEO = seal.
+              // The last classic record is at offset seal - 1.
+              partition.truncateTo(seal, isFuture = false)
+            }
+            if (partition.isLeader) {
+              if (log.logEndOffset >= seal && log.highWatermark < seal) {
+                log.maybeUpdateHighWatermark(seal)
+                stateChangeLogger.info(s"Stale high watermark detected: advanced high watermark to seal offset $seal for " +
+                  s"switched leader partition $tp")
+              } else if (log.logEndOffset < seal) {
+                // This is unreachable in normal operation
+                stateChangeLogger.error(s"Leader partition $tp has LEO ${log.logEndOffset} below " +
+                  s"classic-to-diskless start offset $seal; cannot catch up from another replica. " +
+                  s"Marking the partition offline as its local log is corrupt below the committed seal.")
+                markPartitionOffline(tp)
+              }
+            }
+          } catch {
+            case e: KafkaStorageException =>
+              stateChangeLogger.error(s"Unable to reconcile switched partition $tp " +
+                s"with topic ID $topicId due to a storage error ${e.getMessage}", e)
+              markPartitionOffline(tp)
+          }
+        case None =>
+          stateChangeLogger.warn(s"Skipping switched partition reconciliation for $tp " +
+            s"with topic ID $topicId because the local log is not available")
+      }
+    }
+  }
+
+  /**
+   * Followers truncate only when the seal is committed in the current metadata delta. A switched
+   * follower may later grow past the seal as part of normal consolidation/catch-up flows; unlike a
+   * promoted leader, a follower should not be re-truncated on every unrelated metadata change.
    */
-  private def maybeTruncateNewlySwitchedPartition(tp: TopicPartition,
-                                                  topicId: Uuid,
-                                                  newRegistration: PartitionRegistration,
-                                                  delta: TopicsDelta): Unit = {
+  private def maybeTruncateNewlySwitchedFollower(tp: TopicPartition,
+                                                 topicId: Uuid,
+                                                 newRegistration: PartitionRegistration,
+                                                 delta: TopicsDelta): Unit = {
     val seal = newRegistration.classicToDisklessStartOffset
     if (seal < 0) return
 
@@ -3325,27 +3376,23 @@ class ReplicaManager(val config: KafkaConfig,
     if (!sealJustCommitted) return
 
     onlinePartition(tp).foreach { partition =>
-      try {
-        val log = partition.localLogOrException
-        if (log.logEndOffset > seal) {
-          stateChangeLogger.info(s"Truncating switched partition $tp from LEO ${log.logEndOffset} " +
-            s"to classic-to-diskless start offset $seal")
-          // Seal is the classicToDisklessStartOffset, the first offset owned by diskless storage.
-          // truncateTo(seal) removes all entries with offset >= seal, leaving LEO = seal.
-          // The last classic record is at offset seal - 1.
-          partition.truncateTo(seal, isFuture = false)
-        } else if (log.logEndOffset < seal && partition.isLeader) {
-          // This is unreachable in normal operation
-          stateChangeLogger.error(s"Leader partition $tp has LEO ${log.logEndOffset} below " +
-            s"classic-to-diskless start offset $seal; cannot catch up from another replica. " +
-            s"Marking the partition offline as its local log is corrupt below the committed seal.")
-          markPartitionOffline(tp)
-        }
-      } catch {
-        case e: KafkaStorageException =>
-          stateChangeLogger.error(s"Unable to reconcile switched partition $tp " +
-            s"with topic ID $topicId due to a storage error ${e.getMessage}", e)
-          markPartitionOffline(tp)
+      partition.log match {
+        case Some(log) =>
+          try {
+            if (log.logEndOffset > seal) {
+              stateChangeLogger.info(s"Truncating switched follower partition $tp from LEO ${log.logEndOffset} " +
+                s"to classic-to-diskless start offset $seal")
+              partition.truncateTo(seal, isFuture = false)
+            }
+          } catch {
+            case e: KafkaStorageException =>
+              stateChangeLogger.error(s"Unable to reconcile switched follower partition $tp " +
+                s"with topic ID $topicId due to a storage error ${e.getMessage}", e)
+              markPartitionOffline(tp)
+          }
+        case None =>
+          stateChangeLogger.warn(s"Skipping switched follower partition reconciliation for $tp " +
+            s"with topic ID $topicId because the local log is not available")
       }
     }
   }
@@ -3526,16 +3573,6 @@ class ReplicaManager(val config: KafkaConfig,
             // local log; seal() here just marks the partition sealed for the HW restore below.
             partition.makeLeader(info.partition, isNew, offsetCheckpoints, Some(info.topicId), partitionAssignedDirectoryId)
             partition.seal()
-            // makeLeader reloads HW from the on-disk checkpoint, which may be stale
-            // (unclean shutdown, or checkpoint interval hadn't fired). Since the
-            // partition is sealed, no produces or follower fetches can ever advance
-            // HW naturally — restore it to the seal offset so consumers can read.
-            val sealOffset = _inklessMetadataView.getClassicToDisklessStartOffset(tp)
-            if (sealOffset >= 0 && partition.localLogOrException.highWatermark < sealOffset) {
-              partition.localLogOrException.maybeUpdateHighWatermark(sealOffset)
-              stateChangeLogger.info(s"Advanced high watermark to seal offset $sealOffset for " +
-                s"switched leader partition $tp after restart")
-            }
             changedPartitions.add(partition)
           } catch {
             case e: KafkaStorageException =>
@@ -3547,13 +3584,11 @@ class ReplicaManager(val config: KafkaConfig,
       }
     }
 
-    // Truncate any partition whose seal was just committed in this delta and whose local log still
-    // runs past it (uncommitted records beyond the committed seal). This must happen after
-    // makeLeader (so the log exists) and before the consolidation fetchers below start, so
-    // consolidation initializes against the truncated LEO. For the broker that drove the switch
-    // this is a no-op (it seals at its own HW == LEO); it is kept here as defense.
+    // Reconcile switched leaders after makeLeader so promoted leaders with stale checkpointed HW
+    // are immediately readable up to the seal, and corrupt/incomplete prefixes are fenced before
+    // serving reads.
     localLeaders.foreachEntry { (tp, info) =>
-      maybeTruncateNewlySwitchedPartition(tp, info.topicId, info.partition, delta)
+      maybeReconcileSwitchedLeader(tp, info.topicId, info.partition)
     }
 
     if (consolidatingDisklessPartitionsToStartFetching.nonEmpty) {
@@ -3683,7 +3718,7 @@ class ReplicaManager(val config: KafkaConfig,
     // truncated LEO. The fetch decisions above key off the high watermark (which never exceeds
     // the seal), so truncating here does not change which partitions need a catch-up fetcher.
     localFollowers.foreachEntry { (tp, info) =>
-      maybeTruncateNewlySwitchedPartition(tp, info.topicId, info.partition, delta)
+      maybeTruncateNewlySwitchedFollower(tp, info.topicId, info.partition, delta)
     }
 
     if (partitionsToStartFetching.nonEmpty) {

core/src/test/scala/unit/kafka/server/DisklessSwitchInvariantsTest.scala

@@ -29,6 +29,7 @@ import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.record.{MemoryRecords, SimpleRecord}
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.image._
+import org.apache.kafka.metadata.{InitDisklessLogFields, PartitionRegistration}
 import org.apache.kafka.server.common.{KRaftVersion, MetadataVersion}
 import org.apache.kafka.server.config.ServerConfigs
 import org.apache.kafka.server.PartitionFetchState
@@ -84,18 +85,21 @@ class DisklessSwitchInvariantsTest {
 
   // --- Invariant: Sealed leader HW >= seal offset ---
 
-  @ParameterizedTest(name = "leader checkpoint hw={0}, seal={1}")
+  @ParameterizedTest(name = "leader checkpoint hw={0}, seal={1}, leo={2}")
   @MethodSource(Array("sealedLeaderScenarios"))
-  def testSealedLeaderHwInvariant(checkpointHw: Long, sealOffset: Long, expectedHw: Long): Unit = {
+  def testSealedLeaderHwInvariant(
+      checkpointHw: Long,
+      sealOffset: Long,
+      localLeo: Long,
+      expectedHw: Long,
+      expectedLeo: Long): Unit = {
     val topicName = "switched-topic"
     val topicId = Uuid.randomUuid()
     val tp = new TopicPartition(topicName, 0)
-    // For committed seals, LEO == seal (fully replicated). For pending/none, use a fixed LEO.
-    val leo = if (sealOffset >= 0) sealOffset else 10L
 
     replicaManager = createReplicaManager(List(topicName))
     val log = replicaManager.logManager.getOrCreateLog(tp, isNew = true, topicId = Optional.of(topicId))
-    populateLocalLogAtLeoAndCheckpointedHwm(replicaManager, tp, log, leo, checkpointHw)
+    populateLocalLogAtLeoAndCheckpointedHwm(replicaManager, tp, log, localLeo, checkpointHw)
 
     when(replicaManager.inklessMetadataView().getClassicToDisklessStartOffset(tp))
       .thenReturn(sealOffset)
@@ -110,6 +114,8 @@ class DisklessSwitchInvariantsTest {
     assertEquals(expectedHw, partition.localLogOrException.highWatermark,
       s"Sealed leader invariant: HW must be >= sealOffset ($sealOffset) " +
       s"regardless of checkpoint state ($checkpointHw)")
+    assertEquals(expectedLeo, partition.localLogOrException.logEndOffset,
+      s"Sealed leader invariant: LEO must be reconciled to expected offset $expectedLeo")
   }
 
   // --- Invariant: Sealed follower with HW < seal triggers catch-up fetcher ---
@@ -156,33 +162,44 @@ class DisklessSwitchInvariantsTest {
   private def leaderDelta(topicName: String, topicId: Uuid): TopicsDelta = {
     val delta = new TopicsDelta(TopicsImage.EMPTY)
     delta.replay(new TopicRecord().setName(topicName).setTopicId(topicId))
-    delta.replay(new PartitionRecord()
+    val partitionRecord = new PartitionRecord()
       .setPartitionId(0)
       .setTopicId(topicId)
       .setReplicas(util.Arrays.asList(brokerId, brokerId + 1))
       .setIsr(util.Arrays.asList(brokerId, brokerId + 1))
       .setLeader(brokerId)
       .setLeaderEpoch(0)
       .setPartitionEpoch(0)
-    )
+    addClassicToDisklessStartOffsetIfPresent(topicName, partitionRecord)
+    delta.replay(partitionRecord)
     delta
   }
 
   private def followerDelta(topicName: String, topicId: Uuid, leaderId: Int): TopicsDelta = {
     val delta = new TopicsDelta(TopicsImage.EMPTY)
     delta.replay(new TopicRecord().setName(topicName).setTopicId(topicId))
-    delta.replay(new PartitionRecord()
+    val partitionRecord = new PartitionRecord()
       .setPartitionId(0)
       .setTopicId(topicId)
       .setReplicas(util.Arrays.asList(brokerId, leaderId))
       .setIsr(util.Arrays.asList(brokerId, leaderId))
       .setLeader(leaderId)
       .setLeaderEpoch(0)
       .setPartitionEpoch(0)
-    )
+    addClassicToDisklessStartOffsetIfPresent(topicName, partitionRecord)
+    delta.replay(partitionRecord)
     delta
   }
 
+  private def addClassicToDisklessStartOffsetIfPresent(topicName: String, partitionRecord: PartitionRecord): Unit = {
+    val sealOffset = replicaManager.inklessMetadataView().getClassicToDisklessStartOffset(
+      new TopicPartition(topicName, partitionRecord.partitionId()))
+    if (sealOffset != PartitionRegistration.NO_CLASSIC_TO_DISKLESS_START_OFFSET) {
+      partitionRecord.unknownTaggedFields().add(
+        InitDisklessLogFields.encodeClassicToDisklessStartOffset(sealOffset))
+    }
+  }
+
   private def populateLocalLogAtLeoAndCheckpointedHwm(
       replicaManager: ReplicaManager, tp: TopicPartition,
       log: UnifiedLog, leo: Long, hw: Long): Unit = {
@@ -263,28 +280,30 @@ object DisklessSwitchInvariantsTest {
   import org.apache.kafka.metadata.PartitionRegistration._
 
   def sealedLeaderScenarios(): java.util.stream.Stream[Arguments] = {
-    // (checkpointHw, sealOffset, expectedHwAfterApplyDelta)
+    // (checkpointHw, sealOffset, localLeo, expectedHwAfterApplyDelta, expectedLeoAfterApplyDelta)
     // Committed seal (>= 0): HW is always advanced to sealOffset regardless of checkpoint.
     // The post-restart path (getOrCreatePartition + makeLeader) creates a new Partition object
     // that re-initializes HW from the checkpoint via LazyOffsetCheckpoints. However, since the
     // log already exists, LogManager returns it without re-reading the checkpoint. The seal
     // offset advancement in applyLocalLeadersDelta (sealOffset >= 0 && HW < sealOffset guard)
-    // compensates by forcing HW = sealOffset.
+    // compensates by forcing HW = sealOffset, even when an uncommitted tail must first be
+    // truncated from LEO > seal to LEO == seal.
     java.util.stream.Stream.of(
-      Arguments.of(0L: java.lang.Long,  10L: java.lang.Long, 10L: java.lang.Long),  // unclean shutdown
-      Arguments.of(5L: java.lang.Long,  10L: java.lang.Long, 10L: java.lang.Long),  // partial flush
-      Arguments.of(10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long),  // clean shutdown
+      Arguments.of(0L: java.lang.Long,  10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long),  // unclean shutdown
+      Arguments.of(5L: java.lang.Long,  10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long),  // partial flush
+      Arguments.of(5L: java.lang.Long,  10L: java.lang.Long, 13L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long),  // stale HW with uncommitted tail
+      Arguments.of(10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long, 10L: java.lang.Long),  // clean shutdown
       // Pending seal (-2): partition is sealed. The HW is NOT advanced because the seal
       // offset is not yet committed. The post-restart path does not restore HW from
       // checkpoint for pending seals — the pending fetcher (follower side) will eventually
       // propagate the correct HW when the seal commits.
-      Arguments.of(0L: java.lang.Long,  CLASSIC_TO_DISKLESS_SWITCH_PENDING: java.lang.Long, 0L: java.lang.Long),
-      Arguments.of(5L: java.lang.Long,  CLASSIC_TO_DISKLESS_SWITCH_PENDING: java.lang.Long, 0L: java.lang.Long),
+      Arguments.of(0L: java.lang.Long,  CLASSIC_TO_DISKLESS_SWITCH_PENDING: java.lang.Long, 10L: java.lang.Long, 0L: java.lang.Long, 10L: java.lang.Long),
+      Arguments.of(5L: java.lang.Long,  CLASSIC_TO_DISKLESS_SWITCH_PENDING: java.lang.Long, 10L: java.lang.Long, 0L: java.lang.Long, 10L: java.lang.Long),
       // No switch (-1): not realistic today but documents defensive behavior of the
       // post-restart path which seals unconditionally. Relevant if diskless topics
       // eventually use local logs as a read cache.
-      Arguments.of(0L: java.lang.Long,  NO_CLASSIC_TO_DISKLESS_START_OFFSET: java.lang.Long, 0L: java.lang.Long),
-      Arguments.of(5L: java.lang.Long,  NO_CLASSIC_TO_DISKLESS_START_OFFSET: java.lang.Long, 0L: java.lang.Long),
+      Arguments.of(0L: java.lang.Long,  NO_CLASSIC_TO_DISKLESS_START_OFFSET: java.lang.Long, 10L: java.lang.Long, 0L: java.lang.Long, 10L: java.lang.Long),
+      Arguments.of(5L: java.lang.Long,  NO_CLASSIC_TO_DISKLESS_START_OFFSET: java.lang.Long, 10L: java.lang.Long, 0L: java.lang.Long, 10L: java.lang.Long),
     )
   }