From 063f24c84bbd12094d42428351ccdfe573d6f322 Mon Sep 17 00:00:00 2001
From: "hanshuhang.lxy" <hanshuhang.lxy@bytedance.com>
Date: Fri, 21 Aug 2026 10:35:28 +0800
Subject: [PATCH] [KYUUBI #7645][SERVER] Fix operation state metrics transition

---
 .../kyuubi/operation/KyuubiOperation.scala    |  11 +-
 .../operation/KyuubiOperationSuite.scala      | 100 ++++++++++++++++++
 2 files changed, 107 insertions(+), 4 deletions(-)
 create mode 100644 kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationSuite.scala

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 897c90c2396..dfa822258fa 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 00000000000..68f14c57e9b
--- /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 = {}
+}