OmniaGM commented on code in PR #23639:
URL: https://github.com/apache/kafka/pull/23639#discussion_r4455727926


##########
core/src/main/scala/kafka/zk/KafkaZkClient.scala:
##########
@@ -1764,11 +1764,46 @@ class KafkaZkClient private[zk] (
     initialState.withMigrationZkVersion(0)
   }
 
+  /**
+   * When there is a ConnectionLossException while updating the /migration 
znode, the update may have actually succeeded
+   * on the server before the connection dropped, so the retried request can 
fail with BadVersionException even though
+   * our own write went through. This checker reads back the current data and, 
if its migration-relevant fields match
+   * what we intended to write, treats it as our own write having succeeded, 
returning the real current version.
+   */
+  private def checkMigrationStateZkData(
+                                         zkClient: KafkaZkClient,
+                                         path: String,
+                                         expectedData: Array[Byte]
+                                       ): (Boolean, Int) = {
+    try {
+      val (writtenDataOpt, writtenStat) = zkClient.getDataAndStat(path)
+      val expected = MigrationZNode.decode(expectedData, 
writtenStat.getVersion, -1)
+      val succeeded = writtenDataOpt.exists { writtenData =>
+        val written = MigrationZNode.decode(writtenData, 
writtenStat.getVersion, -1)
+        expected.kraftControllerId() == written.kraftControllerId() &&
+          expected.kraftControllerEpoch() == written.kraftControllerEpoch() &&
+          expected.kraftMetadataOffset() == written.kraftMetadataOffset() &&
+          expected.kraftMetadataEpoch() == written.kraftMetadataEpoch()
+      }
+      if (succeeded) (true, writtenStat.getVersion) else (false, 
ZkVersion.UnknownVersion)
+    } catch {
+      case _: Exception => (false, ZkVersion.UnknownVersion)
+    }
+  }
+
   def updateMigrationState(migrationState: ZkMigrationLeadershipState): 
ZkMigrationLeadershipState = {
-    val req = SetDataRequest(MigrationZNode.path, 
MigrationZNode.encode(migrationState), migrationState.migrationZkVersion())
-    val resp = retryRequestUntilConnected(req)
-    resp.maybeThrow()
-    migrationState.withMigrationZkVersion(resp.stat.getVersion)
+    val path = MigrationZNode.path
+    val data = MigrationZNode.encode(migrationState)
+    val expectedVersion = migrationState.migrationZkVersion()
+    val (succeeded, newVersion) = conditionalUpdatePath(
+      path, data, expectedVersion, Some(checkMigrationStateZkData))
+
+    if (succeeded && (expectedVersion == ZkVersion.MatchAnyVersion || 
expectedVersion + 1 == newVersion)) {
+      migrationState.withMigrationZkVersion(newVersion)

Review Comment:
   Just to clarify, was switching to `conditionalUpdatePath` here intentional? 
I think it changes what happens when the `/migration` znode is missing. Before, 
`resp.maybeThrow()` threw `NoNodeException`, which the driver retried. Now 
`conditionalUpdatePath` returns `(false, ZkVersion.UnknownVersion)` for 
`Code.NONODE`, so we throw `MigrationConflictException` below and the 
controller renounces. The message, `Conditional update on KRaft Migration ZNode 
failed. Sent zkVersion = …`, sounds misleading in that case.
   
   I'm also not sure what happens next. After the renounce, the new leader's 
`KRaftLeaderEvent` calls `getOrCreateMigrationRecoveryState`, which recreates 
`/migration` from that node's in-memory state. Is that safe, or should `NONODE` 
keep its old behaviour here? WDYT?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to