8362573: Incorrect weight of the first ObjectAllocationSample JFR event (regression)

src/hotspot/share/jfr/recorder/checkpoint/jfrCheckpointManager.cpp

@@ -566,7 +566,7 @@ size_t JfrCheckpointManager::write_static_type_set_and_threads() {
 void JfrCheckpointManager::on_rotation() {
   assert(SafepointSynchronize::is_at_safepoint(), "invariant");
   JfrTypeManager::on_rotation();
-  notify_threads();
+  notify_threads(false);
 }
 
 void JfrCheckpointManager::clear_type_set() {
@@ -676,18 +676,36 @@ void JfrCheckpointManager::write_simplified_vthread_checkpoint(traceid vtid) {
   JfrTypeManager::write_simplified_vthread_checkpoint(vtid);
 }
 
-class JfrNotifyClosure : public ThreadClosure {
+// Reset thread local state used for object allocation sampling.
+static void clear_last_allocated_bytes(JavaThread* jt) {
+  assert(jt != nullptr, "invariant");
+  assert(!JfrRecorder::is_recording(), "invariant");
+  JfrThreadLocal* const tl = jt->jfr_thread_local();
+  assert(tl != nullptr, "invariant");
+  if (tl->last_allocated_bytes() != 0) {
+    tl->clear_last_allocated_bytes();
+  }
+  assert(tl->last_allocated_bytes() == 0, "invariant");
+}
+
+class JfrNotifyClosure : public StackObj {
+ private:
+  bool _clear;
  public:
-  void do_thread(Thread* thread) {
-    assert(thread != nullptr, "invariant");
+  JfrNotifyClosure(bool clear) : _clear(clear) {}
+  void do_thread(JavaThread* jt) {
+    assert(jt != nullptr, "invariant");
     assert_locked_or_safepoint(Threads_lock);
-    JfrJavaEventWriter::notify(JavaThread::cast(thread));
+    JfrJavaEventWriter::notify(jt);
+    if (_clear) {
+      clear_last_allocated_bytes(jt);
+    }
   }
 };
 
-void JfrCheckpointManager::notify_threads() {
+void JfrCheckpointManager::notify_threads(bool clear /* true */) {
   assert(SafepointSynchronize::is_at_safepoint(), "invariant");
-  JfrNotifyClosure tc;
+  JfrNotifyClosure tc(clear);
   JfrJavaThreadIterator iter;
   while (iter.has_next()) {
     tc.do_thread(iter.next());

src/hotspot/share/jfr/recorder/checkpoint/jfrCheckpointManager.hpp

@@ -88,7 +88,7 @@ class JfrCheckpointManager : public JfrCHeapObj {
 
   size_t clear();
   size_t write();
-  void notify_threads();
+  void notify_threads(bool clear = true);
 
   size_t write_static_type_set(Thread* thread);
   size_t write_threads(JavaThread* thread);

src/hotspot/share/jfr/recorder/service/jfrRecorderService.cpp

@@ -122,30 +122,6 @@ const Thread* JfrRotationLock::_owner_thread = nullptr;
 const int JfrRotationLock::retry_wait_millis = 10;
 volatile int JfrRotationLock::_lock = 0;
 
-// Reset thread local state used for object allocation sampling.
-class ClearObjectAllocationSampling : public ThreadClosure {
- public:
-  void do_thread(Thread* t) {
-    assert(t != nullptr, "invariant");
-    t->jfr_thread_local()->clear_last_allocated_bytes();
-  }
-};
-
-template <typename Iterator>
-static inline void iterate(Iterator& it, ClearObjectAllocationSampling& coas) {
-  while (it.has_next()) {
-    coas.do_thread(it.next());
-  }
-}
-
-static void clear_object_allocation_sampling() {
-  ClearObjectAllocationSampling coas;
-  JfrJavaThreadIterator jit;
-  iterate(jit, coas);
-  JfrNonJavaThreadIterator njit;
-  iterate(njit, coas);
-}
-
 template <typename Instance, size_t(Instance::*func)()>
 class Content {
  private:
@@ -472,7 +448,6 @@ void JfrRecorderService::clear() {
 }
 
 void JfrRecorderService::pre_safepoint_clear() {
-  clear_object_allocation_sampling();
   _storage.clear();
   JfrStackTraceRepository::clear();
 }

src/hotspot/share/jfr/support/jfrAllocationTracer.cpp

@@ -22,6 +22,7 @@
 *
 */
 
+#include "jfr/jfrEvents.hpp"
 #include "jfr/leakprofiler/leakProfiler.hpp"
 #include "jfr/support/jfrAllocationTracer.hpp"
 #include "jfr/support/jfrObjectAllocationSample.hpp"
@@ -31,5 +32,7 @@ JfrAllocationTracer::JfrAllocationTracer(const Klass* klass, HeapWord* obj, size
   if (LeakProfiler::is_running()) {
     LeakProfiler::sample(obj, alloc_size, thread);
   }
-  JfrObjectAllocationSample::send_event(klass, alloc_size, outside_tlab, thread);
+  if (EventObjectAllocationSample::is_enabled()) {
+    JfrObjectAllocationSample::send_event(klass, alloc_size, outside_tlab, thread);
+  }
 }

src/hotspot/share/jfr/support/jfrObjectAllocationSample.cpp

@@ -26,7 +26,7 @@
 #include "gc/shared/tlab_globals.hpp"
 #include "jfr/jfrEvents.hpp"
 #include "jfr/support/jfrObjectAllocationSample.hpp"
-#include "jfr/support/jfrThreadLocal.hpp"
+#include "runtime/javaThread.hpp"
 #include "utilities/globalDefinitions.hpp"
 
 inline bool send_allocation_sample(const Klass* klass, int64_t allocated_bytes, JfrThreadLocal* tl) {
@@ -43,35 +43,42 @@ inline bool send_allocation_sample(const Klass* klass, int64_t allocated_bytes,
   return false;
 }
 
-inline int64_t estimate_tlab_size_bytes(Thread* thread) {
-  const size_t desired_tlab_size_bytes = thread->tlab().desired_size() * HeapWordSize;
-  const size_t alignment_reserve_bytes = thread->tlab().alignment_reserve_in_bytes();
+inline int64_t estimate_tlab_size_bytes(JavaThread* jt) {
+  const size_t desired_tlab_size_bytes = jt->tlab().desired_size() * HeapWordSize;
+  const size_t alignment_reserve_bytes = jt->tlab().alignment_reserve_in_bytes();
   assert(desired_tlab_size_bytes >= alignment_reserve_bytes, "invariant");
   return static_cast<int64_t>(desired_tlab_size_bytes - alignment_reserve_bytes);
 }
 
-inline int64_t load_allocated_bytes(JfrThreadLocal* tl, Thread* thread) {
-  const int64_t allocated_bytes = thread->allocated_bytes();
-  return allocated_bytes == tl->last_allocated_bytes() ? 0 : allocated_bytes;
+inline int64_t load_allocated_bytes(JfrThreadLocal* tl, JavaThread* jt) {
+  const int64_t allocated_bytes = jt->allocated_bytes();
+  const int64_t last_allocated_bytes = tl->last_allocated_bytes();
+  assert(allocated_bytes >= last_allocated_bytes, "invariant");
+  if (last_allocated_bytes == 0) {
+    // Initialization.
+    tl->set_last_allocated_bytes(allocated_bytes);
+    return 0;
+  }
+  return allocated_bytes == last_allocated_bytes ? 0 : allocated_bytes;
 }
 
 // To avoid large objects from being undersampled compared to the regular TLAB samples,
 // the data amount is normalized as if it was a TLAB, giving a number of TLAB sampling attempts to the large object.
-static void normalize_as_tlab_and_send_allocation_samples(const Klass* klass, int64_t obj_alloc_size_bytes, JfrThreadLocal* tl, Thread* thread) {
-  const int64_t allocated_bytes = load_allocated_bytes(tl, thread);
+static void normalize_as_tlab_and_send_allocation_samples(const Klass* klass,
+                                                          int64_t obj_alloc_size_bytes,
+                                                          int64_t allocated_bytes,
+                                                          JfrThreadLocal* tl,
+                                                          JavaThread* jt) {
   assert(allocated_bytes > 0, "invariant"); // obj_alloc_size_bytes is already attributed to allocated_bytes at this point.
   if (!UseTLAB) {
     send_allocation_sample(klass, allocated_bytes, tl);
     return;
   }
-  const int64_t tlab_size_bytes = estimate_tlab_size_bytes(thread);
-  if (tlab_size_bytes <= 0) {
+  const int64_t tlab_size_bytes = estimate_tlab_size_bytes(jt);
+  if (tlab_size_bytes <= 0 || allocated_bytes - tl->last_allocated_bytes() < tlab_size_bytes) {
     // We don't get a TLAB, avoid endless loop below.
     return;
   }
-  if (allocated_bytes - tl->last_allocated_bytes() < tlab_size_bytes) {
-    return;
-  }
   assert(obj_alloc_size_bytes > 0, "invariant");
   do {
     if (send_allocation_sample(klass, allocated_bytes, tl)) {
@@ -81,17 +88,17 @@ static void normalize_as_tlab_and_send_allocation_samples(const Klass* klass, in
   } while (obj_alloc_size_bytes > 0);
 }
 
-void JfrObjectAllocationSample::send_event(const Klass* klass, size_t alloc_size, bool outside_tlab, Thread* thread) {
-  assert(thread != nullptr, "invariant");
-  JfrThreadLocal* const tl = thread->jfr_thread_local();
+void JfrObjectAllocationSample::send_event(const Klass* klass, size_t alloc_size, bool outside_tlab, JavaThread* jt) {
+  assert(klass != nullptr, "invariant");
+  assert(jt != nullptr, "invariant");
+  JfrThreadLocal* const tl = jt->jfr_thread_local();
   assert(tl != nullptr, "invariant");
-  if (outside_tlab) {
-    normalize_as_tlab_and_send_allocation_samples(klass, static_cast<int64_t>(alloc_size), tl, thread);
-    return;
-  }
-  const int64_t allocated_bytes = load_allocated_bytes(tl, thread);
-  if (allocated_bytes == 0) {
-    return;
+  const int64_t allocated_bytes = load_allocated_bytes(tl, jt);
+  if (allocated_bytes > 0) {
+    if (outside_tlab) {
+      normalize_as_tlab_and_send_allocation_samples(klass, static_cast<int64_t>(alloc_size), allocated_bytes, tl, jt);
+      return;
+    }
+    send_allocation_sample(klass, allocated_bytes, tl);
   }
-  send_allocation_sample(klass, allocated_bytes, tl);
 }

src/hotspot/share/jfr/support/jfrObjectAllocationSample.hpp

@@ -27,12 +27,12 @@
 
 #include "memory/allStatic.hpp"
 
+class JavaThread;
 class Klass;
-class Thread;
 
 class JfrObjectAllocationSample : AllStatic {
   friend class JfrAllocationTracer;
-  static void send_event(const Klass* klass, size_t alloc_size, bool outside_tlab, Thread* thread);
+  static void send_event(const Klass* klass, size_t alloc_size, bool outside_tlab, JavaThread* jt);
 };
 
 #endif // SHARE_JFR_SUPPORT_JFROBJECTALLOCATIONSAMPLE_HPP

test/jdk/jdk/jfr/event/allocation/TestObjectAllocationSampleEventInitialWeight.java

@@ -0,0 +1,100 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+package jdk.jfr.event.allocation;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.CountDownLatch;
+import jdk.jfr.consumer.RecordedEvent;
+import jdk.jfr.consumer.RecordingStream;
+import jdk.test.lib.jfr.EventNames;
+import jdk.test.lib.jfr.Events;
+
+/**
+ * @test
+ * @summary Tests that the VM maintains proper initialization state for ObjectAllocationSampleEvent.
+ * @requires vm.hasJFR
+ * @library /test/lib
+ * @run main/othervm -XX:+UseTLAB -XX:TLABSize=2k -XX:-ResizeTLAB jdk.jfr.event.allocation.TestObjectAllocationSampleEventInitialWeight
+ */
+public class TestObjectAllocationSampleEventInitialWeight {
+    private static final String EVENT_NAME = EventNames.ObjectAllocationSample;
+    private static final int OBJECT_SIZE = 4 * 1024;
+    private static final int OBJECTS_TO_ALLOCATE = 16;
+    private static final int OBJECTS_TO_ALLOCATE_BEFORE_RECORDING = 1024;
+    private static final long BEFORE_RECORDING_SAMPLE_WEIGHT = OBJECT_SIZE * OBJECTS_TO_ALLOCATE_BEFORE_RECORDING;
+    private static final String BYTE_ARRAY_CLASS_NAME = new byte[0].getClass().getName();
+
+    private static AtomicBoolean onError = null;
+
+    // Make sure allocation isn't dead code eliminated.
+    public static byte[] tmp;
+
+    public static void main(String... args) throws Exception {
+        test();
+        // Test again to ensure reset logic works correctly for subsequent physical recordings.
+        test();
+    }
+
+    private static void test() throws Exception {
+        CountDownLatch delivered = new CountDownLatch(1);
+        onError = new AtomicBoolean();
+        Thread current = Thread.currentThread();
+        allocateBeforeRecording();
+        try (RecordingStream rs = new RecordingStream()) {
+            rs.enable(EVENT_NAME);
+            rs.onEvent(EVENT_NAME, e -> {
+                if (verify(e, current)) {
+                    delivered.countDown();
+                }
+            });
+            rs.startAsync();
+            allocate(OBJECTS_TO_ALLOCATE);
+            delivered.await();
+            if (onError.get()) {
+                throw new RuntimeException("Sample weight is not below " + BEFORE_RECORDING_SAMPLE_WEIGHT);
+            }
+        }
+    }
+
+    private static void allocateBeforeRecording() throws Exception {
+        allocate(OBJECTS_TO_ALLOCATE_BEFORE_RECORDING);
+    }
+
+    private static void allocate(int number) throws Exception {
+        for (int i = 0; i < number; ++i) {
+            tmp = new byte[OBJECT_SIZE];
+        }
+    }
+
+    private static boolean verify(RecordedEvent event, Thread thread) {
+        if (thread.getId() != event.getThread().getJavaThreadId()) {
+            return false;
+        }
+        System.out.println(event);
+        if (event.getLong("weight") >= BEFORE_RECORDING_SAMPLE_WEIGHT) {
+            onError.set(true);
+        }
+        return true;
+    }
+}