bazel syntax: delete Callstack (used by memory allocation tracker)

The Callstack was a copy of the thread's stack saved in thread-local
storage. It was notified not only of every function call entry and
exit, but the same for every expression too. This made it not only
intrusive but incompatible with a compiled representation.

This change deletes Callstack and replaces it with an interface,
Debug.ThreadHook, that lets a single client be notified when a
thread pushes its first frame or pops its last.
The allocation tracker uses these notifications to maintain a mapping
of Java threads to StarlarkThreads (similar to CpuProfiler),
so that when the memory allocator hook is called, it can ascertain
which Starlark thread caused the allocation.

The tests have been reworked to call the interpreter instead of
peeking and poking into the callstack directly.

BUG=149023294
PiperOrigin-RevId: 293629363
diff --git a/src/main/java/com/google/devtools/build/lib/packages/RuleFunction.java b/src/main/java/com/google/devtools/build/lib/packages/RuleFunction.java
index 3b73342..1be4d07 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/RuleFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/RuleFunction.java
@@ -14,7 +14,7 @@
 
 package com.google.devtools.build.lib.packages;
 
-/** Marker interface for a native or Skylark rule function. */
+/** Interface for a native or Starlark rule function. */
 public interface RuleFunction {
   RuleClass getRuleClass();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java b/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java
index 1955202..54d8726 100644
--- a/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java
+++ b/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java
@@ -15,7 +15,6 @@
 package com.google.devtools.build.lib.profiler.memory;
 
 import com.google.common.base.Objects;
-import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.MapMaker;
@@ -25,9 +24,9 @@
 import com.google.devtools.build.lib.packages.AspectClass;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.RuleFunction;
-import com.google.devtools.build.lib.syntax.Callstack;
-import com.google.devtools.build.lib.syntax.Node;
+import com.google.devtools.build.lib.syntax.Debug;
 import com.google.devtools.build.lib.syntax.StarlarkCallable;
+import com.google.devtools.build.lib.syntax.StarlarkThread;
 import com.google.monitoring.runtime.instrumentation.Sampler;
 import com.google.perftools.profiles.ProfileProto.Function;
 import com.google.perftools.profiles.ProfileProto.Line;
@@ -38,7 +37,6 @@
 import java.io.IOException;
 import java.time.Instant;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import java.util.zip.GZIPOutputStream;
@@ -46,18 +44,34 @@
 
 /** Tracks allocations for memory reporting. */
 @ConditionallyThreadCompatible
-public final class AllocationTracker implements Sampler {
+@SuppressWarnings("ThreadLocalUsage") // the AllocationTracker is effectively a global
+public final class AllocationTracker implements Sampler, Debug.ThreadHook {
+
+  // A mapping from Java thread to StarlarkThread.
+  // Used to effect a hidden StarlarkThread parameter to sampleAllocation.
+  // TODO(adonovan): opt: merge the three different ThreadLocals in use here.
+  private final ThreadLocal<StarlarkThread> starlarkThread = new ThreadLocal<>();
+
+  @Override
+  public void onPushFirst(StarlarkThread thread) {
+    starlarkThread.set(thread);
+  }
+
+  @Override
+  public void onPopLast(StarlarkThread thread) {
+    starlarkThread.remove();
+  }
 
   private static class AllocationSample {
     @Nullable final RuleClass ruleClass; // Current rule being analysed, if any
     @Nullable final AspectClass aspectClass; // Current aspect being analysed, if any
-    final List<Object> callstack; // Skylark callstack, if any
+    final ImmutableList<Frame> callstack; // Starlark callstack, if any
     final long bytes;
 
     AllocationSample(
         @Nullable RuleClass ruleClass,
         @Nullable AspectClass aspectClass,
-        List<Object> callstack,
+        ImmutableList<Frame> callstack,
         long bytes) {
       this.ruleClass = ruleClass;
       this.aspectClass = aspectClass;
@@ -66,6 +80,18 @@
     }
   }
 
+  private static class Frame {
+    final String name;
+    final Location loc;
+    @Nullable final RuleFunction ruleFunction;
+
+    Frame(String name, Location loc, @Nullable RuleFunction ruleFunction) {
+      this.name = name;
+      this.loc = loc;
+      this.ruleFunction = ruleFunction;
+    }
+  }
+
   private final Map<Object, AllocationSample> allocations = new MapMaker().weakKeys().makeMap();
   private final int samplePeriod;
   private final int sampleVariance;
@@ -87,19 +113,60 @@
     this.sampleVariance = variance;
   }
 
+  // Called by instrumentation.recordAllocation, which is in turn called
+  // by an instrumented version of the application assembled on the fly
+  // by instrumentation.AllocationInstrumenter.
+  // The instrumenter inserts a call to recordAllocation after every
+  // memory allocation instruction in the original class.
+  //
+  // This function runs within 'new', so is not supposed to allocate memory;
+  // see Sampler interface. In fact it allocates in nearly a dozen places.
+  // TODO(adonovan): suppress reentrant calls by setting a thread-local flag.
   @Override
   @ThreadSafe
   public void sampleAllocation(int count, String desc, Object newObj, long size) {
     if (!enabled) {
       return;
     }
-    List<Object> callstack = Callstack.get();
+
+    @Nullable StarlarkThread thread = starlarkThread.get();
+
+    // Calling Debug.getCallStack is a dubious operation here.
+    // First it allocates memory, which breaks the Sampler contract.
+    // Second, the allocation could in principle occur while the thread's
+    // representation invariants are temporarily broken (that is, during
+    // the call to ArrayList.add when pushing a new stack frame).
+    // For now at least, the allocation done by ArrayList.add occurs before
+    // the representation of the ArrayList is changed, so it is safe,
+    // but this is a fragile assumption.
+    ImmutableList<Debug.Frame> callstack =
+        thread != null ? Debug.getCallStack(thread) : ImmutableList.of();
+
     RuleClass ruleClass = CurrentRuleTracker.getRule();
     AspectClass aspectClass = CurrentRuleTracker.getAspect();
+
     // Should we bother sampling?
     if (callstack.isEmpty() && ruleClass == null && aspectClass == null) {
       return;
     }
+
+    // Convert the thread's stack right away to our internal form.
+    // It is not safe to inspect Debug.Frame references once the thread resumes,
+    // and keeping StarlarkCallable values live defeats garbage collection.
+    ImmutableList.Builder<Frame> frames = ImmutableList.builderWithExpectedSize(callstack.size());
+    for (Debug.Frame fr : callstack) {
+      // The frame's PC location is currently not updated at every step,
+      // only at function calls, so the leaf frame's line number may be
+      // slightly off; see the tests.
+      // TODO(b/149023294): remove comment when we move to a compiled representation.
+      StarlarkCallable fn = fr.getFunction();
+      frames.add(
+          new Frame(
+              fn.getName(),
+              fr.getLocation(),
+              fn instanceof RuleFunction ? (RuleFunction) fn : null));
+    }
+
     // If we start getting stack overflows here, it's because the memory sampling
     // implementation has changed to call back into the sampling method immediately on
     // every allocation. Since thread locals can allocate, this can in this case lead
@@ -113,9 +180,7 @@
     }
     bytesValue.value = 0;
     nextSampleBytes.set(getNextSample());
-    allocations.put(
-        newObj,
-        new AllocationSample(ruleClass, aspectClass, ImmutableList.copyOf(callstack), bytes));
+    allocations.put(newObj, new AllocationSample(ruleClass, aspectClass, frames.build(), bytes));
   }
 
   private long getNextSample() {
@@ -165,13 +230,11 @@
     }
   }
 
+  // If the topmost stack entry is a call to a rule function, returns it.
   @Nullable
-  private static RuleFunction getRuleCreationCall(AllocationSample allocationSample) {
-    Object topOfCallstack = Iterables.getLast(allocationSample.callstack, null);
-    if (topOfCallstack instanceof RuleFunction) {
-      return (RuleFunction) topOfCallstack;
-    }
-    return null;
+  private static RuleFunction getRule(AllocationSample sample) {
+    Frame top = Iterables.getLast(sample.callstack, null);
+    return top != null ? top.ruleFunction : null;
   }
 
   /**
@@ -185,29 +248,28 @@
     System.gc();
 
     // Get loading phase memory for rules.
-    for (AllocationSample allocationSample : allocations.values()) {
-      RuleFunction ruleCreationCall = getRuleCreationCall(allocationSample);
-      if (ruleCreationCall != null) {
-        RuleClass ruleClass = ruleCreationCall.getRuleClass();
+    for (AllocationSample sample : allocations.values()) {
+      RuleFunction rule = getRule(sample);
+      if (rule != null) {
+        RuleClass ruleClass = rule.getRuleClass();
         String key = ruleClass.getKey();
         RuleBytes ruleBytes = rules.computeIfAbsent(key, k -> new RuleBytes(ruleClass.getName()));
-        rules.put(key, ruleBytes.addBytes(allocationSample.bytes));
+        rules.put(key, ruleBytes.addBytes(sample.bytes));
       }
     }
     // Get analysis phase memory for rules and aspects
-    for (AllocationSample allocationSample : allocations.values()) {
-      if (allocationSample.ruleClass != null) {
-        String key = allocationSample.ruleClass.getKey();
+    for (AllocationSample sample : allocations.values()) {
+      if (sample.ruleClass != null) {
+        String key = sample.ruleClass.getKey();
         RuleBytes ruleBytes =
-            rules.computeIfAbsent(key, k -> new RuleBytes(allocationSample.ruleClass.getName()));
-        rules.put(key, ruleBytes.addBytes(allocationSample.bytes));
+            rules.computeIfAbsent(key, k -> new RuleBytes(sample.ruleClass.getName()));
+        rules.put(key, ruleBytes.addBytes(sample.bytes));
       }
-      if (allocationSample.aspectClass != null) {
-        String key = allocationSample.aspectClass.getKey();
+      if (sample.aspectClass != null) {
+        String key = sample.aspectClass.getKey();
         RuleBytes ruleBytes =
-            aspects.computeIfAbsent(
-                key, k -> new RuleBytes(allocationSample.aspectClass.getName()));
-        aspects.put(key, ruleBytes.addBytes(allocationSample.bytes));
+            aspects.computeIfAbsent(key, k -> new RuleBytes(sample.aspectClass.getName()));
+        aspects.put(key, ruleBytes.addBytes(sample.bytes));
       }
     }
 
@@ -237,38 +299,16 @@
             .setType(stringTable.get("memory"))
             .setUnit(stringTable.get("bytes"))
             .build());
-    for (AllocationSample allocationSample : allocations.values()) {
+    for (AllocationSample sample : allocations.values()) {
       // Skip empty callstacks
-      if (allocationSample.callstack.isEmpty()) {
+      if (sample.callstack.isEmpty()) {
         continue;
       }
-      Sample.Builder sample = Sample.newBuilder().addValue(allocationSample.bytes);
-      int line = -1;
-      String file = null;
-      for (int i = allocationSample.callstack.size() - 1; i >= 0; --i) {
-        Object object = allocationSample.callstack.get(i);
-        if (line == -1) {
-          final Location location;
-          if (object instanceof Node) {
-            location = ((Node) object).getStartLocation();
-          } else if (object instanceof StarlarkCallable) {
-            location = ((StarlarkCallable) object).getLocation();
-          } else {
-            throw new IllegalStateException(
-                "Unknown node type: " + object.getClass().getSimpleName());
-          }
-          file = location.file();
-          line = location.line();
-        }
-        if (object instanceof StarlarkCallable) {
-          String function = ((StarlarkCallable) object).getName();
-          sample.addLocationId(
-              locationTable.get(Strings.nullToEmpty(file), Strings.nullToEmpty(function), line));
-          line = -1;
-          file = null;
-        }
+      Sample.Builder b = Sample.newBuilder().addValue(sample.bytes);
+      for (Frame fr : sample.callstack.reverse()) {
+        b.addLocationId(locationTable.get(fr.loc.file(), fr.name, fr.loc.line()));
       }
-      profile.addSample(sample.build());
+      profile.addSample(b.build());
     }
     profile.setTimeNanos(Instant.now().getEpochSecond() * 1000000000);
     return profile.build();
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerModule.java b/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerModule.java
index 7ae9d55..e0d1fd7 100644
--- a/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerModule.java
+++ b/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerModule.java
@@ -21,7 +21,7 @@
 import com.google.devtools.build.lib.runtime.BlazeModule;
 import com.google.devtools.build.lib.runtime.BlazeRuntime;
 import com.google.devtools.build.lib.runtime.WorkspaceBuilder;
-import com.google.devtools.build.lib.syntax.Callstack;
+import com.google.devtools.build.lib.syntax.Debug;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.common.options.OptionsParsingResult;
 import java.util.UUID;
@@ -69,7 +69,7 @@
     enabled = memoryTrackerPropery != null && memoryTrackerPropery.equals("1");
     if (enabled) {
       tracker = new AllocationTracker(SAMPLE_SIZE, VARIANCE);
-      Callstack.setEnabled(true);
+      Debug.setThreadHook(tracker);
       CurrentRuleTracker.setEnabled(true);
       AllocationTrackerInstaller.installAllocationTracker(tracker);
     }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/BUILD b/src/main/java/com/google/devtools/build/lib/syntax/BUILD
index 82dddf8..2aac73b 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/syntax/BUILD
@@ -92,7 +92,6 @@
         "BaseFunction.java",
         "BuiltinCallable.java",
         "CallUtils.java",
-        "Callstack.java",
         "ClassObject.java",
         "CpuProfiler.java",
         "Debug.java",
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Callstack.java b/src/main/java/com/google/devtools/build/lib/syntax/Callstack.java
deleted file mode 100644
index 3c92f8b..0000000
--- a/src/main/java/com/google/devtools/build/lib/syntax/Callstack.java
+++ /dev/null
@@ -1,63 +0,0 @@
-// Copyright 2017 The Bazel Authors. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//    http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package com.google.devtools.build.lib.syntax;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Preconditions;
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * Holds the Skylark callstack in thread-local storage. Contains all Expressions and functions
- * currently being evaluated.
- *
- * <p>This is needed for memory tracking, since the evaluator is not available in the context of
- * instrumentation. It should not be used by normal Skylark interpreter logic.
- */
-public class Callstack {
-  // This field should be read, but not written, directly in order to avoid the overhead of a
-  // method call.
-  static boolean enabled;
-  private static final ThreadLocal<List<Object>> callstack =
-      ThreadLocal.withInitial(ArrayList::new);
-
-  public static void setEnabled(boolean enabled) {
-    Callstack.enabled = enabled;
-  }
-
-  public static void push(Node node) {
-    callstack.get().add(node);
-  }
-
-  public static void push(StarlarkCallable fn) {
-    callstack.get().add(fn);
-  }
-
-  public static void pop() {
-    List<Object> threadStack = callstack.get();
-    threadStack.remove(threadStack.size() - 1);
-  }
-
-  public static List<Object> get() {
-    Preconditions.checkState(enabled, "Must call Callstack#setEnabled before getting");
-    return callstack.get();
-  }
-
-  @VisibleForTesting
-  public static void resetStateForTest() {
-    enabled = false;
-    callstack.get().clear();
-  }
-}
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Debug.java b/src/main/java/com/google/devtools/build/lib/syntax/Debug.java
index f91066e..a24362a 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Debug.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Debug.java
@@ -62,4 +62,25 @@
     /** Returns the local environment of this frame. */
     ImmutableMap<String, Object> getLocals();
   }
+
+  /**
+   * Interface by which debugging tools are notified of a thread entering or leaving its top-level
+   * frame.
+   */
+  public interface ThreadHook {
+    void onPushFirst(StarlarkThread thread);
+
+    void onPopLast(StarlarkThread thread);
+  }
+
+  static ThreadHook threadHook = null;
+
+  /**
+   * Installs a global hook that is notified each time a thread pushes or pops its top-level frame.
+   * This interface is provided to support special tools; ordinary clients should have no need for
+   * it.
+   */
+  public static void setThreadHook(ThreadHook hook) {
+    threadHook = hook;
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Eval.java b/src/main/java/com/google/devtools/build/lib/syntax/Eval.java
index 1643e85..6cb3133 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Eval.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Eval.java
@@ -421,23 +421,10 @@
 
   private static Object eval(StarlarkThread.Frame fr, Expression expr)
       throws EvalException, InterruptedException {
-    // TODO(adonovan): don't push and pop all the time. We should only need the stack of function
-    // call frames, and we should recycle them.
-    // TODO(adonovan): put the StarlarkThread into the Java thread-local store
-    // once only, in push, and undo this in pop.
     try {
-      if (Callstack.enabled) {
-        Callstack.push(expr);
-      }
-      try {
-        return doEval(fr, expr);
-      } catch (EvalException ex) {
-        throw maybeTransformException(expr, ex);
-      }
-    } finally {
-      if (Callstack.enabled) {
-        Callstack.pop();
-      }
+      return doEval(fr, expr);
+    } catch (EvalException ex) {
+      throw maybeTransformException(expr, ex);
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/StarlarkThread.java b/src/main/java/com/google/devtools/build/lib/syntax/StarlarkThread.java
index 6c2dbad..b354107 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/StarlarkThread.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/StarlarkThread.java
@@ -367,10 +367,9 @@
     Frame fr = new Frame(this, fn);
     callstack.add(fr);
 
-    // Push the function onto the allocation tracker's stack.
-    // TODO(adonovan): optimize it out of existence.
-    if (Callstack.enabled) {
-      Callstack.push(fn);
+    // Notify debug tools of the thread's first push.
+    if (callstack.size() == 1 && Debug.threadHook != null) {
+      Debug.threadHook.onPushFirst(this);
     }
 
     ProfilerTask taskKind;
@@ -432,8 +431,9 @@
       fr.profileSpan.close();
     }
 
-    if (Callstack.enabled) {
-      Callstack.pop();
+    // Notify debug tools of the thread's last pop.
+    if (last == 0 && Debug.threadHook != null) {
+      Debug.threadHook.onPopLast(this);
     }
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerTest.java b/src/test/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerTest.java
index 89c2e9f..a24052c 100644
--- a/src/test/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/profiler/memory/AllocationTrackerTest.java
@@ -17,22 +17,31 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-import com.google.devtools.build.lib.events.Location;
+import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.RuleFunction;
 import com.google.devtools.build.lib.profiler.memory.AllocationTracker.RuleBytes;
-import com.google.devtools.build.lib.syntax.Callstack;
-import com.google.devtools.build.lib.syntax.Expression;
+import com.google.devtools.build.lib.syntax.Debug;
+import com.google.devtools.build.lib.syntax.EvalException;
+import com.google.devtools.build.lib.syntax.EvalUtils;
+import com.google.devtools.build.lib.syntax.HasBinary;
+import com.google.devtools.build.lib.syntax.Module;
+import com.google.devtools.build.lib.syntax.Mutability;
 import com.google.devtools.build.lib.syntax.ParserInput;
+import com.google.devtools.build.lib.syntax.Starlark;
 import com.google.devtools.build.lib.syntax.StarlarkCallable;
+import com.google.devtools.build.lib.syntax.StarlarkThread;
 import com.google.devtools.build.lib.syntax.SyntaxError;
+import com.google.devtools.build.lib.syntax.TokenKind;
 import com.google.perftools.profiles.ProfileProto.Function;
 import com.google.perftools.profiles.ProfileProto.Profile;
 import com.google.perftools.profiles.ProfileProto.Sample;
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.List;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -43,156 +52,84 @@
 @RunWith(JUnit4.class)
 public final class AllocationTrackerTest {
 
-  private AllocationTracker allocationTracker;
+  // These tests are quite artificial as they call sampleAllocation explicitly.
+  // In reality, a call could occur after any 'new' operation.
 
-  // makeExpr returns an expression located at the specified file and line.
-  private static Expression makeExpr(String file, int line) throws SyntaxError {
-    ParserInput input =
-        ParserInput.create(new String(new char[line - 1]).replace('\u0000', '\n') + "None", file);
-    return Expression.parse(input);
-  }
+  private AllocationTracker tracker;
+  private final ArrayList<Object> live = new ArrayList<>();
 
-  private static class TestFunction implements StarlarkCallable {
-    private final String name;
-    private final Location location;
-
-    TestFunction(String file, String name, int line) {
-      this.name = name;
-      this.location = Location.fromFileLineColumn(file, line, 0);
-    }
-
+  // A Starlark value whose plus operator "x + 123" simulates allocation of 123 bytes.
+  // (We trigger allocation with an operator not a function call so as not to change the stack.)
+  private class SamplerValue implements HasBinary {
     @Override
-    public String getName() {
-      return name;
-    }
-
-    @Override
-    public Location getLocation() {
-      return location;
+    public Object binaryOp(TokenKind op, Object that, boolean thisLeft) throws EvalException {
+      if (op == TokenKind.PLUS && thisLeft && that instanceof Integer) {
+        int size = (Integer) that;
+        Object obj = new Object();
+        live.add(obj); // ensure that obj outlives the test assertions
+        tracker.sampleAllocation(1, "", obj, size);
+        return Starlark.NONE;
+      }
+      return null;
     }
   }
 
-  static class TestRuleFunction extends TestFunction implements RuleFunction {
-
-    private final RuleClass ruleClass;
-
-    TestRuleFunction(String file, String name, int line) {
-      super(file, name, line);
-      this.ruleClass = mock(RuleClass.class);
-      when(ruleClass.getName()).thenReturn(name);
-      when(ruleClass.getKey()).thenReturn(name);
-    }
-
-    @Override
-    public RuleClass getRuleClass() {
-      return ruleClass;
-    }
+  private static RuleClass myRuleClass() {
+    RuleClass myrule = mock(RuleClass.class);
+    when(myrule.getName()).thenReturn("myrule");
+    when(myrule.getKey()).thenReturn("myrule");
+    return myrule;
   }
 
   @Before
   public void setup() {
-    Callstack.setEnabled(true);
     CurrentRuleTracker.setEnabled(true);
-    allocationTracker = new AllocationTracker(1, 0);
+    tracker = new AllocationTracker(1, 0);
+    Debug.setThreadHook(tracker);
   }
 
   @After
   public void tearDown() {
-    Callstack.resetStateForTest();
+    Debug.setThreadHook(null);
     CurrentRuleTracker.setEnabled(false);
   }
 
   @Test
-  public void testSimpleMemoryProfile() throws Exception {
-    Object allocation = new Object();
-    Callstack.push(new TestFunction("fileA", "fn", 120));
-    Callstack.push(makeExpr("fileA", 10));
-    allocationTracker.sampleAllocation(1, "", allocation, 12);
-    Callstack.pop();
-    Callstack.pop();
+  public void testMemoryProfileDuringExecution() throws Exception {
+    // The nop() calls force the frame PC location to be updated.
+    // It is not updated for a + operation on the assumption that
+    // the stack is unobservable to an implementation of the +
+    // operator... but the AllocationTracker sneaks a peek at it
+    // using thread-local storage.
+    // TODO(b/149023294): update this when we use a compiled representation.
+    exec(
+        "def nop(): pass",
+        "def g():",
+        "  nop(); sample + 12", // sample[0]: 12 bytes
+        "def f():",
+        "  g()",
+        "  nop(); sample + 73", // sample[1]: 73 bytes
+        "f()");
 
     Map<String, RuleBytes> rules = new HashMap<>();
     Map<String, RuleBytes> aspects = new HashMap<>();
-    allocationTracker.getRuleMemoryConsumption(rules, aspects);
+    tracker.getRuleMemoryConsumption(rules, aspects);
     assertThat(rules).isEmpty();
     assertThat(aspects).isEmpty();
 
-    Profile profile = allocationTracker.buildMemoryProfile();
-    assertThat(profile.getSampleList()).hasSize(1);
-    assertThat(sampleToCallstack(profile, profile.getSample(0))).containsExactly("fileA:fn:10");
-  }
-
-  @Test
-  public void testLongerCallstack() throws Exception {
-    Object allocation = new Object();
-    Callstack.push(new TestFunction("fileB", "fnB", 120));
-    Callstack.push(makeExpr("fileB", 10));
-    Callstack.push(makeExpr("fileB", 12));
-    Callstack.push(makeExpr("fileB", 14));
-    Callstack.push(makeExpr("fileB", 18));
-    Callstack.push(new TestFunction("fileA", "fnA", 120));
-    Callstack.push(makeExpr("fileA", 10));
-    allocationTracker.sampleAllocation(1, "", allocation, 12);
-    for (int i = 0; i < 7; ++i) {
-      Callstack.pop();
+    Profile profile = tracker.buildMemoryProfile();
+    assertThat(profile.getSampleList()).hasSize(2);
+    Set<String> lines = new HashSet<>();
+    for (Sample s : profile.getSampleList()) {
+      lines.add(sampleToCallstack(profile, s));
     }
-
-    Profile profile = allocationTracker.buildMemoryProfile();
-    assertThat(profile.getSampleList()).hasSize(1);
-    assertThat(sampleToCallstack(profile, profile.getSample(0)))
-        .containsExactly("fileB:fnB:18", "fileA:fnA:10");
+    assertThat(lines).contains("a.star:f:6, a.star:<toplevel>:7");
+    assertThat(lines).contains("a.star:g:3, a.star:f:5, a.star:<toplevel>:7");
   }
 
-  @Test
-  public void testConfiguredTargetsMemoryAllocation() throws Exception {
-    RuleClass ruleClass = mock(RuleClass.class);
-    when(ruleClass.getName()).thenReturn("rule");
-    when(ruleClass.getKey()).thenReturn("rule");
-    CurrentRuleTracker.beginConfiguredTarget(ruleClass);
-    Object ruleAllocation0 = new Object();
-    Object ruleAllocation1 = new Object();
-    allocationTracker.sampleAllocation(1, "", ruleAllocation0, 10);
-    allocationTracker.sampleAllocation(1, "", ruleAllocation1, 20);
-    CurrentRuleTracker.endConfiguredTarget();
-
-    CurrentRuleTracker.beginConfiguredAspect(() -> "aspect");
-    Object aspectAllocation = new Object();
-    allocationTracker.sampleAllocation(1, "", aspectAllocation, 12);
-    CurrentRuleTracker.endConfiguredAspect();
-
-    Map<String, RuleBytes> rules = new HashMap<>();
-    Map<String, RuleBytes> aspects = new HashMap<>();
-    allocationTracker.getRuleMemoryConsumption(rules, aspects);
-    assertThat(rules).containsExactly("rule", new RuleBytes("rule").addBytes(30L));
-    assertThat(aspects).containsExactly("aspect", new RuleBytes("aspect").addBytes(12L));
-
-    Profile profile = allocationTracker.buildMemoryProfile();
-    assertThat(profile.getSampleList()).isEmpty(); // No callstacks
-  }
-
-  @Test
-  public void testLoadingPhaseRuleAllocations() throws Exception {
-    Object allocation = new Object();
-    Callstack.push(new TestFunction("fileB", "fnB", 120));
-    Callstack.push(makeExpr("fileB", 18));
-    Callstack.push(new TestFunction("fileA", "fnA", 120));
-    Callstack.push(makeExpr("fileA", 10));
-    Callstack.push(new TestRuleFunction("<native>", "proto_library", -1));
-    allocationTracker.sampleAllocation(1, "", allocation, 128);
-    for (int i = 0; i < 5; ++i) {
-      Callstack.pop();
-    }
-
-    Map<String, RuleBytes> rules = new HashMap<>();
-    Map<String, RuleBytes> aspects = new HashMap<>();
-    allocationTracker.getRuleMemoryConsumption(rules, aspects);
-    assertThat(rules)
-        .containsExactly("proto_library", new RuleBytes("proto_library").addBytes(128L));
-  }
-
-  /** Formats a callstack as (file):(method name):(line) */
-  private List<String> sampleToCallstack(Profile profile, Sample sample) {
-    List<String> result = new ArrayList<>();
+  /** Formats a call stack as a comma-separated list of file:function:line elements. */
+  private static String sampleToCallstack(Profile profile, Sample sample) {
+    StringBuilder buf = new StringBuilder();
     for (long locationId : sample.getLocationIdList()) {
       com.google.perftools.profiles.ProfileProto.Location location =
           profile.getLocation((int) locationId - 1);
@@ -204,8 +141,86 @@
       long methodId = function.getName();
       String file = profile.getStringTable((int) fileId);
       String method = profile.getStringTable((int) methodId);
-      result.add(String.format("%s:%s:%d", file, method, line));
+      if (buf.length() > 0) {
+        buf.append(", ");
+      }
+      buf.append(String.format("%s:%s:%d", file, method, line));
     }
-    return result;
+    return buf.toString();
   }
+
+  @Test
+  public void testConfiguredTargetsMemoryAllocation() throws Exception {
+    CurrentRuleTracker.beginConfiguredTarget(myRuleClass());
+    Object ruleAllocation0 = new Object();
+    Object ruleAllocation1 = new Object();
+    tracker.sampleAllocation(1, "", ruleAllocation0, 10);
+    tracker.sampleAllocation(1, "", ruleAllocation1, 20);
+    CurrentRuleTracker.endConfiguredTarget();
+
+    CurrentRuleTracker.beginConfiguredAspect(() -> "aspect");
+    Object aspectAllocation = new Object();
+    tracker.sampleAllocation(1, "", aspectAllocation, 12);
+    CurrentRuleTracker.endConfiguredAspect();
+
+    Map<String, RuleBytes> rules = new HashMap<>();
+    Map<String, RuleBytes> aspects = new HashMap<>();
+    tracker.getRuleMemoryConsumption(rules, aspects);
+    assertThat(rules).containsExactly("myrule", new RuleBytes("myrule").addBytes(30L));
+    assertThat(aspects).containsExactly("aspect", new RuleBytes("aspect").addBytes(12L));
+
+    Profile profile = tracker.buildMemoryProfile();
+    assertThat(profile.getSampleList()).isEmpty(); // no callstacks
+  }
+
+  @Test
+  public void testLoadingPhaseRuleAllocations() throws Exception {
+    exec(
+        "def g():", //
+        "  myrule()",
+        "def f():",
+        "  g()",
+        "f()");
+    Map<String, RuleBytes> rules = new HashMap<>();
+    Map<String, RuleBytes> aspects = new HashMap<>();
+    tracker.getRuleMemoryConsumption(rules, aspects);
+    assertThat(rules).containsExactly("myrule", new RuleBytes("myrule").addBytes(128L));
+  }
+
+  private void exec(String... lines) throws SyntaxError, EvalException, InterruptedException {
+    Mutability mu = Mutability.create("test");
+    StarlarkThread thread =
+        StarlarkThread.builder(mu)
+            .useDefaultSemantics()
+            .setGlobals(
+                Module.createForBuiltins(
+                    ImmutableMap.of(
+                        "sample", new SamplerValue(),
+                        "myrule", new MyRuleFunction())))
+            .build();
+    ParserInput input = ParserInput.create(Joiner.on("\n").join(lines), "a.star");
+    EvalUtils.exec(input, thread);
+  }
+
+  // A fake Bazel rule. The allocation tracker reports retained memory broken down by rule class.
+  private class MyRuleFunction implements RuleFunction, StarlarkCallable {
+    @Override
+    public Object fastcall(StarlarkThread thread, Object[] parameters, Object[] named) {
+      Object obj = new Object();
+      live.add(obj); // ensure that obj outlives the test assertions
+      tracker.sampleAllocation(1, "", obj, 128);
+      return Starlark.NONE;
+    }
+
+    @Override
+    public String getName() {
+      return "myrule";
+    }
+
+    @Override
+    public RuleClass getRuleClass() {
+      return myRuleClass();
+    }
+  }
+
 }