Add MakeProfileOptions::collect_compressibility and internal compressibility scanning support in profile_builder for heap profiles.

tcmalloc/internal/BUILD

@@ -957,13 +957,13 @@ cc_library(
     srcs = ["profile_builder.cc"],
     hdrs = ["profile_builder.h"],
     copts = TCMALLOC_DEFAULT_COPTS,
-    visibility = [
-        "//tcmalloc:__subpackages__",
-    ],
+    visibility = ["//tcmalloc:__subpackages__"],
     deps = [
         ":logging",
         ":pageflags",
+        ":parameter_accessors",
         ":residency",
+        ":util",
         "//tcmalloc:malloc_extension",
         "//tcmalloc/internal:profile_cc_proto",
         "@com_google_absl//absl/base",

tcmalloc/internal/parameter_accessors.h

@@ -31,6 +31,11 @@ struct TracerSizeClassInfo {
   size_t span_size_in_bytes;
   size_t num_objects_to_move;
 };
+
+struct CompressibleBuffer {
+  const void* ptr;
+  size_t size;
+};
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc
 
@@ -120,6 +125,10 @@ ABSL_ATTRIBUTE_WEAK void TCMalloc_Internal_GetSizeClasses(
     std::vector<tcmalloc::tcmalloc_internal::TracerSizeClassInfo>* absl_nonnull
         size_classes);
 ABSL_ATTRIBUTE_WEAK size_t TCMalloc_Internal_GetPageSize();
+
+ABSL_ATTRIBUTE_WEAK size_t TCMalloc_Internal_ZstdCompress(
+    const tcmalloc::tcmalloc_internal::CompressibleBuffer* buffers,
+    size_t num_buffers, void* dst, size_t dst_capacity);
 }
 
 #endif  // TCMALLOC_INTERNAL_PARAMETER_ACCESSORS_H_

tcmalloc/internal/profile_builder.cc

@@ -44,6 +44,7 @@
 #include <string>
 #include <tuple>
 #include <utility>
+#include <vector>
 
 #include "tcmalloc/internal/profile.pb.h"
 #include "absl/base/attributes.h"
@@ -52,7 +53,9 @@
 #include "absl/strings/escaping.h"
 #include "absl/strings/string_view.h"
 #include "tcmalloc/internal/logging.h"
+#include "tcmalloc/internal/parameter_accessors.h"
 #include "tcmalloc/internal/residency.h"
+#include "tcmalloc/internal/util.h"
 
 namespace tcmalloc {
 namespace tcmalloc_internal {
@@ -111,6 +114,9 @@ struct SampleMergedData {
   std::optional<size_t> stale_size;
   std::optional<size_t> locked_size;
   std::optional<uint64_t> stale_scan_period;
+  size_t standalone_compressed_size = 0;
+  size_t grouped_compressed_size = 0;
+  size_t zero_size = 0;
 };
 
 // The equality and hash methods of Profile::Sample only use a subset of its
@@ -152,7 +158,7 @@ SampleMergedMap MergeProfileSamplesAndMaybeGetResidencyInfo(
     data.count += entry.count;
     data.sum += entry.sum;
     if (residency) {
-      auto residency_info =
+      std::optional<Residency::Info> residency_info =
           residency->Get(entry.span_start_address, entry.allocated_size);
       // As long as `residency_info` provides data in some samples, the merged
       // data will have their sums.
@@ -199,6 +205,185 @@ SampleMergedMap MergeProfileSamplesAndMaybeGetResidencyInfo(
   return map;
 }
 
+// Calculates zero-byte count, standalone compressed size, and grouped
+// compressed size for all sampled objects in `profile`, updating `merged_map`
+// with the results. Uses SafeCopyMemory (process_vm_readv) to inspect
+// allocation memory without crashing if memory was freed or unmapped
+// concurrently. Keeps peak temporary memory strictly bounded to <= 256 KB.
+void AddCompressibilityInfo(const tcmalloc::Profile& profile,
+                            SampleMergedMap& merged_map) {
+  if (TCMalloc_Internal_ZstdCompress == nullptr) {
+    return;
+  }
+
+  // ---------------------------------------------------------------------------
+  // Phase 1: Standalone Compression & Zero Bytes (Peak RAM: ~128.5 KB)
+  // ---------------------------------------------------------------------------
+  struct SampleRef {
+    const tcmalloc::Profile::Sample* sample;
+    SampleMergedData* merged_data;
+    size_t scan_size;
+    size_t standalone_compressed_size_one;
+  };
+
+  std::vector<SampleRef> samples;
+  samples.reserve(1024);
+
+  std::string shared_in(/*count=*/64 * 1024, /*ch=*/'\0');
+  std::string shared_out(/*count=*/64 * 1024 + 512, /*ch=*/'\0');
+
+  profile.Iterate([&](const tcmalloc::Profile::Sample& entry) {
+    auto it = merged_map.find(entry);
+    if (it == merged_map.end()) {
+      return;
+    }
+    SampleMergedData& data = it->second;
+
+    if (entry.span_start_address == nullptr || entry.requested_size == 0) {
+      return;
+    }
+    size_t scan_size = std::min(entry.requested_size, size_t{64 * 1024});
+
+    if (!SafeCopyMemory(/*src=*/entry.span_start_address,
+                        /*dst=*/shared_in.data(), /*size=*/scan_size)) {
+      return;
+    }
+
+    size_t zero_bytes = 0;
+    const uint8_t* p = reinterpret_cast<const uint8_t*>(shared_in.data());
+    for (size_t i = 0; i < scan_size; ++i) {
+      if (p[i] == 0) {
+        zero_bytes++;
+      }
+    }
+    if (entry.requested_size > scan_size) {
+      zero_bytes = static_cast<size_t>(static_cast<double>(zero_bytes) *
+                                       entry.requested_size / scan_size);
+    }
+    data.zero_size += entry.count * zero_bytes;
+
+    size_t comp_size = 0;
+    CompressibleBuffer buf = {shared_in.data(), scan_size};
+    size_t comp = TCMalloc_Internal_ZstdCompress(
+        /*buffers=*/&buf, /*num_buffers=*/1, /*dst=*/shared_out.data(),
+        /*dst_capacity=*/shared_out.size());
+    if (comp > 0) {
+      comp_size = comp;
+      if (entry.requested_size > scan_size) {
+        comp_size = static_cast<size_t>(static_cast<double>(comp_size) *
+                                        entry.requested_size / scan_size);
+      }
+      data.standalone_compressed_size += entry.count * comp_size;
+    }
+
+    samples.push_back({&entry, &data, scan_size, comp_size});
+  });
+
+  // Reclaim Phase 1 scratch buffer before Phase 2 to bound total RAM <= 256 KB.
+  shared_in.clear();
+  shared_in.shrink_to_fit();
+
+  // ---------------------------------------------------------------------------
+  // Phase 2: Grouped Compression by Callstack (Peak RAM: ~192.5 KB <= 256 KB)
+  // ---------------------------------------------------------------------------
+  struct CallstackHash {
+    size_t operator()(const tcmalloc::Profile::Sample* s) const {
+      return absl::HashOf(absl::MakeConstSpan(s->stack, s->depth), s->depth);
+    }
+  };
+  struct CallstackEq {
+    bool operator()(const tcmalloc::Profile::Sample* a,
+                    const tcmalloc::Profile::Sample* b) const {
+      if (a->depth != b->depth) {
+        return false;
+      }
+      return std::equal(a->stack, a->stack + a->depth, b->stack);
+    }
+  };
+
+  absl::flat_hash_map<const tcmalloc::Profile::Sample*, std::vector<size_t>,
+                      CallstackHash, CallstackEq>
+      groups;
+  for (size_t i = 0; i < samples.size(); ++i) {
+    groups[samples[i].sample].push_back(i);
+  }
+
+  std::string grouped_in_pool(/*count=*/128 * 1024, /*ch=*/'\0');
+
+  for (const auto& group_pair : groups) {
+    std::vector<size_t> indices = group_pair.second;
+    std::sort(indices.begin(), indices.end(), [&](size_t i1, size_t i2) {
+      return samples[i1].sample->requested_size <
+             samples[i2].sample->requested_size;
+    });
+
+    if (indices.empty()) {
+      continue;
+    }
+
+    size_t idx_pos = 0;
+    while (idx_pos < indices.size()) {
+      std::vector<CompressibleBuffer> buffers;
+      std::vector<size_t> chunk_indices;
+      size_t pool_used = 0;
+      size_t chunk_uncompressed_total = 0;
+
+      while (idx_pos < indices.size()) {
+        size_t idx = indices[idx_pos];
+        SampleRef& s = samples[idx];
+        if (s.scan_size == 0 || s.sample->span_start_address == nullptr) {
+          ++idx_pos;
+          continue;
+        }
+        if (pool_used + s.scan_size > grouped_in_pool.size() &&
+            !chunk_indices.empty()) {
+          break;
+        }
+        size_t copy_len =
+            std::min(s.scan_size, grouped_in_pool.size() - pool_used);
+        if (SafeCopyMemory(/*src=*/s.sample->span_start_address,
+                           /*dst=*/grouped_in_pool.data() + pool_used,
+                           /*size=*/copy_len)) {
+          buffers.push_back({grouped_in_pool.data() + pool_used, copy_len});
+          chunk_indices.push_back(idx);
+          pool_used += copy_len;
+          chunk_uncompressed_total += copy_len;
+        } else {
+          s.scan_size = 0;
+          s.merged_data->grouped_compressed_size +=
+              s.sample->count * s.standalone_compressed_size_one;
+        }
+        ++idx_pos;
+        if (pool_used == grouped_in_pool.size()) {
+          break;
+        }
+      }
+
+      if (chunk_uncompressed_total > 0 && !buffers.empty()) {
+        size_t grouped_comp_total = TCMalloc_Internal_ZstdCompress(
+            /*buffers=*/buffers.data(), /*num_buffers=*/buffers.size(),
+            /*dst=*/shared_out.data(), /*dst_capacity=*/shared_out.size());
+
+        for (size_t idx : chunk_indices) {
+          SampleRef& s = samples[idx];
+          if (s.sample->requested_size > 0 && s.scan_size > 0) {
+            size_t grouped_comp = static_cast<size_t>(
+                static_cast<double>(s.scan_size) / chunk_uncompressed_total *
+                grouped_comp_total);
+            if (s.sample->requested_size > s.scan_size) {
+              grouped_comp =
+                  static_cast<size_t>(static_cast<double>(grouped_comp) *
+                                      s.sample->requested_size / s.scan_size);
+            }
+            s.merged_data->grouped_compressed_size +=
+                s.sample->count * grouped_comp;
+          }
+        }
+      }
+    }
+  }
+}
+
 }  // namespace
 
 #if defined(__linux__)
@@ -778,8 +963,8 @@ std::unique_ptr<perftools::profiles::Profile> ProfileBuilder::Finalize() && {
 }
 
 absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
-    const ::tcmalloc::Profile& profile, PageFlagsBase* pageflags,
-    Residency* residency) {
+    const ::tcmalloc::Profile& profile, const MakeProfileOptions& options,
+    PageFlagsBase* pageflags, Residency* residency) {
   if (profile.Type() == ProfileType::kDoNotUse) {
 #if defined(ABSL_HAVE_ADDRESS_SANITIZER) || \
     defined(ABSL_HAVE_LEAK_SANITIZER) ||    \
@@ -815,6 +1000,11 @@ absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
   const int swapped_space_id = builder.InternString("swapped_space");
   const int stale_space_id = builder.InternString("stale_space");
   const int locked_space_id = builder.InternString("locked_space");
+  const int standalone_compressed_space_id =
+      builder.InternString("standalone_compressed_space");
+  const int grouped_compressed_space_id =
+      builder.InternString("grouped_compressed_space");
+  const int zero_space_id = builder.InternString("zero_space");
 
   perftools::profiles::Profile& converted = builder.profile();
 
@@ -861,6 +1051,23 @@ absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
     sample_type->set_unit(bytes_id);
   }
 
+  const bool exporting_compressibility =
+      (profile.Type() == tcmalloc::ProfileType::kHeap &&
+       options.collect_compressibility);
+  if (exporting_compressibility) {
+    perftools::profiles::ValueType* sample_type = converted.add_sample_type();
+    sample_type->set_type(standalone_compressed_space_id);
+    sample_type->set_unit(bytes_id);
+
+    sample_type = converted.add_sample_type();
+    sample_type->set_type(grouped_compressed_space_id);
+    sample_type->set_unit(bytes_id);
+
+    sample_type = converted.add_sample_type();
+    sample_type->set_type(zero_space_id);
+    sample_type->set_unit(bytes_id);
+  }
+
   int default_sample_type_id;
   switch (profile.Type()) {
     case tcmalloc::ProfileType::kFragmentation:
@@ -879,6 +1086,9 @@ absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
 
   SampleMergedMap samples = MergeProfileSamplesAndMaybeGetResidencyInfo(
       profile, pageflags, residency);
+  if (exporting_compressibility) {
+    AddCompressibilityInfo(profile, samples);
+  }
   for (const auto& [entry, data] : samples) {
     perftools::profiles::Profile& profile = builder.profile();
     perftools::profiles::Sample& sample = *profile.add_sample();
@@ -894,6 +1104,11 @@ absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
       sample.add_value(data.stale_size.value_or(0));
       sample.add_value(data.locked_size.value_or(0));
     }
+    if (exporting_compressibility) {
+      sample.add_value(data.standalone_compressed_size);
+      sample.add_value(data.grouped_compressed_size);
+      sample.add_value(data.zero_size);
+    }
 
     // add fields that are common to all memory profiles
     auto add_label = [&](int key, int unit, size_t value) {
@@ -928,7 +1143,7 @@ absl::Status ProfileBuilder::SetDocURL(absl::string_view url) {
 }
 
 absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
-    const ::tcmalloc::Profile& profile) {
+    const ::tcmalloc::Profile& profile, const MakeProfileOptions& options) {
   // Used to populate residency info in heap profile.
   std::optional<PageFlags> pageflags;
   std::optional<ResidencyPageMap> residency;
@@ -941,7 +1156,7 @@ absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
     r = &residency.emplace();
   }
 
-  return MakeProfileProto(profile, p, r);
+  return MakeProfileProto(profile, options, p, r);
 }
 
 }  // namespace tcmalloc_internal

tcmalloc/internal/profile_builder.h

@@ -83,17 +83,24 @@ class ProfileBuilder {
 
 extern const absl::string_view kProfileDropFrames;
 
+// Options controlling protobuf serialization of profiles in MakeProfileProto.
+struct MakeProfileOptions {
+  // If true, collects zero-byte count and compression statistics for each
+  // sampled object when serializing a heap profile.
+  bool collect_compressibility = false;
+};
+
 absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
-    const ::tcmalloc::Profile& profile);
+    const ::tcmalloc::Profile& profile, const MakeProfileOptions& options = {});
 
 class PageFlagsBase;
 class PageFlags;
 class Residency;
 
 // Exposed to facilitate testing.
 absl::StatusOr<std::unique_ptr<perftools::profiles::Profile>> MakeProfileProto(
-    const ::tcmalloc::Profile& profile, PageFlagsBase* pageflags,
-    Residency* residency);
+    const ::tcmalloc::Profile& profile, const MakeProfileOptions& options,
+    PageFlagsBase* pageflags, Residency* residency);
 
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc

tcmalloc/internal/profile_builder_test.cc

@@ -505,7 +505,7 @@ perftools::profiles::Profile MakeTestProfile(
   fake_profile->SetStartTime(start_time);
   fake_profile->SetSamples(std::move(samples));
   Profile profile = ProfileAccessor::MakeProfile(std::move(fake_profile));
-  auto converted_or = MakeProfileProto(profile, p, r);
+  auto converted_or = MakeProfileProto(profile, MakeProfileOptions{}, p, r);
   CHECK_OK(converted_or.status());
   return **converted_or;
 }