8353184: ZGC: Simplify and correct tlab_used() tracking

Reviewed-by: stefank, aboldtch

Author: kstefanj

src/hotspot/share/gc/z/zAllocator.cpp

@@ -39,10 +39,6 @@ ZAllocatorEden::ZAllocatorEden()
   ZAllocator::_eden = this;
 }
 
-size_t ZAllocatorEden::tlab_used() const {
-  return _object_allocator.used();
-}
-
 size_t ZAllocatorEden::remaining() const {
   return _object_allocator.remaining();
 }

src/hotspot/share/gc/z/zAllocator.hpp

@@ -62,7 +62,6 @@ class ZAllocatorEden : public ZAllocator {
   zaddress alloc_object(size_t size);
 
   // Statistics
-  size_t tlab_used() const;
   size_t remaining() const;
 };
 

src/hotspot/share/gc/z/zArguments.cpp

@@ -143,6 +143,11 @@ void ZArguments::initialize() {
     FLAG_SET_ERGO_IF_DEFAULT(ZCollectionIntervalMajor, ZCollectionInterval);
   }
 
+  // Set an initial TLAB size to avoid depending on the current capacity
+  if (FLAG_IS_DEFAULT(TLABSize)) {
+    FLAG_SET_DEFAULT(TLABSize, 256*K);
+  }
+
   // Set medium page size here because MaxTenuringThreshold may use it.
   ZHeuristics::set_medium_page_size();
 

src/hotspot/share/gc/z/zCollectedHeap.cpp

@@ -238,7 +238,7 @@ size_t ZCollectedHeap::tlab_used(Thread* ignored) const {
 }
 
 size_t ZCollectedHeap::max_tlab_size() const {
-  return _heap.max_tlab_size();
+  return _heap.max_tlab_size() / HeapWordSize;
 }
 
 size_t ZCollectedHeap::unsafe_max_tlab_alloc(Thread* ignored) const {

src/hotspot/share/gc/z/zGeneration.cpp

@@ -837,6 +837,9 @@ void ZGenerationYoung::mark_start() {
   // Change good colors
   flip_mark_start();
 
+  // Reset TLAB usage
+  ZHeap::heap()->reset_tlab_used();
+
   // Retire allocating pages
   ZAllocator::eden()->retire_pages();
   for (ZPageAge i = ZPageAge::survivor1; i <= ZPageAge::survivor14; i = static_cast<ZPageAge>(static_cast<uint>(i) + 1)) {

src/hotspot/share/gc/z/zHeap.cpp

@@ -62,6 +62,7 @@ ZHeap::ZHeap()
     _serviceability(InitialHeapSize, min_capacity(), max_capacity()),
     _old(&_page_table, &_page_allocator),
     _young(&_page_table, _old.forwarding_table(), &_page_allocator),
+    _tlab_usage(),
     _initialized(false) {
 
   // Install global heap instance
@@ -131,11 +132,11 @@ size_t ZHeap::unused() const {
 }
 
 size_t ZHeap::tlab_capacity() const {
-  return capacity();
+  return _tlab_usage.tlab_capacity();
 }
 
 size_t ZHeap::tlab_used() const {
-  return _allocator_eden.tlab_used();
+  return _tlab_usage.tlab_used();
 }
 
 size_t ZHeap::max_tlab_size() const {
@@ -156,6 +157,9 @@ size_t ZHeap::unsafe_max_tlab_alloc() const {
 
   return MIN2(size, max_tlab_size());
 }
+void ZHeap::reset_tlab_used() {
+  _tlab_usage.reset();
+}
 
 bool ZHeap::is_in(uintptr_t addr) const {
   if (addr == 0) {
@@ -219,11 +223,38 @@ void ZHeap::out_of_memory() {
   log_info(gc)("Out Of Memory (%s)", Thread::current()->name());
 }
 
+static bool is_small_eden_page(ZPage* page) {
+  return page->type() == ZPageType::small && page->age() == ZPageAge::eden;
+}
+
+void ZHeap::account_alloc_page(ZPage* page) {
+  // Do TLAB accounting for small eden pages
+  if (is_small_eden_page(page)) {
+    _tlab_usage.increase_used(page->size());
+  }
+}
+
+void ZHeap::account_undo_alloc_page(ZPage* page) {
+  // Increase the undo counter
+  ZStatInc(ZCounterUndoPageAllocation);
+
+  // Undo TLAB accounting for small eden pages
+  if (is_small_eden_page(page)) {
+    _tlab_usage.decrease_used(page->size());
+  }
+
+  log_trace(gc)("Undo page allocation, thread: " PTR_FORMAT " (%s), page: " PTR_FORMAT ", size: %zu",
+                p2i(Thread::current()), ZUtils::thread_name(), p2i(page), page->size());
+}
+
 ZPage* ZHeap::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age) {
   ZPage* const page = _page_allocator.alloc_page(type, size, flags, age);
   if (page != nullptr) {
     // Insert page table entry
     _page_table.insert(page);
+
+    // Do accounting for the allocated page
+    account_alloc_page(page);
   }
 
   return page;
@@ -232,9 +263,8 @@ ZPage* ZHeap::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZP
 void ZHeap::undo_alloc_page(ZPage* page) {
   assert(page->is_allocating(), "Invalid page state");
 
-  ZStatInc(ZCounterUndoPageAllocation);
-  log_trace(gc)("Undo page allocation, thread: " PTR_FORMAT " (%s), page: " PTR_FORMAT ", size: %zu",
-                p2i(Thread::current()), ZUtils::thread_name(), p2i(page), page->size());
+  // Undo accounting for the page being freed
+  account_undo_alloc_page(page);
 
   free_page(page);
 }

src/hotspot/share/gc/z/zHeap.hpp

@@ -33,6 +33,7 @@
 #include "gc/z/zPageTable.hpp"
 #include "gc/z/zPageType.hpp"
 #include "gc/z/zServiceability.hpp"
+#include "gc/z/zTLABUsage.hpp"
 
 class OopFieldClosure;
 
@@ -55,8 +56,14 @@ class ZHeap {
   ZGenerationOld          _old;
   ZGenerationYoung        _young;
 
+  ZTLABUsage              _tlab_usage;
+
   bool                    _initialized;
 
+  // Page allocation accounting
+  void account_alloc_page(ZPage* page);
+  void account_undo_alloc_page(ZPage* page);
+
 public:
   static ZHeap* heap();
 
@@ -81,6 +88,7 @@ class ZHeap {
   size_t tlab_used() const;
   size_t max_tlab_size() const;
   size_t unsafe_max_tlab_alloc() const;
+  void reset_tlab_used();
 
   bool is_in(uintptr_t addr) const;
   bool is_in_page_relaxed(const ZPage* page, zaddress addr) const;

src/hotspot/share/gc/z/zObjectAllocator.cpp

@@ -43,8 +43,6 @@ static const ZStatCounter ZCounterUndoObjectAllocationFailed("Memory", "Undo Obj
 ZObjectAllocator::ZObjectAllocator(ZPageAge age)
   : _age(age),
     _use_per_cpu_shared_small_pages(ZHeuristics::use_per_cpu_shared_small_pages()),
-    _used(0),
-    _undone(0),
     _shared_small_page(nullptr),
     _shared_medium_page(nullptr),
     _medium_page_alloc_lock() {}
@@ -58,23 +56,14 @@ ZPage* const* ZObjectAllocator::shared_small_page_addr() const {
 }
 
 ZPage* ZObjectAllocator::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags) {
-  ZPage* const page = ZHeap::heap()->alloc_page(type, size, flags, _age);
-  if (page != nullptr) {
-    // Increment used bytes
-    Atomic::add(_used.addr(), size);
-  }
-
-  return page;
+  return ZHeap::heap()->alloc_page(type, size, flags, _age);
 }
 
 ZPage* ZObjectAllocator::alloc_page_for_relocation(ZPageType type, size_t size, ZAllocationFlags flags) {
   return ZHeap::heap()->alloc_page(type, size, flags, _age);
 }
 
 void ZObjectAllocator::undo_alloc_page(ZPage* page) {
-  // Increment undone bytes
-  Atomic::add(_undone.addr(), page->size());
-
   ZHeap::heap()->undo_alloc_page(page);
 }
 
@@ -229,23 +218,6 @@ ZPageAge ZObjectAllocator::age() const {
   return _age;
 }
 
-size_t ZObjectAllocator::used() const {
-  size_t total_used = 0;
-  size_t total_undone = 0;
-
-  ZPerCPUConstIterator<size_t> iter_used(&_used);
-  for (const size_t* cpu_used; iter_used.next(&cpu_used);) {
-    total_used += *cpu_used;
-  }
-
-  ZPerCPUConstIterator<size_t> iter_undone(&_undone);
-  for (const size_t* cpu_undone; iter_undone.next(&cpu_undone);) {
-    total_undone += *cpu_undone;
-  }
-
-  return total_used - total_undone;
-}
-
 size_t ZObjectAllocator::remaining() const {
   assert(Thread::current()->is_Java_thread(), "Should be a Java thread");
 
@@ -260,10 +232,6 @@ size_t ZObjectAllocator::remaining() const {
 void ZObjectAllocator::retire_pages() {
   assert(SafepointSynchronize::is_at_safepoint(), "Should be at safepoint");
 
-  // Reset used and undone bytes
-  _used.set_all(0);
-  _undone.set_all(0);
-
   // Reset allocation pages
   _shared_medium_page.set(nullptr);
   _shared_small_page.set_all(nullptr);

src/hotspot/share/gc/z/zObjectAllocator.hpp

@@ -38,8 +38,6 @@ class ZObjectAllocator {
 private:
   ZPageAge           _age;
   const bool         _use_per_cpu_shared_small_pages;
-  ZPerCPU<size_t>    _used;
-  ZPerCPU<size_t>    _undone;
   ZPerCPU<ZPage*>    _shared_small_page;
   ZContended<ZPage*> _shared_medium_page;
   ZLock              _medium_page_alloc_lock;
@@ -80,7 +78,6 @@ class ZObjectAllocator {
 
   ZPageAge age() const;
 
-  size_t used() const;
   size_t remaining() const;
 
   void retire_pages();

src/hotspot/share/gc/z/zTLABUsage.cpp

@@ -0,0 +1,70 @@
+/*
+ * 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.
+ */
+
+#include "gc/z/zTLABUsage.hpp"
+#include "logging/log.hpp"
+#include "runtime/atomic.hpp"
+
+ZTLABUsage::ZTLABUsage()
+  : _used(0),
+    _used_history() {}
+
+void ZTLABUsage::increase_used(size_t size) {
+  Atomic::add(&_used, size, memory_order_relaxed);
+}
+
+void ZTLABUsage::decrease_used(size_t size) {
+  precond(size <= _used);
+
+  Atomic::sub(&_used, size, memory_order_relaxed);
+}
+
+void ZTLABUsage::reset() {
+  const size_t used = Atomic::xchg(&_used, (size_t) 0);
+
+  // Avoid updates when nothing has been allocated since the last YC
+  if (used == 0) {
+    return;
+  }
+
+  // Save the old values for logging
+  const size_t old_tlab_used = tlab_used();
+  const size_t old_tlab_capacity = tlab_capacity();
+
+  // Update the usage history with the current value
+  _used_history.add(used);
+
+  log_debug(gc, tlab)("TLAB usage update: used %zuM -> %zuM, capacity: %zuM -> %zuM",
+                      old_tlab_used / M,
+                      tlab_used() / M,
+                      old_tlab_capacity / M,
+                      tlab_capacity() / M);
+  }
+
+size_t ZTLABUsage::tlab_used() const {
+  return _used_history.last();
+}
+
+size_t ZTLABUsage::tlab_capacity() const {
+  return _used_history.davg();
+}