converted counter to gauge

apache · Sep 23, 2024 · 676a132 · 676a132
1 parent 11db012
commit 676a132
Showing 1 changed file with 7 additions and 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
@@ -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)
   }