Capture unfinished perf logger entries (#44)

google · Aug 1, 2023 · eebc550 · eebc550
1 parent f5fb99f
commit eebc550
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 4 deletions.
diff --git a/src/java/com/google/cloud/bigquery/dwhassessment/hooks/logger/EventRecordConstructor.java b/src/java/com/google/cloud/bigquery/dwhassessment/hooks/logger/EventRecordConstructor.java
@@ -243,11 +243,18 @@ private static <C, G extends Iterable<C>> Optional<String> generateCountersJson(
     return outerObj.length() > 0 ? Optional.of(outerObj.toString()) : Optional.empty();
   }
 
-  private static String dumpPerfData(PerfLogger perfLogger) {
+  private String dumpPerfData(PerfLogger perfLogger) {
     JSONObject perfObj = new JSONObject();
-
-    for (String key : perfLogger.getEndTimes().keySet()) {
-      perfObj.put(key, perfLogger.getDuration(key));
+    long now = clock.millis();
+
+    for (String key : perfLogger.getStartTimes().keySet()) {
+      long duration = perfLogger.getDuration(key);
+      // Some perf logger entries are finished after the hook. Make the best effort to capture them
+      // here with the duration at the current time.
+      if (duration == 0L) {
+        duration = now - perfLogger.getStartTime(key);
+      }
+      perfObj.put(key, duration);
     }
 
     return perfObj.toString();

diff --git a/...ests/com/google/cloud/bigquery/dwhassessment/hooks/logger/EventRecordConstructorTest.java b/...ests/com/google/cloud/bigquery/dwhassessment/hooks/logger/EventRecordConstructorTest.java
@@ -47,6 +47,7 @@
 import org.apache.hadoop.hive.ql.exec.tez.TezTask;
 import org.apache.hadoop.hive.ql.hooks.HookContext;
 import org.apache.hadoop.hive.ql.hooks.HookContext.HookType;
+import org.apache.hadoop.hive.ql.log.PerfLogger;
 import org.apache.hadoop.hive.ql.metadata.Hive;
 import org.apache.hadoop.hive.ql.plan.HiveOperation;
 import org.apache.hadoop.hive.ql.plan.TezWork;
@@ -197,6 +198,22 @@ public void postExecHook_success() {
     assertThat(record).hasValue(TestUtils.createPostExecRecord(EventStatus.SUCCESS));
   }
 
+  @Test
+  public void postExecHook_capturesAllEntriesFromPerfLogger_success() {
+    hookContext.setHookType(HookType.POST_EXEC_HOOK);
+    PerfLogger perfLogger = PerfLogger.getPerfLogger(state.getConf(), false);
+    perfLogger.PerfLogBegin("test_caller", "test_method_1");
+    perfLogger.PerfLogEnd("test_caller", "test_method_1");
+    perfLogger.PerfLogBegin("test_caller", "test_method_2");
+
+    // Act
+    GenericRecord record = eventRecordConstructor.constructEvent(hookContext).get();
+
+    // Assert
+    assertThat((String) record.get("PerfObject")).contains("test_method_1");
+    assertThat((String) record.get("PerfObject")).contains("test_method_2");
+  }
+
   @Test
   public void postExecHook_recordsYarnApplicationDataWhenPossible() {
     hookContext.setHookType(HookType.POST_EXEC_HOOK);