[ VM ] Ensure heap profiler state is updated on TLAB release due to scavenge

Also fixes some race conditions and flaky test behavior.

Fixes https://github.com/dart-lang/sdk/issues/51098

TEST=DartAPI_HeapSampling_* looped in GDB under load

Change-Id: I0c1f297ecb140575e28650837158b707b9dc4b41
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/279590
Reviewed-by: Ryan Macnak <rmacnak@google.com>
Commit-Queue: Ben Konyi <bkonyi@google.com>

runtime/vm/dart_api_impl_test.cc

@@ -10524,6 +10524,7 @@ void InitHeapSampling(Thread* thread) {
 }
 
 TEST_CASE(DartAPI_HeapSampling_UserDefinedClass) {
+  DisableBackgroundCompilationScope scope;
   auto isolate_group_data = Dart_CurrentIsolateGroupData();
   const char* kScriptChars = R"(
     class Bar {}
@@ -10544,9 +10545,8 @@ TEST_CASE(DartAPI_HeapSampling_UserDefinedClass) {
   EXPECT_VALID(result);
   EXPECT(heap_samples > 0);
   EXPECT(heap_samples < 100000);
-  EXPECT_EQ(last_isolate_group_data, isolate_group_data);
-
-  EXPECT_STREQ("Bar", last_allocation_cls);
+  EXPECT(last_isolate_group_data == isolate_group_data);
+  EXPECT_STREQ(last_allocation_cls, "Bar");
 }
 
 TEST_CASE(DartAPI_HeapSampling_APIAllocations) {
@@ -10588,6 +10588,7 @@ TEST_CASE(DartAPI_HeapSampling_APIAllocations) {
 }
 
 TEST_CASE(DartAPI_HeapSampling_NonTrivialSamplingPeriod) {
+  DisableBackgroundCompilationScope scope;
   auto isolate_group_data = Dart_CurrentIsolateGroupData();
   InitHeapSampling(thread);
 
@@ -10630,7 +10631,7 @@ TEST_CASE(DartAPI_HeapSampling_NonTrivialSamplingPeriod) {
   EXPECT_VALID(result);
   EXPECT(heap_samples > 0);
   EXPECT(heap_samples < kNumAllocations);
-  EXPECT_EQ(last_isolate_group_data, isolate_group_data);
+  EXPECT(last_isolate_group_data == isolate_group_data);
 
   Dart_DisableHeapSampling();
 

runtime/vm/heap/page.h

@@ -257,6 +257,9 @@ class Page {
     thread->set_top(0);
     thread->set_end(0);
     thread->set_true_end(0);
+#if !defined(PRODUCT)
+    thread->heap_sampler().HandleReleasedTLAB(Thread::Current());
+#endif
   }
   void Release() {
     if (owner_ != nullptr) {

runtime/vm/heap/sampler.cc

@@ -92,9 +92,11 @@ void HeapProfileSampler::ResetState() {
 }
 
 void HeapProfileSampler::Initialize() {
-  // We're initializing state for this thread.
-  ReadRwLocker locker(Thread::Current(), lock_);
-  UpdateThreadEnable();
+  // Don't grab lock_ here as it can cause a deadlock if thread initialization
+  // occurs when the profiler state is being changed. Instead, let the thread
+  // perform initialization when it's no longer holding the thread registry's
+  // thread_lock().
+  ScheduleUpdateThreadEnable();
 }
 
 void HeapProfileSampler::ScheduleUpdateThreadEnable() {
@@ -141,6 +143,15 @@ void HeapProfileSampler::SetThreadSamplingIntervalLocked() {
   SetNextSamplingIntervalLocked(GetNextSamplingIntervalLocked());
 }
 
+void HeapProfileSampler::HandleReleasedTLAB(Thread* thread) {
+  ReadRwLocker locker(thread, lock_);
+  if (!enabled_) {
+    return;
+  }
+  interval_to_next_sample_ = remaining_TLAB_interval();
+  next_tlab_offset_ = kUninitialized;
+}
+
 void HeapProfileSampler::HandleNewTLAB(intptr_t old_tlab_remaining_space,
                                        bool is_first_tlab) {
   ASSERT_THREAD_STATE(thread_);
@@ -347,7 +358,6 @@ void HeapProfileSampler::SetNextSamplingIntervalLocked(intptr_t next_interval) {
     thread_->set_end(new_end);
     ASSERT_TLAB_BOUNDARIES_VALID(thread_);
   }
-  ASSERT(interval_to_next_sample_ == kUninitialized);
   interval_to_next_sample_ = next_interval;
 }
 

runtime/vm/heap/sampler.h

@@ -99,6 +99,11 @@ class HeapProfileSampler {
   // profiler instance to avoid concurrent modification of the thread's TLAB.
   void SetThreadSamplingInterval();
 
+  // Updates internal book keeping tracking the remaining size of the sampling
+  // interval. This method must be called when a TLAB is torn down to ensure
+  // that a future TLAB is initialized with the correct sampling interval.
+  void HandleReleasedTLAB(Thread* thread);
+
   // Handles the creation of a new TLAB by updating its boundaries based on the
   // remaining sampling interval.
   //

runtime/vm/profiler_test.cc

@@ -48,21 +48,6 @@ class DisableNativeProfileScope : public ValueObject {
   const bool FLAG_profile_vm_allocation_;
 };
 
-class DisableBackgroundCompilationScope : public ValueObject {
- public:
-  DisableBackgroundCompilationScope()
-      : FLAG_background_compilation_(FLAG_background_compilation) {
-    FLAG_background_compilation = false;
-  }
-
-  ~DisableBackgroundCompilationScope() {
-    FLAG_background_compilation = FLAG_background_compilation_;
-  }
-
- private:
-  const bool FLAG_background_compilation_;
-};
-
 // Temporarily adjust the maximum profile depth.
 class MaxProfileDepthScope : public ValueObject {
  public:

runtime/vm/unit_test.h

@@ -755,6 +755,21 @@ class SetFlagScope : public ValueObject {
   T original_value_;
 };
 
+class DisableBackgroundCompilationScope : public ValueObject {
+ public:
+  DisableBackgroundCompilationScope()
+      : FLAG_background_compilation_(FLAG_background_compilation) {
+    FLAG_background_compilation = false;
+  }
+
+  ~DisableBackgroundCompilationScope() {
+    FLAG_background_compilation = FLAG_background_compilation_;
+  }
+
+ private:
+  const bool FLAG_background_compilation_;
+};
+
 }  // namespace dart
 
 #endif  // RUNTIME_VM_UNIT_TEST_H_