Copilot commented on code in PR #7409:
URL: https://github.com/apache/kyuubi/pull/7409#discussion_r3108091399
##########
kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiBatchService.scala:
##########
@@ -60,57 +60,77 @@ class KyuubiBatchService(
}
override def start(): Unit = {
+ val UNINITIALIZED_BATCH_ID = "UNINITIALIZED_BATCH_ID"
assert(running.compareAndSet(false, true))
val submitTask: Runnable = () => {
restFrontend.waitForServerStarted()
while (running.get) {
- metadataManager.pickBatchForSubmitting(kyuubiInstance) match {
- case None => Thread.sleep(1000)
- case Some(metadata) =>
- val batchId = metadata.identifier
- info(s"$batchId is picked for submission.")
- val batchSession = sessionManager.createBatchSession(
- metadata.username,
- "anonymous",
- metadata.ipAddress,
- metadata.requestConf,
- metadata.engineType,
- Option(metadata.requestName),
- metadata.resource,
- metadata.className,
- metadata.requestArgs,
- Some(metadata),
- fromRecovery = false)
- sessionManager.openBatchSession(batchSession)
- var submitted = false
- while (!submitted) { // block until batch job submitted
- submitted = metadataManager.getBatchSessionMetadata(batchId)
match {
- case Some(metadata) if
OperationState.isTerminal(metadata.opState) =>
- true
- case Some(metadata) if metadata.opState ==
OperationState.RUNNING =>
- metadata.appState match {
- // app that is not submitted to resource manager
- case None | Some(ApplicationState.NOT_FOUND) => false
- // app that is pending in resource manager while the local
startup
- // process is alive. For example, in Spark YARN cluster
mode, if set
- // spark.yarn.submit.waitAppCompletion=false, the local
spark-submit
- // process exits immediately once Application goes
ACCEPTED status,
- // even no resource could be allocated for the AM
container.
- case Some(ApplicationState.PENDING) if
batchSession.startupProcessAlive =>
- false
- // not sure, added for safe
- case Some(ApplicationState.UNKNOWN) => false
- case _ => true
- }
- case Some(_) =>
- false
- case None =>
- error(s"$batchId does not existed in metastore, assume it is
finished")
- true
+ var batchId = UNINITIALIZED_BATCH_ID
+ try {
+ metadataManager.pickBatchForSubmitting(kyuubiInstance) match {
+ case None => Thread.sleep(1000)
+ case Some(metadata) =>
+ batchId = metadata.identifier
+ info(s"$batchId is picked for submission.")
+ val batchSession = sessionManager.createBatchSession(
+ metadata.username,
+ "anonymous",
+ metadata.ipAddress,
+ metadata.requestConf,
+ metadata.engineType,
+ Option(metadata.requestName),
+ metadata.resource,
+ metadata.className,
+ metadata.requestArgs,
+ Some(metadata),
+ fromRecovery = false)
+ sessionManager.openBatchSession(batchSession)
+ var submitted = false
+ while (!submitted) { // block until batch job submitted
+ submitted = metadataManager.getBatchSessionMetadata(batchId)
match {
+ case Some(metadata) if
OperationState.isTerminal(metadata.opState) =>
+ true
+ case Some(metadata) if metadata.opState ==
OperationState.RUNNING =>
+ metadata.appState match {
+ // app that is not submitted to resource manager
+ case None | Some(ApplicationState.NOT_FOUND) => false
+ // app that is pending in resource manager while the
local startup
+ // process is alive. For example, in Spark YARN cluster
mode, if set
+ // spark.yarn.submit.waitAppCompletion=false, the local
spark-submit
+ // process exits immediately once Application goes
ACCEPTED status,
+ // even no resource could be allocated for the AM
container.
+ case Some(ApplicationState.PENDING) if
batchSession.startupProcessAlive =>
+ false
+ // not sure, added for safe
+ case Some(ApplicationState.UNKNOWN) => false
+ case _ => true
+ }
+ case Some(_) =>
+ false
+ case None =>
+ error(s"$batchId does not existed in metastore, assume it
is finished")
+ true
+ }
+ if (!submitted) Thread.sleep(1000)
+ }
+ info(s"$batchId is submitted or finished.")
+ }
+ } catch {
+ // If the batch session failed to open, reinitialize the batch state
to ERROR
+ // This can be due to a DB error or batch_connection_limits exceeded
+ case e: Exception =>
+ if (batchId == UNINITIALIZED_BATCH_ID) {
+ error(s"Error picking batch for submission", e)
+ } else {
+ error(s"Error opening batch session for $batchId", e)
+ try {
+ metadataManager.failScheduledBatch(batchId)
+ } catch {
+ case ex: Exception =>
+ error(s"Unable to modify metadata for $batchId to ERROR", ex)
}
- if (!submitted) Thread.sleep(1000)
}
- info(s"$batchId is submitted or finished.")
+ Thread.sleep(1000)
}
Review Comment:
The catch-all `case e: Exception` will also swallow `InterruptedException`
(including interrupts used during shutdown) and then calls `Thread.sleep`,
which can delay/impede service stop. Handle `InterruptedException` separately
(restore interrupt flag and exit the loop/task) and prefer
`scala.util.control.NonFatal` for the general case instead of catching all
`Exception`.
##########
kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala:
##########
@@ -72,6 +73,72 @@ class BatchesV2ResourceSuite extends
BatchesResourceSuiteBase {
Utils.tryLogNonFatalError { sessionManager.closeSession(session.handle) }
}
}
+
+ test("KyuubiBatchService catch block when openBatchSession fails during
metadata update") {
+ val sessionManager =
fe.be.sessionManager.asInstanceOf[KyuubiSessionManager]
+ val realMetadataManager = sessionManager.metadataManager.get
+
+ val wrapperMetadataManager = new MetadataManager {
+ override def updateMetadata(metadata: Metadata, asyncRetryOnError:
Boolean = true): Unit = {
+ throw new RuntimeException("test metadata update failure")
+ }
+ override def insertMetadata(metadata: Metadata, asyncRetryOnError:
Boolean = true): Unit = {
+ realMetadataManager.insertMetadata(metadata, asyncRetryOnError)
+ }
+ override def getBatch(batchId: String) =
realMetadataManager.getBatch(batchId)
+ }
+ wrapperMetadataManager.initialize(sessionManager.getConf)
+
+ val originalMetadataManager = sessionManager.metadataManager
+ try {
+ sessionManager.metadataManager = Some(wrapperMetadataManager)
+
+ val requestObj = newSparkBatchRequest(Map("spark.master" -> "local"))
+ val response = webTarget.path("api/v1/batches")
+ .request(MediaType.APPLICATION_JSON_TYPE)
+ .header(AUTHORIZATION_HEADER, basicAuthorizationHeader("anonymous"))
+ .post(Entity.entity(requestObj, MediaType.APPLICATION_JSON_TYPE))
+ assert(response.getStatus == 200)
+ val batch = response.readEntity(classOf[Batch])
+ val batchId = batch.getId
+ assert(batch.getState === "INITIALIZED")
+
+ eventually(timeout(15.seconds), interval(1.second)) {
+ val batchInfoResponse = webTarget.path(s"api/v1/batches/$batchId")
+ .request(MediaType.APPLICATION_JSON_TYPE)
+ .header(AUTHORIZATION_HEADER, basicAuthorizationHeader("anonymous"))
+ .get()
+ assert(batchInfoResponse.getStatus == 200)
+ val batchInfo = batchInfoResponse.readEntity(classOf[Batch])
+ assert(
+ batchInfo.getState === "INITIALIZED",
+ "Batch should remain INITIALIZED after openBatchSession failed and
catch block ran " +
+ "(failScheduledBatch only transitions PENDING->ERROR)")
Review Comment:
The assertion that the batch stays `INITIALIZED` looks inconsistent with the
v2 flow: `pickMetadata` transitions `INITIALIZED -> PENDING` when the submitter
picks it, and the new catch path calls `failScheduledBatch` to transition
`PENDING -> ERROR`. This test should validate the batch eventually becomes
`ERROR` (and, importantly, is not left stuck in `PENDING`) rather than
expecting it to remain `INITIALIZED`.
##########
kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala:
##########
@@ -72,6 +73,72 @@ class BatchesV2ResourceSuite extends
BatchesResourceSuiteBase {
Utils.tryLogNonFatalError { sessionManager.closeSession(session.handle) }
}
}
+
+ test("KyuubiBatchService catch block when openBatchSession fails during
metadata update") {
+ val sessionManager =
fe.be.sessionManager.asInstanceOf[KyuubiSessionManager]
+ val realMetadataManager = sessionManager.metadataManager.get
+
+ val wrapperMetadataManager = new MetadataManager {
+ override def updateMetadata(metadata: Metadata, asyncRetryOnError:
Boolean = true): Unit = {
+ throw new RuntimeException("test metadata update failure")
+ }
+ override def insertMetadata(metadata: Metadata, asyncRetryOnError:
Boolean = true): Unit = {
+ realMetadataManager.insertMetadata(metadata, asyncRetryOnError)
+ }
+ override def getBatch(batchId: String) =
realMetadataManager.getBatch(batchId)
+ }
+ wrapperMetadataManager.initialize(sessionManager.getConf)
Review Comment:
`wrapperMetadataManager.initialize(...)` will create a new
`JDBCMetadataStore` (and `HikariDataSource`) via `MetadataManager.initialize`,
but the wrapper service is never `stop()`'d/closed in this test. This can leak
connection pools/threads and cause flakiness across the test suite. Either
avoid calling `initialize` (the overridden methods don’t need it) or call
`wrapperMetadataManager.stop()` (or otherwise close resources) in `finally`.
##########
kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiBatchService.scala:
##########
@@ -60,57 +60,77 @@ class KyuubiBatchService(
}
override def start(): Unit = {
+ val UNINITIALIZED_BATCH_ID = "UNINITIALIZED_BATCH_ID"
assert(running.compareAndSet(false, true))
val submitTask: Runnable = () => {
restFrontend.waitForServerStarted()
while (running.get) {
- metadataManager.pickBatchForSubmitting(kyuubiInstance) match {
- case None => Thread.sleep(1000)
- case Some(metadata) =>
- val batchId = metadata.identifier
- info(s"$batchId is picked for submission.")
- val batchSession = sessionManager.createBatchSession(
- metadata.username,
- "anonymous",
- metadata.ipAddress,
- metadata.requestConf,
- metadata.engineType,
- Option(metadata.requestName),
- metadata.resource,
- metadata.className,
- metadata.requestArgs,
- Some(metadata),
- fromRecovery = false)
- sessionManager.openBatchSession(batchSession)
- var submitted = false
- while (!submitted) { // block until batch job submitted
- submitted = metadataManager.getBatchSessionMetadata(batchId)
match {
- case Some(metadata) if
OperationState.isTerminal(metadata.opState) =>
- true
- case Some(metadata) if metadata.opState ==
OperationState.RUNNING =>
- metadata.appState match {
- // app that is not submitted to resource manager
- case None | Some(ApplicationState.NOT_FOUND) => false
- // app that is pending in resource manager while the local
startup
- // process is alive. For example, in Spark YARN cluster
mode, if set
- // spark.yarn.submit.waitAppCompletion=false, the local
spark-submit
- // process exits immediately once Application goes
ACCEPTED status,
- // even no resource could be allocated for the AM
container.
- case Some(ApplicationState.PENDING) if
batchSession.startupProcessAlive =>
- false
- // not sure, added for safe
- case Some(ApplicationState.UNKNOWN) => false
- case _ => true
- }
- case Some(_) =>
- false
- case None =>
- error(s"$batchId does not existed in metastore, assume it is
finished")
- true
+ var batchId = UNINITIALIZED_BATCH_ID
+ try {
+ metadataManager.pickBatchForSubmitting(kyuubiInstance) match {
+ case None => Thread.sleep(1000)
+ case Some(metadata) =>
+ batchId = metadata.identifier
+ info(s"$batchId is picked for submission.")
+ val batchSession = sessionManager.createBatchSession(
+ metadata.username,
+ "anonymous",
+ metadata.ipAddress,
+ metadata.requestConf,
+ metadata.engineType,
+ Option(metadata.requestName),
+ metadata.resource,
+ metadata.className,
+ metadata.requestArgs,
+ Some(metadata),
+ fromRecovery = false)
+ sessionManager.openBatchSession(batchSession)
+ var submitted = false
+ while (!submitted) { // block until batch job submitted
+ submitted = metadataManager.getBatchSessionMetadata(batchId)
match {
+ case Some(metadata) if
OperationState.isTerminal(metadata.opState) =>
+ true
+ case Some(metadata) if metadata.opState ==
OperationState.RUNNING =>
+ metadata.appState match {
+ // app that is not submitted to resource manager
+ case None | Some(ApplicationState.NOT_FOUND) => false
+ // app that is pending in resource manager while the
local startup
+ // process is alive. For example, in Spark YARN cluster
mode, if set
+ // spark.yarn.submit.waitAppCompletion=false, the local
spark-submit
+ // process exits immediately once Application goes
ACCEPTED status,
+ // even no resource could be allocated for the AM
container.
+ case Some(ApplicationState.PENDING) if
batchSession.startupProcessAlive =>
+ false
+ // not sure, added for safe
+ case Some(ApplicationState.UNKNOWN) => false
+ case _ => true
+ }
+ case Some(_) =>
+ false
+ case None =>
+ error(s"$batchId does not existed in metastore, assume it
is finished")
Review Comment:
Typo/grammar in the log message: "does not existed" should be "does not
exist" for clarity when diagnosing missing metadata records.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
