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


##########
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)
+    } else {
+      throw new MigrationConflictException(s"Conditional update on KRaft 
Migration ZNode failed. Sent zkVersion = " +

Review Comment:
   Nit: the old RuntimeException messages included the failed payload. Could we 
add the current version and data to this message too? It would make it easier 
to tell a lost reply from a real second writer when reading the logs.



-- 
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