[vm] Make heap snapshot writer also visit isolate stacks as roots

Heap snapshots currently produced don't visit isolate stacks. As such
analyzing such snapshots may lead one to conclude there is a lot of
garbage while objects are actually reachable.

=> This CL makes us visit isolate stacks when building heap snapshots.

Furthermore we add a new `VMInternals.writeHeapSnapshotToFile` helper
that can be used to programmatically write snapshots and can be handy
for internal use at times. (We also use this helper in a test)

TEST=vm/dart{,_2}/heap_snapshot_test

Change-Id: I976544b7f6d20863764af9a40bf1ffb3c319bbce
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/253785
Reviewed-by: Ryan Macnak <rmacnak@google.com>
Commit-Queue: Martin Kustermann <kustermann@google.com>

runtime/lib/object.cc

@@ -11,6 +11,7 @@
 #include "vm/heap/heap.h"
 #include "vm/native_entry.h"
 #include "vm/object.h"
+#include "vm/object_graph.h"
 #include "vm/object_store.h"
 #include "vm/resolver.h"
 #include "vm/stack_frame.h"
@@ -313,6 +314,22 @@ DEFINE_NATIVE_ENTRY(Internal_collectAllGarbage, 0, 0) {
   return Object::null();
 }
 
+DEFINE_NATIVE_ENTRY(Internal_writeHeapSnapshotToFile, 0, 1) {
+#if !defined(PRODUCT)
+  const String& filename =
+      String::CheckedHandle(zone, arguments->NativeArgAt(0));
+  {
+    FileHeapSnapshotWriter file_writer(thread, filename.ToCString());
+    HeapSnapshotWriter writer(thread, &file_writer);
+    writer.Write();
+  }
+#else
+  Exceptions::ThrowUnsupportedError(
+      "Heap snapshots are only supported in non-product mode.");
+#endif  // !defined(PRODUCT)
+  return Object::null();
+}
+
 DEFINE_NATIVE_ENTRY(Internal_deoptimizeFunctionsOnStack, 0, 0) {
   DeoptimizeFunctionsOnStack();
   return Object::null();

runtime/tests/vm/dart/heap_snapshot_test.dart

@@ -0,0 +1,108 @@
+// Copyright (c) 2022, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'dart:io';
+import 'dart:_internal';
+
+import 'package:expect/expect.dart';
+import 'package:path/path.dart' as path;
+import 'package:vm_service/vm_service.dart';
+
+import 'use_flag_test_helper.dart';
+
+final bool alwaysTrue = int.parse('1') == 1;
+
+@pragma('vm:entry-point') // Prevent name mangling
+class Foo {}
+
+var global = null;
+
+main() async {
+  if (const bool.fromEnvironment('dart.vm.product')) {
+    var exception;
+    try {
+      await runTest();
+    } catch (e) {
+      exception = e;
+    }
+    Expect.contains(
+        'Heap snapshots are only supported in non-product mode.', '$exception');
+    return;
+  }
+
+  await runTest();
+}
+
+Future runTest() async {
+  await withTempDir('heap_snapshot_test', (String dir) async {
+    final state1 = path.join(dir, 'state1.heapsnapshot');
+    final state2 = path.join(dir, 'state2.heapsnapshot');
+    final state3 = path.join(dir, 'state3.heapsnapshot');
+
+    var local;
+
+    VMInternalsForTesting.writeHeapSnapshotToFile(state1);
+    if (alwaysTrue) {
+      global = Foo();
+      local = Foo();
+    }
+    VMInternalsForTesting.writeHeapSnapshotToFile(state2);
+    if (alwaysTrue) {
+      global = null;
+      local = null;
+    }
+    VMInternalsForTesting.writeHeapSnapshotToFile(state3);
+
+    final int count1 = countFooInstances(
+        findReachableObjects(loadHeapSnapshotFromFile(state1)));
+    final int count2 = countFooInstances(
+        findReachableObjects(loadHeapSnapshotFromFile(state2)));
+    final int count3 = countFooInstances(
+        findReachableObjects(loadHeapSnapshotFromFile(state3)));
+
+    Expect.equals(0, count1);
+    Expect.equals(2, count2);
+    Expect.equals(0, count3);
+
+    reachabilityFence(local);
+    reachabilityFence(global);
+  });
+}
+
+HeapSnapshotGraph loadHeapSnapshotFromFile(String filename) {
+  final bytes = File(filename).readAsBytesSync();
+  return HeapSnapshotGraph.fromChunks([bytes.buffer.asByteData()]);
+}
+
+Set<HeapSnapshotObject> findReachableObjects(HeapSnapshotGraph graph) {
+  const int rootObjectIdx = 1;
+
+  final reachableObjects = Set<HeapSnapshotObject>();
+  final worklist = <HeapSnapshotObject>[];
+
+  final rootObject = graph.objects[rootObjectIdx];
+
+  reachableObjects.add(rootObject);
+  worklist.add(rootObject);
+
+  while (worklist.isNotEmpty) {
+    final objectToExpand = worklist.removeLast();
+
+    for (final successor in objectToExpand.successors) {
+      if (!reachableObjects.contains(successor)) {
+        reachableObjects.add(successor);
+        worklist.add(successor);
+      }
+    }
+  }
+  return reachableObjects;
+}
+
+int countFooInstances(Set<HeapSnapshotObject> reachableObjects) {
+  int count = 0;
+  for (final object in reachableObjects) {
+    if (object.klass.name == 'Foo') count++;
+  }
+  return count;
+}

runtime/tests/vm/dart_2/heap_snapshot_test.dart

@@ -0,0 +1,110 @@
+// Copyright (c) 2022, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+// @dart=2.9
+
+import 'dart:io';
+import 'dart:_internal';
+
+import 'package:expect/expect.dart';
+import 'package:path/path.dart' as path;
+import 'package:vm_service/vm_service.dart';
+
+import 'use_flag_test_helper.dart';
+
+final bool alwaysTrue = int.parse('1') == 1;
+
+@pragma('vm:entry-point') // Prevent name mangling
+class Foo {}
+
+var global = null;
+
+main() async {
+  if (const bool.fromEnvironment('dart.vm.product')) {
+    var exception;
+    try {
+      await runTest();
+    } catch (e) {
+      exception = e;
+    }
+    Expect.contains(
+        'Heap snapshots are only supported in non-product mode.', '$exception');
+    return;
+  }
+
+  await runTest();
+}
+
+Future runTest() async {
+  await withTempDir('heap_snapshot_test', (String dir) async {
+    final state1 = path.join(dir, 'state1.heapsnapshot');
+    final state2 = path.join(dir, 'state2.heapsnapshot');
+    final state3 = path.join(dir, 'state3.heapsnapshot');
+
+    var local;
+
+    VMInternalsForTesting.writeHeapSnapshotToFile(state1);
+    if (alwaysTrue) {
+      global = Foo();
+      local = Foo();
+    }
+    VMInternalsForTesting.writeHeapSnapshotToFile(state2);
+    if (alwaysTrue) {
+      global = null;
+      local = null;
+    }
+    VMInternalsForTesting.writeHeapSnapshotToFile(state3);
+
+    final int count1 = countFooInstances(
+        findReachableObjects(loadHeapSnapshotFromFile(state1)));
+    final int count2 = countFooInstances(
+        findReachableObjects(loadHeapSnapshotFromFile(state2)));
+    final int count3 = countFooInstances(
+        findReachableObjects(loadHeapSnapshotFromFile(state3)));
+
+    Expect.equals(0, count1);
+    Expect.equals(2, count2);
+    Expect.equals(0, count3);
+
+    reachabilityFence(local);
+    reachabilityFence(global);
+  });
+}
+
+HeapSnapshotGraph loadHeapSnapshotFromFile(String filename) {
+  final bytes = File(filename).readAsBytesSync();
+  return HeapSnapshotGraph.fromChunks([bytes.buffer.asByteData()]);
+}
+
+Set<HeapSnapshotObject> findReachableObjects(HeapSnapshotGraph graph) {
+  const int rootObjectIdx = 1;
+
+  final reachableObjects = Set<HeapSnapshotObject>();
+  final worklist = <HeapSnapshotObject>[];
+
+  final rootObject = graph.objects[rootObjectIdx];
+
+  reachableObjects.add(rootObject);
+  worklist.add(rootObject);
+
+  while (worklist.isNotEmpty) {
+    final objectToExpand = worklist.removeLast();
+
+    for (final successor in objectToExpand.successors) {
+      if (!reachableObjects.contains(successor)) {
+        reachableObjects.add(successor);
+        worklist.add(successor);
+      }
+    }
+  }
+  return reachableObjects;
+}
+
+int countFooInstances(Set<HeapSnapshotObject> reachableObjects) {
+  int count = 0;
+  for (final object in reachableObjects) {
+    if (object.klass.name == 'Foo') count++;
+  }
+  return count;
+}

runtime/vm/bootstrap_natives.h

@@ -319,6 +319,7 @@ namespace dart {
   V(Internal_unsafeCast, 1)                                                    \
   V(Internal_nativeEffect, 1)                                                  \
   V(Internal_collectAllGarbage, 0)                                             \
+  V(Internal_writeHeapSnapshotToFile, 1)                                       \
   V(Internal_makeListFixedLength, 1)                                           \
   V(Internal_makeFixedListUnmodifiable, 1)                                     \
   V(Internal_extractTypeArguments, 2)                                          \

runtime/vm/isolate.cc

@@ -2825,6 +2825,13 @@ void Isolate::VisitObjectPointers(ObjectPointerVisitor* visitor,
   }
 }
 
+void Isolate::VisitStackPointers(ObjectPointerVisitor* visitor,
+                                 ValidationPolicy validate_frames) {
+  if (mutator_thread_ != nullptr) {
+    mutator_thread_->VisitObjectPointers(visitor, validate_frames);
+  }
+}
+
 void IsolateGroup::ReleaseStoreBuffers() {
   thread_registry()->ReleaseStoreBuffers();
 }
@@ -3016,9 +3023,7 @@ void IsolateGroup::VisitStackPointers(ObjectPointerVisitor* visitor,
   for (Isolate* isolate : isolates_) {
     // Visit mutator thread, even if the isolate isn't entered/scheduled
     // (there might be live API handles to visit).
-    if (isolate->mutator_thread_ != nullptr) {
-      isolate->mutator_thread_->VisitObjectPointers(visitor, validate_frames);
-    }
+    isolate->VisitStackPointers(visitor, validate_frames);
   }
 
   visitor->clear_gc_root_type();

runtime/vm/object_graph.cc

@@ -656,11 +656,12 @@ void HeapSnapshotWriter::EnsureAvailable(intptr_t needed) {
   ASSERT(buffer_ == nullptr);
 
   intptr_t chunk_size = kPreferredChunkSize;
-  if (chunk_size < needed + kMetadataReservation) {
-    chunk_size = needed + kMetadataReservation;
+  const intptr_t reserved_prefix = writer_->ReserveChunkPrefixSize();
+  if (chunk_size < (reserved_prefix + needed)) {
+    chunk_size = reserved_prefix + needed;
   }
   buffer_ = reinterpret_cast<uint8_t*>(malloc(chunk_size));
-  size_ = kMetadataReservation;
+  size_ = reserved_prefix;
   capacity_ = chunk_size;
 }
 
@@ -669,28 +670,8 @@ void HeapSnapshotWriter::Flush(bool last) {
     return;
   }
 
-  JSONStream js;
-  {
-    JSONObject jsobj(&js);
-    jsobj.AddProperty("jsonrpc", "2.0");
-    jsobj.AddProperty("method", "streamNotify");
-    {
-      JSONObject params(&jsobj, "params");
-      params.AddProperty("streamId", Service::heapsnapshot_stream.id());
-      {
-        JSONObject event(&params, "event");
-        event.AddProperty("type", "Event");
-        event.AddProperty("kind", "HeapSnapshot");
-        event.AddProperty("isolate", thread()->isolate());
-        event.AddPropertyTimeMillis("timestamp", OS::GetCurrentTimeMillis());
-        event.AddProperty("last", last);
-      }
-    }
-  }
+  writer_->WriteChunk(buffer_, size_, last);
 
-  Service::SendEventWithData(Service::heapsnapshot_stream.id(), "HeapSnapshot",
-                             kMetadataReservation, js.buffer()->buffer(),
-                             js.buffer()->length(), buffer_, size_);
   buffer_ = nullptr;
   size_ = 0;
   capacity_ = 0;
@@ -1200,6 +1181,60 @@ class CollectStaticFieldNames : public ObjectVisitor {
   DISALLOW_COPY_AND_ASSIGN(CollectStaticFieldNames);
 };
 
+void VmServiceHeapSnapshotChunkedWriter::WriteChunk(uint8_t* buffer,
+                                                    intptr_t size,
+                                                    bool last) {
+  JSONStream js;
+  {
+    JSONObject jsobj(&js);
+    jsobj.AddProperty("jsonrpc", "2.0");
+    jsobj.AddProperty("method", "streamNotify");
+    {
+      JSONObject params(&jsobj, "params");
+      params.AddProperty("streamId", Service::heapsnapshot_stream.id());
+      {
+        JSONObject event(&params, "event");
+        event.AddProperty("type", "Event");
+        event.AddProperty("kind", "HeapSnapshot");
+        event.AddProperty("isolate", thread()->isolate());
+        event.AddPropertyTimeMillis("timestamp", OS::GetCurrentTimeMillis());
+        event.AddProperty("last", last);
+      }
+    }
+  }
+
+  Service::SendEventWithData(Service::heapsnapshot_stream.id(), "HeapSnapshot",
+                             kMetadataReservation, js.buffer()->buffer(),
+                             js.buffer()->length(), buffer, size);
+}
+
+FileHeapSnapshotWriter::FileHeapSnapshotWriter(Thread* thread,
+                                               const char* filename)
+    : ChunkedWriter(thread) {
+  auto open = Dart::file_open_callback();
+  if (open != nullptr) {
+    file_ = open(filename, /*write=*/true);
+  }
+}
+FileHeapSnapshotWriter::~FileHeapSnapshotWriter() {
+  auto close = Dart::file_close_callback();
+  if (close != nullptr) {
+    close(file_);
+  }
+}
+
+void FileHeapSnapshotWriter::WriteChunk(uint8_t* buffer,
+                                        intptr_t size,
+                                        bool last) {
+  if (file_ != nullptr) {
+    auto write = Dart::file_write_callback();
+    if (write != nullptr) {
+      write(buffer, size, file_);
+    }
+  }
+  free(buffer);
+}
+
 void HeapSnapshotWriter::Write() {
   HeapIterationScope iteration(thread());
 
@@ -1393,6 +1428,8 @@ void HeapSnapshotWriter::Write() {
             ++object_count_;
             isolate->VisitObjectPointers(&visitor,
                                          ValidationPolicy::kDontValidateFrames);
+            isolate->VisitStackPointers(&visitor,
+                                        ValidationPolicy::kDontValidateFrames);
             ++num_isolates;
           },
           /*at_safepoint=*/true);
@@ -1449,9 +1486,13 @@ void HeapSnapshotWriter::Write() {
           visitor.DoCount();
           isolate->VisitObjectPointers(&visitor,
                                        ValidationPolicy::kDontValidateFrames);
+          isolate->VisitStackPointers(&visitor,
+                                      ValidationPolicy::kDontValidateFrames);
           visitor.DoWrite();
           isolate->VisitObjectPointers(&visitor,
                                        ValidationPolicy::kDontValidateFrames);
+          isolate->VisitStackPointers(&visitor,
+                                      ValidationPolicy::kDontValidateFrames);
         },
         /*at_safepoint=*/true);
 

runtime/vm/object_graph.h

@@ -117,11 +117,45 @@ class ObjectGraph : public ThreadStackResource {
   DISALLOW_IMPLICIT_CONSTRUCTORS(ObjectGraph);
 };
 
+class ChunkedWriter : public ThreadStackResource {
+ public:
+  explicit ChunkedWriter(Thread* thread) : ThreadStackResource(thread) {}
+
+  virtual intptr_t ReserveChunkPrefixSize() { return 0; }
+
+  // Takes ownership of [buffer], must be freed with [malloc].
+  virtual void WriteChunk(uint8_t* buffer, intptr_t size, bool last) = 0;
+};
+
+class FileHeapSnapshotWriter : public ChunkedWriter {
+ public:
+  FileHeapSnapshotWriter(Thread* thread, const char* filename);
+  ~FileHeapSnapshotWriter();
+
+  virtual void WriteChunk(uint8_t* buffer, intptr_t size, bool last);
+
+ private:
+  void* file_ = nullptr;
+};
+
+class VmServiceHeapSnapshotChunkedWriter : public ChunkedWriter {
+ public:
+  explicit VmServiceHeapSnapshotChunkedWriter(Thread* thread)
+      : ChunkedWriter(thread) {}
+
+  virtual intptr_t ReserveChunkPrefixSize() { return kMetadataReservation; }
+  virtual void WriteChunk(uint8_t* buffer, intptr_t size, bool last);
+
+ private:
+  static const intptr_t kMetadataReservation = 512;
+};
+
 // Generates a dump of the heap, whose format is described in
 // runtime/vm/service/heap_snapshot.md.
 class HeapSnapshotWriter : public ThreadStackResource {
  public:
-  explicit HeapSnapshotWriter(Thread* thread) : ThreadStackResource(thread) {}
+  HeapSnapshotWriter(Thread* thread, ChunkedWriter* writer)
+      : ThreadStackResource(thread), writer_(writer) {}
 
   void WriteSigned(int64_t value) {
     EnsureAvailable((sizeof(value) * kBitsPerByte) / 7 + 1);
@@ -191,7 +225,6 @@ class HeapSnapshotWriter : public ThreadStackResource {
  private:
   static uint32_t GetHashHelper(Thread* thread, ObjectPtr obj);
 
-  static const intptr_t kMetadataReservation = 512;
   static const intptr_t kPreferredChunkSize = MB;
 
   void SetupCountingPages();
@@ -201,6 +234,8 @@ class HeapSnapshotWriter : public ThreadStackResource {
   void EnsureAvailable(intptr_t needed);
   void Flush(bool last = false);
 
+  ChunkedWriter* writer_ = nullptr;
+
   uint8_t* buffer_ = nullptr;
   intptr_t size_ = 0;
   intptr_t capacity_ = 0;

runtime/vm/service.cc

@@ -4545,7 +4545,8 @@ static const MethodParameter* const request_heap_snapshot_params[] = {
 
 static void RequestHeapSnapshot(Thread* thread, JSONStream* js) {
   if (Service::heapsnapshot_stream.enabled()) {
-    HeapSnapshotWriter writer(thread);
+    VmServiceHeapSnapshotChunkedWriter vmservice_writer(thread);
+    HeapSnapshotWriter writer(thread, &vmservice_writer);
     writer.Write();
   }
   // TODO(koda): Provide some id that ties this request to async response(s).

sdk/lib/_internal/vm/lib/internal_patch.dart

@@ -177,6 +177,9 @@ abstract class VMInternalsForTesting {
   @pragma("vm:external-name", "Internal_collectAllGarbage")
   external static void collectAllGarbage();
 
+  @pragma("vm:external-name", "Internal_writeHeapSnapshotToFile")
+  external static void writeHeapSnapshotToFile(String filename);
+
   @pragma("vm:external-name", "Internal_deoptimizeFunctionsOnStack")
   external static void deoptimizeFunctionsOnStack();