(kyuubi) branch master updated: [KYUUBI #7645][SERVER] Fix operation state metrics transition

[wa...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wangzhigang1999 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new c24e02a65d [KYUUBI #7645][SERVER] Fix operation state metrics transition
c24e02a65d is described below

commit c24e02a65d73a411cb08fd69b6ca279b41fd269a
Author: Shuhang Han <14...@users.noreply.github.com>
AuthorDate: Tue Sep 8 20:56:03 2026 +0800

    [KYUUBI #7645][SERVER] Fix operation state metrics transition
    
    ### Why are the changes needed?
    
    `KyuubiOperation#setState` updated operation state metrics before delegating to `AbstractOperation#setState`, where the state transition is validated and the actual operation state is changed. It also did not protect the whole transition-and-metrics-update sequence with the operation lock.
    
    This can cause inaccurate operation state metrics in two cases:
    
    1. A stale or invalid state transition may update metrics before being rejected.
    
        For example, an operation may have already moved from `RUNNING` to `FINISHED`, while a delayed cancel/close path still tries to update it to another terminal state such as `CANCELED`. In the old implementation, the target state metric was updated before `OperationState#validateTransition` rejected the transition, so the operation remained in the original state but the metrics had already been changed.
    
    2. Concurrent state transitions for the same operation may update metrics based on the same stale old state.
    
        For example, one thread may observe the remote query as finished and call `KyuubiOperation#setState(FINISHED)`, while another thread concurrently handles cancel/close/timeout/error and calls `KyuubiOperation#setState(CANCELED)` or another terminal state. Since the old implementation read `state` and updated metrics outside a common operation lock, both paths could observe the old state as `RUNNING` and both update the `RUNNING` metric, making `kyuubi.operation.state.ExecuteStateme [...]
    
    This patch makes `KyuubiOperation#setState` execute under the operation lock, captures the old state once, delegates to `AbstractOperation#setState` first, and updates operation state metrics only after the state transition succeeds.
    
    ### How was this patch tested?
    
    Added a unit test in `KyuubiOperationSuite` to verify that a stale terminal transition does not update operation state metrics.
    ```
    build/mvn test -pl kyuubi-server -am \
      -Pspark-provided -Pflink-provided -Phive-provided \
      -Dtest=none \
      -DwildcardSuites=org.apache.kyuubi.operation.KyuubiOperationSuite
    ```
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Assisted-by: TraeCode with GPT-5
    
    Closes #7642 from HanShuhang/hsh/fix-state.
    
    Closes #7645
    
    063f24c84 [hanshuhang.lxy] [KYUUBI #7645][SERVER] Fix operation state metrics transition
    
    Lead-authored-by: Shuhang Han <14...@users.noreply.github.com>
    Co-authored-by: hanshuhang.lxy <ha...@bytedance.com>
    Signed-off-by: wangzhigang <ia...@gmail.com>
---
 .../apache/kyuubi/operation/KyuubiOperation.scala  |  11 ++-
 .../kyuubi/operation/KyuubiOperationSuite.scala    | 100 +++++++++++++++++++++
 2 files changed, 107 insertions(+), 4 deletions(-)

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
index 897c90c239..dfa822258f 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
@@ -213,15 +213,18 @@ abstract class KyuubiOperation(session: Session) extends AbstractOperation(sessi
 
   if (eventEnabled) EventBus.post(getOperationEvent)
 
-  override def setState(newState: OperationState): Unit = {
+  override def setState(newState: OperationState): Unit = withLockRequired {
+    val oldState = state
+    super.setState(newState)
     MetricsSystem.tracing { ms =>
-      if (!OperationState.isTerminal(state)) {
-        ms.markMeter(MetricRegistry.name(OPERATION_STATE, opType, state.toString.toLowerCase), -1)
+      if (!OperationState.isTerminal(oldState)) {
+        ms.markMeter(
+          MetricRegistry.name(OPERATION_STATE, opType, oldState.toString.toLowerCase),
+          -1)
       }
       ms.markMeter(MetricRegistry.name(OPERATION_STATE, opType, newState.toString.toLowerCase))
       ms.markMeter(MetricRegistry.name(OPERATION_STATE, newState.toString.toLowerCase))
     }
-    super.setState(newState)
     if (eventEnabled) EventBus.post(getOperationEvent)
   }
 
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationSuite.scala
new file mode 100644
index 0000000000..68f14c57e9
--- /dev/null
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationSuite.scala
@@ -0,0 +1,100 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.operation
+
+import com.codahale.metrics.MetricRegistry
+import org.mockito.Mockito.when
+import org.scalatestplus.mockito.MockitoSugar.mock
+
+import org.apache.kyuubi.{KyuubiFunSuite, KyuubiSQLException}
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.metrics.{MetricsConf, MetricsConstants, MetricsSystem}
+import org.apache.kyuubi.operation.OperationState.OperationState
+import org.apache.kyuubi.session.{Session, SessionHandle, SessionManager}
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V10
+
+class KyuubiOperationSuite extends KyuubiFunSuite {
+
+  private var metricsSystem: MetricsSystem = _
+
+  override def beforeEach(): Unit = {
+    super.beforeEach()
+    metricsSystem = new MetricsSystem()
+    metricsSystem.initialize(KyuubiConf()
+      .set(MetricsConf.METRICS_REPORTERS, Set.empty[String]))
+    metricsSystem.start()
+  }
+
+  override def afterEach(): Unit = {
+    try {
+      if (metricsSystem != null) {
+        metricsSystem.stop()
+        metricsSystem = null
+      }
+    } finally {
+      super.afterEach()
+    }
+  }
+
+  test("do not update operation state metrics on stale terminal state transition") {
+    val operation = new TestKyuubiOperation(mockSession())
+
+    operation.transitState(OperationState.RUNNING)
+    operation.transitState(OperationState.FINISHED)
+
+    val canceledMetric = stateMetric(OperationState.CANCELED)
+    val runningMetric = stateMetric(OperationState.RUNNING)
+    val canceledCount = MetricsSystem.meterValue(canceledMetric).getOrElse(0L)
+    val runningCount = MetricsSystem.meterValue(runningMetric).getOrElse(0L)
+
+    intercept[KyuubiSQLException] {
+      operation.transitState(OperationState.CANCELED)
+    }
+
+    assert(MetricsSystem.meterValue(canceledMetric).getOrElse(0L) === canceledCount)
+    assert(MetricsSystem.meterValue(runningMetric).getOrElse(0L) === runningCount)
+    assert(operation.getStatus.state === OperationState.FINISHED)
+  }
+
+  private def mockSession(): Session = {
+    val conf = KyuubiConf()
+    val sessionManager = mock[SessionManager]
+    when(sessionManager.getConf).thenReturn(conf)
+
+    val session = mock[Session]
+    when(session.protocol).thenReturn(HIVE_CLI_SERVICE_PROTOCOL_V10)
+    when(session.handle).thenReturn(SessionHandle())
+    when(session.user).thenReturn("kyuubi")
+    when(session.sessionManager).thenReturn(sessionManager)
+    session
+  }
+
+  private def stateMetric(state: OperationState): String = {
+    MetricRegistry.name(
+      MetricsConstants.OPERATION_STATE,
+      classOf[TestKyuubiOperation].getSimpleName,
+      state.toString.toLowerCase)
+  }
+}
+
+private class TestKyuubiOperation(session: Session) extends KyuubiOperation(session) {
+
+  def transitState(newState: OperationState): Unit = setState(newState)
+
+  override protected def runInternal(): Unit = {}
+}