diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala index c5182979ee7..914668af0ea 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala @@ -337,16 +337,22 @@ class MetadataManager extends AbstractService("MetadataManager") { info(s"Retrying metadata requests for $id") var request = ref.metadataRequests.peek() while (request != null) { - request match { - case insert: InsertMetadata => - insertMetadata(insert.metadata, asyncRetryOnError = false) - case update: UpdateMetadata => - updateMetadata(update.metadata, asyncRetryOnError = false) + try { + request match { + case insert: InsertMetadata => + insertMetadata(insert.metadata, asyncRetryOnError = false) + case update: UpdateMetadata => + updateMetadata(update.metadata, asyncRetryOnError = false) + } + removeRetryRequest(ref, request) + } catch { + case e: MetadataUpdatePostconditionException => + error( + s"Discarding metadata update retry for $id after " + + "its postcondition check failed", + e) + removeRetryRequest(ref, request) } - ref.metadataRequests.remove(request) - MetricsSystem.tracing(_.markMeter( - MetricsConstants.METADATA_REQUEST_RETRYING, - -1L)) request = ref.metadataRequests.peek() } } catch { @@ -379,6 +385,13 @@ class MetadataManager extends AbstractService("MetadataManager") { requestsRetryInterval, TimeUnit.MILLISECONDS) } + + private def removeRetryRequest( + ref: MetadataRequestsRetryRef, + request: MetadataRequest): Unit = { + ref.metadataRequests.remove(request) + MetricsSystem.tracing(_.markMeter(MetricsConstants.METADATA_REQUEST_RETRYING, -1L)) + } } object MetadataManager extends Logging { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala index 492b286e3b4..4fb13280f43 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala @@ -19,6 +19,7 @@ package org.apache.kyuubi.server.metadata import java.io.Closeable +import org.apache.kyuubi.KyuubiException import org.apache.kyuubi.server.metadata.api.{KubernetesEngineInfo, Metadata, MetadataFilter} trait MetadataStore extends Closeable { @@ -118,3 +119,17 @@ trait MetadataStore extends Closeable { */ def cleanupKubernetesEngineInfoByAge(maxAge: Long, limit: Int): Int } + +sealed abstract private[metadata] class MetadataUpdatePostconditionException(message: String) + extends KyuubiException(message) + +private[metadata] class MetadataRowNotFoundException(identifier: String) + extends MetadataUpdatePostconditionException( + s"Metadata row $identifier was not found after an update returned 0") + +private[metadata] class MetadataUpdateMismatchException( + identifier: String, + val mismatchedColumns: Seq[String]) + extends MetadataUpdatePostconditionException( + s"Metadata row $identifier does not match the requested update for columns: " + + mismatchedColumns.mkString(", ")) diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala index a9ee8fa9569..f13687bdff1 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala @@ -34,7 +34,7 @@ import com.zaxxer.hikari.{HikariConfig, HikariDataSource} import org.apache.kyuubi.{KyuubiException, Logging, Utils} import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.operation.OperationState -import org.apache.kyuubi.server.metadata.MetadataStore +import org.apache.kyuubi.server.metadata.{MetadataRowNotFoundException, MetadataStore, MetadataUpdateMismatchException} import org.apache.kyuubi.server.metadata.api.{KubernetesEngineInfo, Metadata, MetadataFilter} import org.apache.kyuubi.server.metadata.jdbc.DatabaseType._ import org.apache.kyuubi.server.metadata.jdbc.JDBCMetadataStoreConf._ @@ -332,76 +332,97 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging { override def updateMetadata(metadata: Metadata): Unit = { val queryBuilder = new StringBuilder - val params = ListBuffer[Any]() + val updateFields = ListBuffer[(String, Any)]() queryBuilder.append(s"UPDATE $METADATA_TABLE") - val setClauses = ListBuffer[String]() Option(metadata.kyuubiInstance).foreach { _ => - setClauses += "kyuubi_instance = ?" - params += metadata.kyuubiInstance + updateFields += (("kyuubi_instance", metadata.kyuubiInstance)) } Option(metadata.state).foreach { _ => - setClauses += "state = ?" - params += metadata.state + updateFields += (("state", metadata.state)) } Option(metadata.requestConf).filter(_.nonEmpty).foreach { _ => - setClauses += "request_conf =?" - params += valueAsString(metadata.requestConf) + updateFields += (("request_conf", valueAsString(metadata.requestConf))) } metadata.clusterManager.foreach { cm => - setClauses += "cluster_manager = ?" - params += cm + updateFields += (("cluster_manager", cm)) } if (metadata.endTime > 0) { - setClauses += "end_time = ?" - params += metadata.endTime + updateFields += (("end_time", metadata.endTime)) } if (metadata.engineOpenTime > 0) { - setClauses += "engine_open_time = ?" - params += metadata.engineOpenTime + updateFields += (("engine_open_time", metadata.engineOpenTime)) } Option(metadata.engineId).foreach { _ => - setClauses += "engine_id = ?" - params += metadata.engineId + updateFields += (("engine_id", metadata.engineId)) } Option(metadata.engineName).foreach { _ => - setClauses += "engine_name = ?" - params += metadata.engineName + updateFields += (("engine_name", metadata.engineName)) } Option(metadata.engineUrl).foreach { _ => - setClauses += "engine_url = ?" - params += metadata.engineUrl + updateFields += (("engine_url", metadata.engineUrl)) } Option(metadata.engineState).foreach { _ => - setClauses += "engine_state = ?" - params += metadata.engineState + updateFields += (("engine_state", metadata.engineState)) } metadata.engineError.foreach { error => - setClauses += "engine_error = ?" - params += error + updateFields += (("engine_error", error)) } if (metadata.peerInstanceClosed) { - setClauses += "peer_instance_closed = ?" - params += metadata.peerInstanceClosed + updateFields += (("peer_instance_closed", metadata.peerInstanceClosed)) } - if (setClauses.nonEmpty) { - queryBuilder.append(setClauses.mkString(" SET ", ", ", "")) + if (updateFields.nonEmpty) { + queryBuilder.append(updateFields.map { case (column, _) => + s"$column = ?" + }.mkString(" SET ", ", ", "")) } queryBuilder.append(" WHERE identifier = ?") - params += metadata.identifier + val params = updateFields.map(_._2) :+ metadata.identifier val query = queryBuilder.toString() JdbcUtils.withConnection { connection => withUpdateCount(connection, query, params.toSeq: _*) { updateCount => if (updateCount == 0) { - throw new KyuubiException( - s"Error updating metadata for ${metadata.identifier} by SQL: $query, " + - s"with params: ${params.mkString(", ")}") + verifyMetadataUpdatePostcondition(connection, metadata.identifier, updateFields.toSeq) } } } } + private def verifyMetadataUpdatePostcondition( + connection: Connection, + identifier: String, + updateFields: Seq[(String, Any)]): Unit = { + val query = s"SELECT ${updateFields.map(_._1).mkString(", ")} " + + s"FROM $METADATA_TABLE WHERE identifier = ?" + withResultSet(connection, query, identifier) { resultSet => + if (!resultSet.next()) { + throw new MetadataRowNotFoundException(identifier) + } + + val mismatchedColumns = updateFields.zipWithIndex.collect { + case ((column, expected), index) + if getResultSetValue(resultSet, index + 1, expected) != expected => column + } + if (mismatchedColumns.nonEmpty) { + throw new MetadataUpdateMismatchException(identifier, mismatchedColumns) + } + } + } + + private def getResultSetValue(resultSet: ResultSet, index: Int, expected: Any): Any = { + val value = expected match { + case _: String => resultSet.getString(index) + case _: Int => resultSet.getInt(index) + case _: Long => resultSet.getLong(index) + case _: Double => resultSet.getDouble(index) + case _: Float => resultSet.getFloat(index) + case _: Boolean => resultSet.getBoolean(index) + case _ => resultSet.getObject(index) + } + if (resultSet.wasNull()) null else value + } + override def cleanupMetadataByIdentifier(identifier: String): Unit = { val query = s"DELETE FROM $METADATA_TABLE WHERE identifier = ?" JdbcUtils.withConnection { connection => diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/MetadataManagerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/MetadataManagerSuite.scala index de2e651f44a..65dd5bd83bf 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/MetadataManagerSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/MetadataManagerSuite.scala @@ -97,6 +97,22 @@ class MetadataManagerSuite extends KyuubiFunSuite { } } + test("stop retrying an update when its row remains missing") { + withMetadataManager(Map( + METADATA_REQUEST_ASYNC_RETRY_ENABLED.key -> "true", + METADATA_REQUEST_RETRY_INTERVAL.key -> "100")) { metadataManager => + val metadata = newMetadata() + metadataManager.updateMetadata(metadata) + val retryRef = metadataManager.getMetadataRequestsRetryRef(metadata.identifier) + + assert(retryRef != null) + eventually(timeout(3.seconds)) { + assert(!retryRef.hasRemainingRequests()) + assert(metadataManager.getMetadataRequestsRetryRef(metadata.identifier) == null) + } + } + } + test("async metadata request metrics") { withMetadataManager(Map( METADATA_REQUEST_ASYNC_RETRY_ENABLED.key -> "true", diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala index 5cbc13baa05..e88d009dd3d 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala @@ -22,10 +22,10 @@ import java.util.UUID import org.scalatest.concurrent.PatienceConfiguration.Timeout import org.scalatest.time.SpanSugar._ -import org.apache.kyuubi.{KyuubiException, KyuubiFunSuite} +import org.apache.kyuubi.KyuubiFunSuite import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.engine.ApplicationState -import org.apache.kyuubi.server.metadata.MetadataManager +import org.apache.kyuubi.server.metadata.{MetadataManager, MetadataRowNotFoundException, MetadataUpdateMismatchException} import org.apache.kyuubi.server.metadata.api.{KubernetesEngineInfo, Metadata, MetadataFilter} import org.apache.kyuubi.server.metadata.jdbc.JDBCMetadataStoreConf._ import org.apache.kyuubi.session.SessionType @@ -274,13 +274,64 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite { jdbcMetadataStore.cleanupMetadataByIdentifier(batchId) } - test("throw exception if update count is 0") { + test("report a missing metadata row after update count is 0") { val metadata = Metadata(identifier = UUID.randomUUID().toString, state = "RUNNING") - intercept[KyuubiException] { + intercept[MetadataRowNotFoundException] { jdbcMetadataStore.updateMetadata(metadata) } } + test("accept update count 0 if the metadata already matches") { + val batchId = UUID.randomUUID().toString + val triggerName = s"ignore_metadata_update_${batchId.replace("-", "")}" + val metadata = Metadata( + identifier = batchId, + sessionType = SessionType.BATCH, + realUser = "kyuubi", + username = "kyuubi", + state = "RUNNING", + createTime = System.currentTimeMillis(), + engineType = "spark") + val metadataToUpdate = Metadata( + identifier = batchId, + kyuubiInstance = "localhost:10099", + state = "RUNNING", + requestConf = Map("spark.master" -> "local"), + clusterManager = Some("kubernetes"), + engineOpenTime = System.currentTimeMillis(), + engineId = "app_id", + engineName = "app_name", + engineUrl = "app_url", + engineState = "FAILED", + engineError = Some("engine_error"), + endTime = System.currentTimeMillis(), + peerInstanceClosed = true) + + try { + jdbcMetadataStore.insertMetadata(metadata) + jdbcMetadataStore.updateMetadata(metadataToUpdate) + executeSql( + s""" + |CREATE TRIGGER $triggerName + |BEFORE UPDATE ON metadata + |WHEN OLD.identifier = '$batchId' + |BEGIN + | SELECT RAISE(IGNORE); + |END + |""".stripMargin) + + jdbcMetadataStore.updateMetadata(metadataToUpdate) + + val error = intercept[MetadataUpdateMismatchException] { + jdbcMetadataStore.updateMetadata(metadataToUpdate.copy(state = "FINISHED")) + } + assert(error.mismatchedColumns === Seq("state")) + } finally { + executeSql(s"DROP TRIGGER IF EXISTS $triggerName") + jdbcMetadataStore.cleanupMetadataByIdentifier(batchId) + } + } + test("get schema urls with correct version ordering") { val url1 = "metadata-store-schema-1.7.0.mysql.sql" val url2 = "metadata-store-schema-1.7.1.mysql.sql" @@ -363,4 +414,18 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite { jdbcMetadataStore.cleanupKubernetesEngineInfoByIdentifier(tag) assert(jdbcMetadataStore.getKubernetesMetaEngineInfo(tag) == null) } + + private def executeSql(sql: String): Unit = { + val connection = jdbcMetadataStore.hikariDataSource.getConnection + try { + val statement = connection.createStatement() + try { + statement.execute(sql) + } finally { + statement.close() + } + } finally { + connection.close() + } + } }