From 676a1320421b8dc87be437104de2cc013b4ac67e Mon Sep 17 00:00:00 2001
From: Sanskar Modi <sanskarmodi97@gmail.com>
Date: Mon, 23 Sep 2024 14:03:50 +0530
Subject: [PATCH] converted counter to gauge

---
 .../apache/celeborn/service/deploy/worker/Worker.scala    | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index c8d0a442ce..dcb457b834 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -435,6 +435,13 @@ private[celeborn] class Worker(
       0
     }
   }
+  workerSource.addGauge(WorkerSource.UNRELEASED_SHUFFLE_COUNT) { () =>
+    if (shutdown.get() && workerStatusManager.currentWorkerStatus.getState == State.Exit) {
+      storageManager.shuffleKeySet().size
+    } else {
+      0
+    }
+  }
   workerSource.addGauge(WorkerSource.CLEAN_TASK_QUEUE_SIZE) { () =>
     cleanTaskQueue.size()
   }
@@ -959,7 +966,6 @@ private[celeborn] class Worker(
     } else {
       logWarning(s"Waiting for all shuffle expired cost ${waitTime}ms, " +
         s"unreleased shuffle: \n${unreleasedShuffleKeys.asScala.mkString("[", ", ", "]")}")
-      workerSource.incCounter(UNRELEASED_SHUFFLE_COUNT, unreleasedShuffleKeys.size)
     }
     workerStatusManager.transitionState(State.Exit)
   }