From e33721044df33db4932d6b13583902ddee7e7445 Mon Sep 17 00:00:00 2001
From: Nathaniel Bauernfeind <natebauernfeind@deephaven.io>
Date: Mon, 13 Nov 2023 17:01:14 -0700
Subject: [PATCH] The Fixes

---
 .../impl/perf/QueryPerformanceNugget.java     | 23 +++++++++++++++----
 .../perf/QueryPerformanceRecorderImpl.java    |  4 ++++
 py/server/deephaven/perfmon.py                |  2 +-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java
index 160fa90d1d2..3d6928a7adc 100644
--- a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java
+++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java
@@ -222,10 +222,7 @@ protected QueryPerformanceNugget(
         final RuntimeMemory runtimeMemory = RuntimeMemory.getInstance();
         runtimeMemory.read(startMemorySample);
 
-        startClockEpochNanos = DateTimeUtils.millisToNanos(System.currentTimeMillis());
-        onBaseEntryStart();
-
-        state = QueryState.RUNNING;
+        state = QueryState.NOT_STARTED;
         shouldLogThisAndStackParents = false;
     }
 
@@ -260,6 +257,24 @@ public void markStartTime() {
         startClockEpochNanos = DateTimeUtils.millisToNanos(System.currentTimeMillis());
     }
 
+    @Override
+    public synchronized void onBaseEntryStart() {
+        super.onBaseEntryStart();
+        if (state == QueryState.RUNNING) {
+            throw new IllegalStateException("Nugget was already started");
+        }
+        state = QueryState.RUNNING;
+    }
+
+    @Override
+    public synchronized void onBaseEntryEnd() {
+        if (state != QueryState.RUNNING) {
+            throw new IllegalStateException("Nugget isn't running");
+        }
+        state = QueryState.SUSPENDED;
+        super.onBaseEntryEnd();
+    }
+
     /**
      * Mark this nugget {@link QueryState#FINISHED} and notify the recorder.
      *
diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceRecorderImpl.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceRecorderImpl.java
index 312c1ef2ce4..37a754e656f 100644
--- a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceRecorderImpl.java
+++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceRecorderImpl.java
@@ -159,6 +159,8 @@ private SafeCloseable resumeInternal() {
 
     private void startCatchAll() {
         catchAllNugget = nuggetFactory.createForCatchAll(queryNugget, operationNuggets.size(), this::releaseNugget);
+        catchAllNugget.markStartTime();
+        catchAllNugget.onBaseEntryStart();
     }
 
     private void stopCatchAll(final boolean abort) {
@@ -200,6 +202,8 @@ public synchronized QueryPerformanceNugget getNugget(@NotNull final String name,
 
         final QueryPerformanceNugget nugget = nuggetFactory.createForOperation(
                 parent, operationNuggets.size(), name, inputSize, this::releaseNugget);
+        nugget.markStartTime();
+        nugget.onBaseEntryStart();
         operationNuggets.add(nugget);
         userNuggetStack.addLast(nugget);
         return nugget;
diff --git a/py/server/deephaven/perfmon.py b/py/server/deephaven/perfmon.py
index 995c77994a6..213ee349c56 100644
--- a/py/server/deephaven/perfmon.py
+++ b/py/server/deephaven/perfmon.py
@@ -11,7 +11,7 @@
 
 from deephaven import DHError
 from deephaven.jcompat import j_map_to_dict
-from deephaven.table import Table
+from deephaven.table import Table, TreeTable
 
 _JPerformanceQueries = jpy.get_type("io.deephaven.engine.table.impl.util.PerformanceQueries")
 _JMetricsManager = jpy.get_type("io.deephaven.util.metrics.MetricsManager")