%time magic #104

restoring CPU time measurements

Author: andrus

jjava-jupyter/src/main/java/org/dflib/jjava/jupyter/instrumentation/WallAndCpuTimer.java

@@ -0,0 +1,78 @@
+package org.dflib.jjava.jupyter.instrumentation;
+
+import java.lang.management.ThreadMXBean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
+
+/**
+ * A timer that makes an attempt to capture CPU and user time in addition to the overall elapsed time.
+ */
+public abstract class WallAndCpuTimer implements EvalTimer {
+
+    private final ThreadMXBean threadMXBean;
+    private final AtomicLong totalWallTime;
+    private final AtomicLong totalUserTime;
+    private final AtomicLong totalCpuTime;
+
+    public static boolean canMeasureCpuTimes(ThreadMXBean threadMXBean) {
+        // according to JavaDocs, isThreadCpuTimeEnabled() can throw in some cases
+        try {
+            return threadMXBean.isCurrentThreadCpuTimeSupported() && threadMXBean.isThreadCpuTimeEnabled();
+        } catch (Exception e) {
+            return false;
+        }
+    }
+
+    public WallAndCpuTimer(ThreadMXBean threadMXBean) {
+        this.threadMXBean = threadMXBean;
+        this.totalWallTime = new AtomicLong();
+        this.totalUserTime = new AtomicLong();
+        this.totalCpuTime = new AtomicLong();
+    }
+
+    @Override
+    public void start() {
+        totalWallTime.set(0);
+        totalUserTime.set(0);
+        totalCpuTime.set(0);
+    }
+
+    @Override
+    public <T> T runAndMeasureStep(Supplier<T> timedAction) {
+
+        long t0 = System.nanoTime();
+        long ut0 = threadMXBean.getCurrentThreadUserTime();
+        long ct0 = threadMXBean.getCurrentThreadCpuTime();
+
+        try {
+            return timedAction.get();
+        } finally {
+
+            long ct1 = threadMXBean.getCurrentThreadCpuTime();
+            long ut1 = threadMXBean.getCurrentThreadUserTime();
+            long t1 = System.nanoTime();
+
+            totalWallTime.addAndGet(t1 - t0);
+
+            // sanity check
+            if (ut0 != -1 && ut1 != -1 && ct0 != -1 && ct1 != -1) {
+                totalUserTime.addAndGet(ut1 - ut0);
+                totalCpuTime.addAndGet(ct1 - ct0);
+            } else {
+                totalUserTime.set(-1);
+                totalCpuTime.set(-1);
+            }
+        }
+    }
+
+    @Override
+    public void stop() {
+        onStop(totalWallTime.get(), totalUserTime.get(), totalCpuTime.get());
+    }
+
+    protected abstract void onStop(
+            long wallTimeNanos,
+            long userTimeNanos,
+            long totalTimeNanos);
+
+}

jjava-kernel/src/main/java/org/dflib/jjava/kernel/magics/TimeMagic.java

@@ -1,12 +1,15 @@
 package org.dflib.jjava.kernel.magics;
 
+import org.dflib.jjava.jupyter.instrumentation.WallAndCpuTimer;
 import org.dflib.jjava.jupyter.instrumentation.WallTimer;
 import org.dflib.jjava.jupyter.kernel.BaseKernel;
 import org.dflib.jjava.jupyter.kernel.display.DisplayData;
 import org.dflib.jjava.jupyter.kernel.magic.CellMagic;
 import org.dflib.jjava.jupyter.kernel.magic.LineMagic;
 import org.dflib.jjava.kernel.JavaKernel;
 
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadMXBean;
 import java.util.List;
 import java.util.UUID;
 
@@ -28,14 +31,25 @@ public DisplayData eval(JavaKernel kernel, List<String> args, String body) throw
     }
 
     private DisplayData timeAndRunCode(JavaKernel kernel, String code) {
-        System.out.println("TimeMagic thread: " + Thread.currentThread().getName());
+        ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
         return kernel.evalBuilder(code)
                 .resolveMagics()
                 .renderResults()
-                .timed(new TMWallTimer(kernel))
+                .timed(WallAndCpuTimer.canMeasureCpuTimes(threadMXBean)
+                        ? new TMWallAndCpuTimer(kernel, threadMXBean)
+                        : new TMWallTimer(kernel))
                 .eval();
     }
 
+    private static void displayCpuTime(BaseKernel kernel, long userTimeNanos, long sysTimeNanos, long totalTimeNanos) {
+        String cpuTimes = String.format("CPU times: user %.3f s, sys %.3f s, total %.3f s",
+                userTimeNanos / NANONS_IN_SEC,
+                sysTimeNanos / NANONS_IN_SEC,
+                totalTimeNanos / NANONS_IN_SEC);
+
+        kernel.display(new DisplayData(cpuTimes).setDisplayId(UUID.randomUUID().toString()));
+    }
+
     private static void displayWallTime(BaseKernel kernel, long wallTimeNanos) {
         String wallTime = String.format("Wall time: %.3f s", wallTimeNanos / NANONS_IN_SEC);
         kernel.display(new DisplayData(wallTime).setDisplayId(UUID.randomUUID().toString()));
@@ -53,4 +67,24 @@ protected void onStop(long wallTimeNanos) {
             displayWallTime(kernel, wallTimeNanos);
         }
     }
+
+
+    static class TMWallAndCpuTimer extends WallAndCpuTimer {
+        private final BaseKernel kernel;
+
+        public TMWallAndCpuTimer(BaseKernel kernel, ThreadMXBean threadMXBean) {
+            super(threadMXBean);
+            this.kernel = kernel;
+        }
+
+        @Override
+        protected void onStop(long wallTimeNanos, long userTimeNanos, long totalTimeNanos) {
+            // sanity check: CPU and user times can't be measured sometimes
+            if (userTimeNanos != -1 && totalTimeNanos != -1) {
+                displayCpuTime(kernel, userTimeNanos, totalTimeNanos - userTimeNanos, totalTimeNanos);
+            }
+
+            displayWallTime(kernel, wallTimeNanos);
+        }
+    }
 }