From ba0e2f581fe20592ba80a00059393378528030a4 Mon Sep 17 00:00:00 2001 From: Todd Lipcon Date: Wed, 22 Jul 2026 20:06:33 -0700 Subject: [PATCH] Add MakeProfileOptions::collect_compressibility and internal compressibility scanning support in profile_builder for heap profiles. PiperOrigin-RevId: 952480217 --- tcmalloc/internal/BUILD | 6 +- tcmalloc/internal/parameter_accessors.h | 9 + tcmalloc/internal/profile_builder.cc | 225 +++++++++++++++++++++- tcmalloc/internal/profile_builder.h | 13 +- tcmalloc/internal/profile_builder_test.cc | 2 +- tcmalloc/internal/util.cc | 18 ++ tcmalloc/internal/util.h | 6 + 7 files changed, 267 insertions(+), 12 deletions(-) diff --git a/tcmalloc/internal/BUILD b/tcmalloc/internal/BUILD index 991f77a5b..680809025 100644 --- a/tcmalloc/internal/BUILD +++ b/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", diff --git a/tcmalloc/internal/parameter_accessors.h b/tcmalloc/internal/parameter_accessors.h index 8d93888cd..2f3f9d9c2 100644 --- a/tcmalloc/internal/parameter_accessors.h +++ b/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* 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_ diff --git a/tcmalloc/internal/profile_builder.cc b/tcmalloc/internal/profile_builder.cc index 12b03f7a0..71dff6d44 100644 --- a/tcmalloc/internal/profile_builder.cc +++ b/tcmalloc/internal/profile_builder.cc @@ -44,6 +44,7 @@ #include #include #include +#include #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 stale_size; std::optional locked_size; std::optional 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->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 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(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(static_cast(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(static_cast(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, + 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 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 buffers; + std::vector 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( + static_cast(s.scan_size) / chunk_uncompressed_total * + grouped_comp_total); + if (s.sample->requested_size > s.scan_size) { + grouped_comp = + static_cast(static_cast(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 ProfileBuilder::Finalize() && { } absl::StatusOr> 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> 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> 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> 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> 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> MakeProfileProto( - const ::tcmalloc::Profile& profile) { + const ::tcmalloc::Profile& profile, const MakeProfileOptions& options) { // Used to populate residency info in heap profile. std::optional pageflags; std::optional residency; @@ -941,7 +1156,7 @@ absl::StatusOr> MakeProfileProto( r = &residency.emplace(); } - return MakeProfileProto(profile, p, r); + return MakeProfileProto(profile, options, p, r); } } // namespace tcmalloc_internal diff --git a/tcmalloc/internal/profile_builder.h b/tcmalloc/internal/profile_builder.h index eebffa38a..63e0a2b1f 100644 --- a/tcmalloc/internal/profile_builder.h +++ b/tcmalloc/internal/profile_builder.h @@ -83,8 +83,15 @@ 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> MakeProfileProto( - const ::tcmalloc::Profile& profile); + const ::tcmalloc::Profile& profile, const MakeProfileOptions& options = {}); class PageFlagsBase; class PageFlags; @@ -92,8 +99,8 @@ class Residency; // Exposed to facilitate testing. absl::StatusOr> 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 diff --git a/tcmalloc/internal/profile_builder_test.cc b/tcmalloc/internal/profile_builder_test.cc index 80790f400..73f86cea5 100644 --- a/tcmalloc/internal/profile_builder_test.cc +++ b/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; } diff --git a/tcmalloc/internal/util.cc b/tcmalloc/internal/util.cc index 133c4d0b6..abb8c6513 100644 --- a/tcmalloc/internal/util.cc +++ b/tcmalloc/internal/util.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include #include @@ -123,6 +124,23 @@ ssize_t signal_safe_read(int fd, char* buf, size_t count, size_t* bytes_read) { return rc; } +bool SafeCopyMemory(const void* src, void* dst, size_t size) { + if (src == nullptr || dst == nullptr || size == 0) { + return false; + } + struct iovec local_iov; + local_iov.iov_base = dst; + local_iov.iov_len = size; + + struct iovec remote_iov; + remote_iov.iov_base = const_cast(src); + remote_iov.iov_len = size; + + ssize_t bytes = process_vm_readv(getpid(), &local_iov, /*liovcnt=*/1, + &remote_iov, /*riovcnt=*/1, /*flags=*/0); + return bytes == static_cast(size); +} + } // namespace tcmalloc_internal } // namespace tcmalloc GOOGLE_MALLOC_SECTION_END diff --git a/tcmalloc/internal/util.h b/tcmalloc/internal/util.h index d7e3b77e0..a59a5ff12 100644 --- a/tcmalloc/internal/util.h +++ b/tcmalloc/internal/util.h @@ -92,6 +92,12 @@ ssize_t signal_safe_read(int fd, char* buf, size_t count, size_t* bytes_read); // involved with the latter. int signal_safe_poll(struct ::pollfd* fds, int nfds, absl::Duration timeout); +// Copies `size` bytes from `src` to `dst` safely using process_vm_readv. +// If `src` points to unmapped or concurrently freed/mprotected memory, this +// function returns false without triggering a SIGSEGV crash. Returns true if +// exactly `size` bytes were successfully copied. +bool SafeCopyMemory(const void* src, void* dst, size_t size); + class ScopedSigmask { public: // Masks all signal handlers. (SIG_SETMASK, All)