Record sizeof(key_type), sizeof(value_type) in hashtable profiles.

This can identify situations where flat_hash_* is suboptimal for large
elements.

PiperOrigin-RevId: 618937993
Change-Id: I2bde069bc3526b14ad1718ba6f50467002aeed16

absl/container/internal/hashtablez_sampler.cc

@@ -79,6 +79,8 @@ HashtablezInfo::~HashtablezInfo() = default;
 
 void HashtablezInfo::PrepareForSampling(int64_t stride,
                                         size_t inline_element_size_value,
+                                        size_t key_size_value,
+                                        size_t value_size_value,
                                         uint16_t soo_capacity_value) {
   capacity.store(0, std::memory_order_relaxed);
   size.store(0, std::memory_order_relaxed);
@@ -99,6 +101,8 @@ void HashtablezInfo::PrepareForSampling(int64_t stride,
   depth = absl::GetStackTrace(stack, HashtablezInfo::kMaxStackDepth,
                               /* skip_count= */ 0);
   inline_element_size = inline_element_size_value;
+  key_size = key_size_value;
+  value_size = value_size_value;
   soo_capacity = soo_capacity_value;
 }
 
@@ -123,12 +127,13 @@ static bool ShouldForceSampling() {
 }
 
 HashtablezInfo* SampleSlow(SamplingState& next_sample,
-                           size_t inline_element_size, uint16_t soo_capacity) {
+                           size_t inline_element_size, size_t key_size,
+                           size_t value_size, uint16_t soo_capacity) {
   if (ABSL_PREDICT_FALSE(ShouldForceSampling())) {
     next_sample.next_sample = 1;
     const int64_t old_stride = exchange(next_sample.sample_stride, 1);
     HashtablezInfo* result = GlobalHashtablezSampler().Register(
-        old_stride, inline_element_size, soo_capacity);
+        old_stride, inline_element_size, key_size, value_size, soo_capacity);
     return result;
   }
 
@@ -158,11 +163,12 @@ HashtablezInfo* SampleSlow(SamplingState& next_sample,
   // that case.
   if (first) {
     if (ABSL_PREDICT_TRUE(--next_sample.next_sample > 0)) return nullptr;
-    return SampleSlow(next_sample, inline_element_size, soo_capacity);
+    return SampleSlow(next_sample, inline_element_size, key_size, value_size,
+                      soo_capacity);
   }
 
   return GlobalHashtablezSampler().Register(old_stride, inline_element_size,
-                                            soo_capacity);
+                                            key_size, value_size, soo_capacity);
 #endif
 }
 

absl/container/internal/hashtablez_sampler.h

@@ -73,6 +73,7 @@ struct HashtablezInfo : public profiling_internal::Sample<HashtablezInfo> {
   // Puts the object into a clean state, fills in the logically `const` members,
   // blocking for any readers that are currently sampling the object.
   void PrepareForSampling(int64_t stride, size_t inline_element_size_value,
+                          size_t key_size, size_t value_size,
                           uint16_t soo_capacity_value)
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(init_mu);
 
@@ -104,6 +105,8 @@ struct HashtablezInfo : public profiling_internal::Sample<HashtablezInfo> {
   uint16_t soo_capacity;
   void* stack[kMaxStackDepth];
   size_t inline_element_size;  // How big is the slot in bytes?
+  size_t key_size;             // sizeof(key_type)
+  size_t value_size;           // sizeof(value_type)
 };
 
 void RecordRehashSlow(HashtablezInfo* info, size_t total_probe_length);
@@ -128,7 +131,8 @@ struct SamplingState {
 };
 
 HashtablezInfo* SampleSlow(SamplingState& next_sample,
-                           size_t inline_element_size, uint16_t soo_capacity);
+                           size_t inline_element_size, size_t key_size,
+                           size_t value_size, uint16_t soo_capacity);
 void UnsampleSlow(HashtablezInfo* info);
 
 #if defined(ABSL_INTERNAL_HASHTABLEZ_SAMPLE)
@@ -218,13 +222,16 @@ extern ABSL_PER_THREAD_TLS_KEYWORD SamplingState global_next_sample;
 // Returns a sampling handle.
 inline HashtablezInfoHandle Sample(
     ABSL_ATTRIBUTE_UNUSED size_t inline_element_size,
+    ABSL_ATTRIBUTE_UNUSED size_t key_size,
+    ABSL_ATTRIBUTE_UNUSED size_t value_size,
     ABSL_ATTRIBUTE_UNUSED uint16_t soo_capacity) {
 #if defined(ABSL_INTERNAL_HASHTABLEZ_SAMPLE)
   if (ABSL_PREDICT_TRUE(--global_next_sample.next_sample > 0)) {
     return HashtablezInfoHandle(nullptr);
   }
-  return HashtablezInfoHandle(
-      SampleSlow(global_next_sample, inline_element_size, soo_capacity));
+  return HashtablezInfoHandle(SampleSlow(global_next_sample,
+                                         inline_element_size, key_size,
+                                         value_size, soo_capacity));
 #else
   return HashtablezInfoHandle(nullptr);
 #endif  // !ABSL_PER_THREAD_TLS

absl/container/internal/raw_hash_set.h

@@ -1786,7 +1786,8 @@ constexpr bool ShouldSampleHashtablezInfo() {
 }
 
 template <bool kSooEnabled>
-HashtablezInfoHandle SampleHashtablezInfo(size_t sizeof_slot,
+HashtablezInfoHandle SampleHashtablezInfo(size_t sizeof_slot, size_t sizeof_key,
+                                          size_t sizeof_value,
                                           size_t old_capacity, bool was_soo,
                                           HashtablezInfoHandle forced_infoz,
                                           CommonFields& c) {
@@ -1794,12 +1795,12 @@ HashtablezInfoHandle SampleHashtablezInfo(size_t sizeof_slot,
   // In SOO, we sample on the first insertion so if this is an empty SOO case
   // (e.g. when reserve is called), then we still need to sample.
   if (kSooEnabled && was_soo && c.size() == 0) {
-    return Sample(sizeof_slot, SooCapacity());
+    return Sample(sizeof_slot, sizeof_key, sizeof_value, SooCapacity());
   }
   // For non-SOO cases, we sample whenever the capacity is increasing from zero
   // to non-zero.
   if (!kSooEnabled && old_capacity == 0) {
-    return Sample(sizeof_slot, 0);
+    return Sample(sizeof_slot, sizeof_key, sizeof_value, 0);
   }
   return c.infoz();
 }
@@ -1902,12 +1903,15 @@ class HashSetResizeHelper {
   template <typename Alloc, size_t SizeOfSlot, bool TransferUsesMemcpy,
             bool SooEnabled, size_t AlignOfSlot>
   ABSL_ATTRIBUTE_NOINLINE bool InitializeSlots(CommonFields& c, Alloc alloc,
-                                               ctrl_t soo_slot_h2) {
+                                               ctrl_t soo_slot_h2,
+                                               size_t key_size,
+                                               size_t value_size) {
     assert(c.capacity());
     HashtablezInfoHandle infoz =
         ShouldSampleHashtablezInfo<Alloc>()
-            ? SampleHashtablezInfo<SooEnabled>(SizeOfSlot, old_capacity_,
-                                               was_soo_, forced_infoz_, c)
+            ? SampleHashtablezInfo<SooEnabled>(SizeOfSlot, key_size, value_size,
+                                               old_capacity_, was_soo_,
+                                               forced_infoz_, c)
             : HashtablezInfoHandle{};
 
     const bool has_infoz = infoz.IsSampled();
@@ -3362,7 +3366,8 @@ class raw_hash_set {
   inline HashtablezInfoHandle try_sample_soo() {
     assert(is_soo());
     if (!ShouldSampleHashtablezInfo<CharAlloc>()) return HashtablezInfoHandle{};
-    return Sample(sizeof(slot_type), SooCapacity());
+    return Sample(sizeof(slot_type), sizeof(key_type), sizeof(value_type),
+                  SooCapacity());
   }
 
   inline void destroy_slots() {
@@ -3458,7 +3463,8 @@ class raw_hash_set {
         resize_helper.InitializeSlots<CharAlloc, sizeof(slot_type),
                                       PolicyTraits::transfer_uses_memcpy(),
                                       SooEnabled(), alignof(slot_type)>(
-            common(), CharAlloc(alloc_ref()), soo_slot_h2);
+            common(), CharAlloc(alloc_ref()), soo_slot_h2, sizeof(key_type),
+            sizeof(value_type));
 
     // In the SooEnabled() case, capacity is never 0 so we don't check.
     if (!SooEnabled() && resize_helper.old_capacity() == 0) {

absl/container/internal/raw_hash_set_test.cc

@@ -2571,6 +2571,8 @@ TYPED_TEST_P(RawHashSamplerTest, Sample) {
         std::memory_order_relaxed)]++;
     reservations[info.max_reserve.load(std::memory_order_relaxed)]++;
     EXPECT_EQ(info.inline_element_size, sizeof(typename TypeParam::value_type));
+    EXPECT_EQ(info.key_size, sizeof(typename TypeParam::key_type));
+    EXPECT_EQ(info.value_size, sizeof(typename TypeParam::value_type));
 
     if (soo_enabled) {
       EXPECT_EQ(info.soo_capacity, SooCapacity());