Update handle assignment error metric (#967)

* Update handle assignment error metric * Address feedback
linkedin · Nov 6, 2023 · 9aac481 · 9aac481
1 parent c082369
commit 9aac481
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java b/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java
@@ -669,7 +669,7 @@ public void onSessionExpired() {
     threadPoolExecutor.submit(() -> {
       Instant start = Instant.now();
       try {
-        getAssignmentsFuture(assignmentChangeFutures, start);
+        getAssignmentsFuture(assignmentChangeFutures, start, false);
       } catch (Exception e) {
         _log.warn("Hit exception while clearing the assignment list", e);
       } finally {
@@ -701,7 +701,8 @@ public void onNewSession() {
     _zkSessionExpired = false;
   }
 
-  private void getAssignmentsFuture(List<Future<Boolean>> assignmentChangeFutures, Instant start)
+  private void getAssignmentsFuture(List<Future<Boolean>> assignmentChangeFutures, Instant start,
+      boolean isDatastreamUpdate)
       throws TimeoutException, InterruptedException {
     for (Future<Boolean> assignmentChangeFuture : assignmentChangeFutures) {
       if (Duration.between(start, Instant.now()).compareTo(ASSIGNMENT_TIMEOUT) > 0) {
@@ -711,6 +712,8 @@ private void getAssignmentsFuture(List<Future<Boolean>> assignmentChangeFutures,
         assignmentChangeFuture.get(ASSIGNMENT_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
       } catch (ExecutionException e) {
         _log.warn("onAssignmentChange call threw exception", e);
+        EventType meter = isDatastreamUpdate ? HANDLE_DATASTREAM_CHANGE_WITH_UPDATE : HANDLE_ASSIGNMENT_CHANGE;
+        _metrics.updateKeyedMeter(CoordinatorMetrics.getKeyedMeter(meter), 1);
       }
     }
   }
@@ -768,11 +771,11 @@ private void retryHandleAssignmentChange(boolean isDatastreamUpdate) {
           "onAssignmentChange retries reached threshold of {}. Not queuing further retry on onAssignmentChange event.",
           _config.getMaxAssignmentRetryCount());
       queueEvent = false;
-    }
 
-    EventType meter = isDatastreamUpdate ? HANDLE_DATASTREAM_CHANGE_WITH_UPDATE : HANDLE_ASSIGNMENT_CHANGE;
-    _log.warn("Updating metric for event " + meter);
-    _metrics.updateKeyedMeter(CoordinatorMetrics.getKeyedMeter(meter), 1);
+      EventType meter = isDatastreamUpdate ? HANDLE_DATASTREAM_CHANGE_WITH_UPDATE : HANDLE_ASSIGNMENT_CHANGE;
+      _log.warn("Updating metric for event={}", meter);
+      _metrics.updateKeyedMeter(CoordinatorMetrics.getKeyedMeter(meter), 1);
+    }
 
     if (queueEvent) {
       _log.warn("Queuing onAssignmentChange event");
@@ -847,7 +850,7 @@ private void handleAssignmentChange(boolean isDatastreamUpdate) throws TimeoutEx
     // Wait till all the futures are complete or timeout.
     Instant start = Instant.now();
     try {
-      getAssignmentsFuture(assignmentChangeFutures, start);
+      getAssignmentsFuture(assignmentChangeFutures, start, isDatastreamUpdate);
     } catch (TimeoutException e) {
       // if it's timeout then we will retry
       _log.warn("Timeout when doing the assignment", e);