Return step stats (#464)

* Use `1.0.0-SNAPSHOT` * Use `long` for CPU Ticks * Expose CPU Ticks * Set steps and ticks on globals on script exit
verygoodsecurity · Mar 7, 2024 · b5d940e · b5d940e
1 parent 3302d56
commit b5d940e
Show file tree

Hide file tree

Showing 8 changed files with 74 additions and 72 deletions.
diff --git a/larky/pom.xml b/larky/pom.xml
@@ -7,7 +7,7 @@
 
     <groupId>com.verygood.security</groupId>
     <artifactId>larky</artifactId>
-    <version>1.0-SNAPSHOT</version>
+    <version>1.0.0-SNAPSHOT</version>
 
     <packaging>jar</packaging>
 
@@ -33,7 +33,7 @@
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
 
-        <libstarlark.version>1.0-SNAPSHOT</libstarlark.version>
+        <libstarlark.version>1.0.0-SNAPSHOT</libstarlark.version>
 
         <apache.commons-text.version>1.11.0</apache.commons-text.version>
         <google.crypto.tink>1.6.1</google.crypto.tink>

diff --git a/larky/src/main/java/com/verygood/security/larky/parser/LarkyEvaluator.java b/larky/src/main/java/com/verygood/security/larky/parser/LarkyEvaluator.java
@@ -5,18 +5,18 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
+import com.verygood.security.larky.ModuleSupplier;
+import com.verygood.security.larky.annot.Library;
+import com.verygood.security.larky.console.Console;
+import com.verygood.security.larky.modules.utils.Reporter;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
-
-import com.verygood.security.larky.ModuleSupplier;
-import com.verygood.security.larky.annot.Library;
-import com.verygood.security.larky.console.Console;
-import com.verygood.security.larky.modules.utils.Reporter;
-
+import lombok.Builder;
+import lombok.Getter;
 import net.starlark.java.annot.StarlarkAnnotations;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.eval.EvalException;
@@ -31,21 +31,18 @@
 import net.starlark.java.syntax.Program;
 import net.starlark.java.syntax.StarlarkFile;
 import net.starlark.java.syntax.SyntaxError;
-
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.VisibleForTesting;
 
-import lombok.Builder;
-import lombok.Data;
-import lombok.Getter;
-
 /**
  * An utility class for traversing and evaluating the config file dependency graph.
  */
 public final class LarkyEvaluator {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  public static final String EXECUTION_STEPS = "_EXECUTION_STEPS_";
+  public static final String EXECUTION_CPU_TICKS = "_EXECUTION_CPU_TICKS_";
 
   private final LinkedHashSet<String> pending = new LinkedHashSet<>();
   private final Map<String, Module> loaded = new HashMap<>();
@@ -74,28 +71,32 @@ public LarkyEvaluator(LarkyScript larkyScript, Console console) {
   }
 
   /**
-   * The output of a Larky script is an evaluated {@link Module} that contains
-   * various attributes such as {@link Module#getGlobals()},
-   * {@link Module#getPredeclaredBindings()}, as well as other various items.
-   *
+   * The output of a Larky script is an evaluated {@link Module} that contains various attributes such as
+   * {@link Module#getGlobals()}, {@link Module#getPredeclaredBindings()}, as well as other various items.
+   * <p>
    * Sometimes, when evaluating a Larky script, there is some output that is generated.
-   *
-   * This interface encapsulates an interface that allows the caller to introspect
-   * the result of a Larky script evaluation.
+   * <p>
+   * This interface encapsulates an interface that allows the caller to introspect the result of a Larky script
+   * evaluation.
    */
   public interface EvaluationResult {
-      boolean hasOutput();
-      boolean hasModule();
 
-      Object getOutput();
-      Module getModule();
+    boolean hasOutput();
+
+    boolean hasModule();
+
+    Object output();
+
+    Module module();
+
   }
 
   @Builder
-  @Data
-  protected static class DefaultEvaluationResult implements EvaluationResult {
-    private Object output;
-    private Module module;
+  protected record DefaultEvaluationResult(
+      Object output,
+      Module module
+      )
+      implements EvaluationResult {
 
     public boolean hasOutput() {
       return output != null;
@@ -114,9 +115,8 @@ public EvaluationResult eval(StarFile content)
     Module module = loaded.get(content.path());
     if (module != null) {
       return DefaultEvaluationResult.builder()
-        .output(null)
-        .module(module)
-        .build();
+          .module(module)
+          .build();
     }
     pending.add(content.path());
 
@@ -139,16 +139,20 @@ public EvaluationResult eval(StarFile content)
       thread.setPrintHandler(reporter::report);
       try {
         starlarkOutput = Starlark.execFileProgram(prog, module, thread);
-      } catch(EvalException cause) {
+      } catch (EvalException cause) {
         throw new StarlarkEvalWrapper.Exc.RuntimeEvalException(cause, thread);
       }
+
+      // Set some statistical information
+      module.setGlobal(EXECUTION_STEPS, thread.getExecutedSteps());
+      module.setGlobal(EXECUTION_CPU_TICKS, thread.getCpuTicks());
     }
     pending.remove(content.path());
     loaded.put(content.path(), module);
     return DefaultEvaluationResult.builder()
-      .output(starlarkOutput)
-      .module(module)
-      .build();
+        .output(starlarkOutput)
+        .module(module)
+        .build();
   }
 
   @VisibleForTesting
@@ -171,7 +175,7 @@ public Module load(String moduleToLoad) {
       Module loadedModule = null;
       try {
         if (!ResourceContentStarFile.startsWithPrefix(moduleToLoad)) {
-          loadedModule = evaluator.eval(content.resolve(moduleToLoad + LarkyScript.STAR_EXTENSION)).getModule();
+          loadedModule = evaluator.eval(content.resolve(moduleToLoad + LarkyScript.STAR_EXTENSION)).module();
           return loadedModule;
         }
 
@@ -184,22 +188,21 @@ public Module load(String moduleToLoad) {
          * Check if the module is in the module set. If it is, return a module with an environment
          * of the module that was passed in via the module set.
          */
-        else if(isNativeJavaModule(targetModule)) {
+        else if (isNativeJavaModule(targetModule)) {
           loadedModule = fromNativeModule(targetModule);
-        }
-        else {
+        } else {
           // try to load from directory...
           ResourceContentStarFile starFile = ResourceContentStarFile.buildStarFile(moduleToLoad);
-          loadedModule = evaluator.eval(starFile).getModule();
+          loadedModule = evaluator.eval(starFile).module();
         }
 
       } catch (IOException | InterruptedException | EvalException e) {
         throw new RuntimeException(
-          String.format(
-            "Encountered error (%s) while attempting to load %s from module: %s.",
-            e.getMessage(),
-            moduleToLoad,
-            this.content.path()), e);
+            String.format(
+                "Encountered error (%s) while attempting to load %s from module: %s.",
+                e.getMessage(),
+                moduleToLoad,
+                this.content.path()), e);
       }
       return loadedModule;
     }
@@ -234,10 +237,10 @@ private Module fromNativeModule(String moduleToLoad) throws IOException, Interru
         StarlarkThread thread = new StarlarkThread(mu, evaluator.getLarkySemantics());
         try {
           Starlark.execFile(
-            ParserInput.fromString(String.format("%1$s = _%1$s", moduleToLoad), "<builtin>"),
-            evaluator.getStarlarkValidationOptions(),
-            newModule,
-            thread
+              ParserInput.fromString(String.format("%1$s = _%1$s", moduleToLoad), "<builtin>"),
+              evaluator.getStarlarkValidationOptions(),
+              newModule,
+              thread
           );
         } catch (InterruptedException | EvalException | SyntaxError.Exception e) {
           throw new StarlarkEvalWrapper.Exc.RuntimeEvalException(e, thread);
@@ -276,7 +279,7 @@ Program compileStarlarkProgram(Module module, ParserInput input, FileOptions opt
       throw new EvalException(
           String.format(
               "Error compiling Starlark program: %1$s%n" +
-                  "%2$s",
+              "%2$s",
               input.getFile(),
               String.join("\n", errs)));
     }
@@ -307,10 +310,10 @@ private RuntimeException throwCycleError(String cycleElement) throws EvalExcepti
   }
 
   /**
-   * Create the environment for all evaluations (will be shared between all the dependent files
-   * loaded).
+   * Create the environment for all evaluations (will be shared between all the dependent files loaded).
    */
-  private ImmutableMap<String, Object> createEnvironment(Iterable<Class<?>> globalModules, Map<String, Object> globals) {
+  private ImmutableMap<String, Object> createEnvironment(Iterable<Class<?>> globalModules,
+      Map<String, Object> globals) {
     Map<String, Object> env = Maps.newHashMap();
 
     for (Class<?> module : globalModules) {

diff --git a/larky/src/main/java/com/verygood/security/larky/parser/LarkyScript.java b/larky/src/main/java/com/verygood/security/larky/parser/LarkyScript.java
@@ -22,26 +22,23 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.verygood.security.larky.LarkySemantics;
+import com.verygood.security.larky.ModuleSupplier;
+import com.verygood.security.larky.ModuleSupplier.ModuleSet;
+import com.verygood.security.larky.console.Console;
 import java.io.IOException;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-
-import com.verygood.security.larky.LarkySemantics;
-import com.verygood.security.larky.ModuleSupplier;
-import com.verygood.security.larky.ModuleSupplier.ModuleSet;
-import com.verygood.security.larky.console.Console;
-
+import lombok.Getter;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Module;
 import net.starlark.java.eval.StarlarkSemantics;
 import net.starlark.java.syntax.FileOptions;
 
-import lombok.Getter;
-
 /**
  * Loads Larky out of Starlark files.
  */
@@ -143,13 +140,13 @@ private LarkyEvaluator.EvaluationResult executeStarlark(StarFile content, Module
   @VisibleForTesting
   public Module executeSkylark(StarFile content, ModuleSet moduleSet, Console console)
       throws IOException, InterruptedException, EvalException {
-    return executeStarlark(content, moduleSet, console).getModule();
+    return executeStarlark(content, moduleSet, console).module();
   }
 
   @VisibleForTesting
   public Object executeSkylarkWithOutput(StarFile content, ModuleSet moduleSet, Console console)
       throws IOException, InterruptedException, EvalException {
-    return executeStarlark(content, moduleSet, console).getOutput();
+    return executeStarlark(content, moduleSet, console).output();
   }
 
   public ParsedStarFile evaluate(StarFile content, Console console)
@@ -193,7 +190,7 @@ private ParsedStarFile loadStarFileInternal(StarFile content, ModuleSet moduleSe
       throws IOException, EvalException {
     Module module;
     try {
-      module = new LarkyEvaluator(this, moduleSet, console).eval(content).getModule();
+      module = new LarkyEvaluator(this, moduleSet, console).eval(content).module();
     } catch (InterruptedException e) {
       // This should not happen since we shouldn't have anything interruptable during loading.
       throw new RuntimeException("Internal error", e);

diff --git a/libstarlark/pom.xml b/libstarlark/pom.xml
@@ -7,7 +7,7 @@
 
     <groupId>net.starlark.java</groupId>
     <artifactId>libstarlark</artifactId>
-    <version>1.0-SNAPSHOT</version>
+    <version>1.0.0-SNAPSHOT</version>
 
     <packaging>jar</packaging>
 

diff --git a/libstarlark/src/main/java/net/starlark/java/eval/CpuProfiler.java b/libstarlark/src/main/java/net/starlark/java/eval/CpuProfiler.java
@@ -163,7 +163,7 @@ static void stop() throws IOException {
   }
 
   /** Records a profile event. */
-  void addEvent(int ticks, ImmutableList<Debug.Frame> stack) {
+  void addEvent(long ticks, ImmutableList<Debug.Frame> stack) {
     pprof.writeEvent(ticks, stack);
   }
 
@@ -289,7 +289,7 @@ private static final class PprofWriter {
       }
     }
 
-    synchronized void writeEvent(int ticks, ImmutableList<Debug.Frame> stack) {
+    synchronized void writeEvent(long ticks, ImmutableList<Debug.Frame> stack) {
       if (this.error == null) {
         try {
           ByteArrayOutputStream sample = new ByteArrayOutputStream();

diff --git a/libstarlark/src/main/java/net/starlark/java/eval/StarlarkThread.java b/libstarlark/src/main/java/net/starlark/java/eval/StarlarkThread.java
@@ -20,7 +20,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.Immutable;
 import net.starlark.java.syntax.Location;
@@ -56,7 +56,7 @@ public final class StarlarkThread {
   // not in the corresponding pop, but when the last frame is popped, because
   // the profiler session might start in the middle of a call and/or run beyond
   // the lifetime of this thread.
-  final AtomicInteger cpuTicks = new AtomicInteger();
+  final AtomicLong cpuTicks = new AtomicLong();
   @Nullable private CpuProfiler profiler;
   private StarlarkThread savedThread; // saved StarlarkThread, when profiling reentrant evaluation
 
@@ -77,6 +77,8 @@ public long getExecutedSteps() {
     return steps;
   }
 
+  public long getCpuTicks() { return cpuTicks.get();}
+
   /**
    * Sets the maximum number of Starlark computation steps that may be executed by this thread (see
    * {@link #getExecutedSteps}). When the step counter reaches or exceeds this value, execution
@@ -258,7 +260,7 @@ void pop() {
     Frame fr = callstack.get(last);
 
     if (profiler != null) {
-      int ticks = cpuTicks.getAndSet(0);
+      long ticks = cpuTicks.getAndSet(0);
       if (ticks > 0) {
         profiler.addEvent(ticks, getDebugCallStack());
       }

diff --git a/pom.xml b/pom.xml
@@ -6,7 +6,7 @@
 
     <groupId>com.verygood.security</groupId>
     <artifactId>starlarky</artifactId>
-    <version>1.0-SNAPSHOT</version>
+    <version>1.0.0-SNAPSHOT</version>
     <packaging>pom</packaging>
 
     <modules>

diff --git a/runlarky/pom.xml b/runlarky/pom.xml
@@ -7,7 +7,7 @@
 
     <groupId>com.verygood.security</groupId>
     <artifactId>runlarky</artifactId>
-    <version>1.0-SNAPSHOT</version>
+    <version>1.0.0-SNAPSHOT</version>
 
     <repositories>
         <repository>