[LIVY-1070] Add LivySessionMetrics Codahale gauges for session monitoring

server/src/main/scala/org/apache/livy/server/LivyServer.scala

@@ -252,6 +252,8 @@ class LivyServer extends Logging {
             }
 
             context.mountMetricsAdminServlet("/metrics")
+            LivySessionMetrics.register(
+              metricRegistry, interactiveSessionManager, batchSessionManager)
 
             mount(context, livyVersionServlet, "/version/*")
           } catch {

server/src/main/scala/org/apache/livy/server/LivySessionMetrics.scala

@@ -0,0 +1,142 @@
+/*
+ * 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.livy.server
+
+import com.codahale.metrics.{Gauge, MetricRegistry}
+
+import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager}
+
+/**
+ * Registers Livy session metrics as Codahale/Dropwizard gauges in the MetricRegistry.
+ *
+ * Uses the existing `io.dropwizard.metrics:metrics-core` dependency (Java package
+ * `com.codahale.metrics`) already declared in `server/pom.xml` and wired through
+ * `scalatra-metrics` (`MetricsBootstrap`) in `LivyServer`. No new metrics library
+ * is introduced by this class; it only adds session gauges to the shared registry
+ * exposed at `/metrics`.
+ *
+ * 18 gauges total:
+ *   - 3 overall (total, active_total, terminal_total)
+ *   - 8 interactive (total, idle, busy, starting, shutting_down, dead, error, killed)
+ *   - 7 batch (total, starting, running, success, dead, error, killed)
+ */
+object LivySessionMetrics {
+
+  /**
+   * Registers session count gauges into `metricRegistry`. Idempotent: existing gauge names
+   * are left unchanged. Gauge callbacks capture the session managers passed here, so no
+   * long-lived helper instance is required.
+   */
+  def register(
+      metricRegistry: MetricRegistry,
+      interactiveSessionManager: InteractiveSessionManager,
+      batchSessionManager: BatchSessionManager): Unit = {
+
+    def registerGauge(name: String, valueFn: => Int): Unit = {
+      if (!metricRegistry.getGauges.containsKey(name)) {
+        metricRegistry.register(name, new Gauge[Int] {
+          override def getValue: Int = {
+            try {
+              valueFn
+            } catch {
+              case _: Throwable => 0
+            }
+          }
+        })
+      }
+    }
+
+    def countInteractiveByState(expectedStates: Set[String]): Int = {
+      interactiveSessionManager.all().count { session =>
+        expectedStates.contains(normalizeState(session.state.toString))
+      }
+    }
+
+    def countBatchByState(expectedStates: Set[String]): Int = {
+      batchSessionManager.all().count { session =>
+        expectedStates.contains(normalizeState(session.state.toString))
+      }
+    }
+
+    def normalizeState(state: String): String = {
+      Option(state).getOrElse("").trim.toLowerCase.replace("-", "_").replace(" ", "_")
+    }
+
+    // ========== Interactive Session Metrics (8) ==========
+
+    registerGauge("livy.sessions.interactive.total",
+      interactiveSessionManager.size)
+
+    registerGauge("livy.sessions.interactive.idle",
+      countInteractiveByState(Set("idle")))
+
+    registerGauge("livy.sessions.interactive.starting",
+      countInteractiveByState(Set("starting")))
+
+    registerGauge("livy.sessions.interactive.busy",
+      countInteractiveByState(Set("busy")))
+
+    registerGauge("livy.sessions.interactive.dead",
+      countInteractiveByState(Set("dead")))
+
+    registerGauge("livy.sessions.interactive.shutting_down",
+      countInteractiveByState(Set("shutting_down", "shuttingdown")))
+
+    registerGauge("livy.sessions.interactive.error",
+      countInteractiveByState(Set("error")))
+
+    registerGauge("livy.sessions.interactive.killed",
+      countInteractiveByState(Set("killed")))
+
+    // ========== Batch Session Metrics (7) ==========
+
+    registerGauge("livy.sessions.batch.total",
+      batchSessionManager.size)
+
+    registerGauge("livy.sessions.batch.starting",
+      countBatchByState(Set("starting")))
+
+    registerGauge("livy.sessions.batch.running",
+      countBatchByState(Set("running")))
+
+    registerGauge("livy.sessions.batch.success",
+      countBatchByState(Set("success", "succeeded")))
+
+    registerGauge("livy.sessions.batch.dead",
+      countBatchByState(Set("dead")))
+
+    registerGauge("livy.sessions.batch.error",
+      countBatchByState(Set("error")))
+
+    registerGauge("livy.sessions.batch.killed",
+      countBatchByState(Set("killed")))
+
+    // ========== Overall Metrics (3) ==========
+
+    registerGauge("livy.sessions.total",
+      interactiveSessionManager.size + batchSessionManager.size)
+
+    registerGauge("livy.sessions.active.total",
+      countInteractiveByState(Set("starting", "idle", "busy")) +
+        countBatchByState(Set("starting", "running")))
+
+    registerGauge("livy.sessions.terminal.total",
+      countInteractiveByState(Set("dead", "error", "killed")) +
+        countBatchByState(Set("success", "succeeded", "dead", "error", "killed")))
+  }
+}

server/src/test/scala/org/apache/livy/server/LivySessionMetricsSpec.scala

@@ -0,0 +1,210 @@
+/*
+ * 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.livy.server
+
+import com.codahale.metrics.{Gauge, MetricRegistry}
+import org.mockito.Mockito._
+import org.scalatest.{FunSpec, Matchers}
+import org.scalatestplus.mockito.MockitoSugar
+
+import org.apache.livy.LivyBaseUnitTestSuite
+import org.apache.livy.server.batch.BatchSession
+import org.apache.livy.server.interactive.InteractiveSession
+import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager, SessionState}
+
+class LivySessionMetricsSpec extends FunSpec with Matchers with MockitoSugar
+  with LivyBaseUnitTestSuite {
+
+  private val expectedGauges = Seq(
+    "livy.sessions.interactive.total",
+    "livy.sessions.interactive.idle",
+    "livy.sessions.interactive.starting",
+    "livy.sessions.interactive.busy",
+    "livy.sessions.interactive.dead",
+    "livy.sessions.interactive.shutting_down",
+    "livy.sessions.interactive.error",
+    "livy.sessions.interactive.killed",
+    "livy.sessions.batch.total",
+    "livy.sessions.batch.starting",
+    "livy.sessions.batch.running",
+    "livy.sessions.batch.success",
+    "livy.sessions.batch.dead",
+    "livy.sessions.batch.error",
+    "livy.sessions.batch.killed",
+    "livy.sessions.total",
+    "livy.sessions.active.total",
+    "livy.sessions.terminal.total"
+  )
+
+  private def gaugeValue(registry: MetricRegistry, name: String): Int = {
+    registry.getGauges.get(name).asInstanceOf[Gauge[Int]].getValue
+  }
+
+  private def mockInteractive(state: SessionState): InteractiveSession = {
+    val session = mock[InteractiveSession]
+    when(session.state).thenReturn(state)
+    session
+  }
+
+  private def mockBatch(state: SessionState): BatchSession = {
+    val session = mock[BatchSession]
+    when(session.state).thenReturn(state)
+    session
+  }
+
+  private def registerMetrics(
+      interactiveSessions: Seq[InteractiveSession] = Seq.empty,
+      batchSessions: Seq[BatchSession] = Seq.empty,
+      registry: MetricRegistry = new MetricRegistry()): Unit = {
+    val interactiveManager = mock[InteractiveSessionManager]
+    val batchManager = mock[BatchSessionManager]
+    when(interactiveManager.size).thenReturn(interactiveSessions.size)
+    when(interactiveManager.all()).thenReturn(interactiveSessions)
+    when(batchManager.size).thenReturn(batchSessions.size)
+    when(batchManager.all()).thenReturn(batchSessions)
+    LivySessionMetrics.register(registry, interactiveManager, batchManager)
+  }
+
+  describe("LivySessionMetrics") {
+
+    it("should register all 18 session gauges") {
+      val registry = new MetricRegistry()
+      registerMetrics(registry = registry)
+
+      expectedGauges.foreach { name =>
+        registry.getGauges.keySet() should contain(name)
+      }
+      registry.getGauges.size() shouldBe 18
+    }
+
+    it("should not register duplicate gauges on the same registry") {
+      val registry = new MetricRegistry()
+      val interactiveManager = mock[InteractiveSessionManager]
+      val batchManager = mock[BatchSessionManager]
+      when(interactiveManager.size).thenReturn(0)
+      when(interactiveManager.all()).thenReturn(Seq.empty)
+      when(batchManager.size).thenReturn(0)
+      when(batchManager.all()).thenReturn(Seq.empty)
+
+      LivySessionMetrics.register(registry, interactiveManager, batchManager)
+      LivySessionMetrics.register(registry, interactiveManager, batchManager)
+
+      registry.getGauges.size() shouldBe 18
+    }
+
+    it("should report zero when there are no sessions") {
+      val registry = new MetricRegistry()
+      registerMetrics(registry = registry)
+
+      expectedGauges.foreach { name =>
+        gaugeValue(registry, name) shouldBe 0
+      }
+    }
+
+    it("should count interactive sessions by state") {
+      val registry = new MetricRegistry()
+      val sessions = Seq(
+        mockInteractive(SessionState.Idle),
+        mockInteractive(SessionState.Idle),
+        mockInteractive(SessionState.Busy),
+        mockInteractive(SessionState.Starting),
+        mockInteractive(SessionState.ShuttingDown),
+        mockInteractive(SessionState.Dead()),
+        mockInteractive(SessionState.Error()),
+        mockInteractive(SessionState.Killed())
+      )
+      registerMetrics(interactiveSessions = sessions, registry = registry)
+
+      gaugeValue(registry, "livy.sessions.interactive.total") shouldBe 8
+      gaugeValue(registry, "livy.sessions.interactive.idle") shouldBe 2
+      gaugeValue(registry, "livy.sessions.interactive.busy") shouldBe 1
+      gaugeValue(registry, "livy.sessions.interactive.starting") shouldBe 1
+      gaugeValue(registry, "livy.sessions.interactive.shutting_down") shouldBe 1
+      gaugeValue(registry, "livy.sessions.interactive.dead") shouldBe 1
+      gaugeValue(registry, "livy.sessions.interactive.error") shouldBe 1
+      gaugeValue(registry, "livy.sessions.interactive.killed") shouldBe 1
+    }
+
+    it("should count batch sessions by state including succeeded alias") {
+      val registry = new MetricRegistry()
+      val sessions = Seq(
+        mockBatch(SessionState.Starting),
+        mockBatch(SessionState.Running),
+        mockBatch(SessionState.Success()),
+        mockBatch(SessionState.Dead()),
+        mockBatch(SessionState.Error()),
+        mockBatch(SessionState.Killed())
+      )
+      registerMetrics(batchSessions = sessions, registry = registry)
+
+      gaugeValue(registry, "livy.sessions.batch.total") shouldBe 6
+      gaugeValue(registry, "livy.sessions.batch.starting") shouldBe 1
+      gaugeValue(registry, "livy.sessions.batch.running") shouldBe 1
+      gaugeValue(registry, "livy.sessions.batch.success") shouldBe 1
+      gaugeValue(registry, "livy.sessions.batch.dead") shouldBe 1
+      gaugeValue(registry, "livy.sessions.batch.error") shouldBe 1
+      gaugeValue(registry, "livy.sessions.batch.killed") shouldBe 1
+    }
+
+    it("should count batch success for succeeded state alias") {
+      val registry = new MetricRegistry()
+      val succeededState = mock[SessionState]
+      when(succeededState.toString).thenReturn("succeeded")
+      val session = mock[BatchSession]
+      when(session.state).thenReturn(succeededState)
+      registerMetrics(batchSessions = Seq(session), registry = registry)
+
+      gaugeValue(registry, "livy.sessions.batch.success") shouldBe 1
+    }
+
+    it("should compute overall totals across interactive and batch sessions") {
+      val registry = new MetricRegistry()
+      val interactive = Seq(
+        mockInteractive(SessionState.Idle),
+        mockInteractive(SessionState.Dead())
+      )
+      val batch = Seq(
+        mockBatch(SessionState.Running),
+        mockBatch(SessionState.Success())
+      )
+      registerMetrics(
+        interactiveSessions = interactive, batchSessions = batch, registry = registry)
+
+      gaugeValue(registry, "livy.sessions.total") shouldBe 4
+      gaugeValue(registry, "livy.sessions.active.total") shouldBe 2
+      gaugeValue(registry, "livy.sessions.terminal.total") shouldBe 2
+    }
+
+    it("should return zero from gauges when session managers throw") {
+      val registry = new MetricRegistry()
+      val interactiveManager = mock[InteractiveSessionManager]
+      val batchManager = mock[BatchSessionManager]
+      when(interactiveManager.size).thenReturn(1)
+      when(interactiveManager.all()).thenThrow(new RuntimeException("interactive failure"))
+      when(batchManager.size).thenReturn(1)
+      when(batchManager.all()).thenThrow(new RuntimeException("batch failure"))
+
+      LivySessionMetrics.register(registry, interactiveManager, batchManager)
+
+      gaugeValue(registry, "livy.sessions.interactive.idle") shouldBe 0
+      gaugeValue(registry, "livy.sessions.batch.running") shouldBe 0
+      gaugeValue(registry, "livy.sessions.active.total") shouldBe 0
+      gaugeValue(registry, "livy.sessions.terminal.total") shouldBe 0
+    }
+  }
+}