From 810bd2dcc1204f1550426453e51975e438c73a28 Mon Sep 17 00:00:00 2001 From: arunkumarm Date: Tue, 18 Aug 2026 19:30:36 +0530 Subject: [PATCH] [LIVY-1070] Add LivySessionMetrics Codahale gauges for session monitoring Register 18 livy.sessions.* session count gauges into the MetricRegistry at server startup via LivySessionMetrics.register(). Gauges reflect interactive and batch session counts by state and are exposed via the existing /metrics endpoint. Includes LivySessionMetricsSpec unit tests. --- .../org/apache/livy/server/LivyServer.scala | 2 + .../livy/server/LivySessionMetrics.scala | 142 ++++++++++++ .../livy/server/LivySessionMetricsSpec.scala | 210 ++++++++++++++++++ 3 files changed, 354 insertions(+) create mode 100644 server/src/main/scala/org/apache/livy/server/LivySessionMetrics.scala create mode 100644 server/src/test/scala/org/apache/livy/server/LivySessionMetricsSpec.scala diff --git a/server/src/main/scala/org/apache/livy/server/LivyServer.scala b/server/src/main/scala/org/apache/livy/server/LivyServer.scala index 73f259d5b..68a3880a9 100644 --- a/server/src/main/scala/org/apache/livy/server/LivyServer.scala +++ b/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 { diff --git a/server/src/main/scala/org/apache/livy/server/LivySessionMetrics.scala b/server/src/main/scala/org/apache/livy/server/LivySessionMetrics.scala new file mode 100644 index 000000000..0619f5216 --- /dev/null +++ b/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"))) + } +} diff --git a/server/src/test/scala/org/apache/livy/server/LivySessionMetricsSpec.scala b/server/src/test/scala/org/apache/livy/server/LivySessionMetricsSpec.scala new file mode 100644 index 000000000..629216a56 --- /dev/null +++ b/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 + } + } +}