1 // Protocol Buffers - Google's data interchange format
2 // Copyright 2008 Google Inc. All rights reserved.
3 //
4 // Use of this source code is governed by a BSD-style
5 // license that can be found in the LICENSE file or at
6 // https://developers.google.com/open-source/licenses/bsd
7
8 #include "google/protobuf/arenaz_sampler.h"
9
10 #include <atomic>
11 #include <cstdint>
12 #include <limits>
13 #include <utility>
14
15
16 // Must be included last.
17 #include "google/protobuf/port_def.inc"
18
19 namespace google {
20 namespace protobuf {
21 namespace internal {
22
GlobalThreadSafeArenazSampler()23 ThreadSafeArenazSampler& GlobalThreadSafeArenazSampler() {
24 static auto* sampler = new ThreadSafeArenazSampler();
25 return *sampler;
26 }
27
UnsampleSlow(ThreadSafeArenaStats * info)28 void UnsampleSlow(ThreadSafeArenaStats* info) {
29 GlobalThreadSafeArenazSampler().Unregister(info);
30 }
31
32 #if defined(PROTOBUF_ARENAZ_SAMPLE)
33 namespace {
34
35 PROTOBUF_CONSTINIT std::atomic<bool> g_arenaz_enabled{true};
36 PROTOBUF_CONSTINIT std::atomic<int32_t> g_arenaz_sample_parameter{1 << 10};
37 PROTOBUF_CONSTINIT std::atomic<ThreadSafeArenazConfigListener>
38 g_arenaz_config_listener{nullptr};
39 PROTOBUF_THREAD_LOCAL absl::profiling_internal::ExponentialBiased
40 g_exponential_biased_generator;
41
TriggerThreadSafeArenazConfigListener()42 void TriggerThreadSafeArenazConfigListener() {
43 auto* listener = g_arenaz_config_listener.load(std::memory_order_acquire);
44 if (listener != nullptr) listener();
45 }
46
47 } // namespace
48
49 PROTOBUF_THREAD_LOCAL SamplingState global_sampling_state = {
50 /*next_sample=*/0, /*sample_stride=*/0};
51
ThreadSafeArenaStats()52 ThreadSafeArenaStats::ThreadSafeArenaStats() { PrepareForSampling(0); }
53 ThreadSafeArenaStats::~ThreadSafeArenaStats() = default;
54
PrepareForSampling()55 void ThreadSafeArenaStats::BlockStats::PrepareForSampling() {
56 num_allocations.store(0, std::memory_order_relaxed);
57 bytes_allocated.store(0, std::memory_order_relaxed);
58 bytes_used.store(0, std::memory_order_relaxed);
59 bytes_wasted.store(0, std::memory_order_relaxed);
60 }
61
PrepareForSampling(int64_t stride)62 void ThreadSafeArenaStats::PrepareForSampling(int64_t stride) {
63 for (auto& blockstats : block_histogram) blockstats.PrepareForSampling();
64 max_block_size.store(0, std::memory_order_relaxed);
65 thread_ids.store(0, std::memory_order_relaxed);
66 weight = stride;
67 // The inliner makes hardcoded skip_count difficult (especially when combined
68 // with LTO). We use the ability to exclude stacks by regex when encoding
69 // instead.
70 depth = absl::GetStackTrace(stack, kMaxStackDepth, /* skip_count= */ 0);
71 }
72
FindBin(size_t bytes)73 size_t ThreadSafeArenaStats::FindBin(size_t bytes) {
74 if (bytes <= kMaxSizeForBinZero) return 0;
75 if (bytes <= kMaxSizeForPenultimateBin) {
76 // absl::bit_width() returns one plus the base-2 logarithm of x, with any
77 // fractional part discarded.
78 return absl::bit_width(absl::bit_ceil(bytes)) - kLogMaxSizeForBinZero - 1;
79 }
80 return kBlockHistogramBins - 1;
81 }
82
MinMaxBlockSizeForBin(size_t bin)83 std::pair<size_t, size_t> ThreadSafeArenaStats::MinMaxBlockSizeForBin(
84 size_t bin) {
85 ABSL_ASSERT(bin < kBlockHistogramBins);
86 if (bin == 0) return {1, kMaxSizeForBinZero};
87 if (bin < kBlockHistogramBins - 1) {
88 return {(1 << (kLogMaxSizeForBinZero + bin - 1)) + 1,
89 1 << (kLogMaxSizeForBinZero + bin)};
90 }
91 return {kMaxSizeForPenultimateBin + 1, std::numeric_limits<size_t>::max()};
92 }
93
RecordAllocateSlow(ThreadSafeArenaStats * info,size_t used,size_t allocated,size_t wasted)94 void RecordAllocateSlow(ThreadSafeArenaStats* info, size_t used,
95 size_t allocated, size_t wasted) {
96 // Update the allocated bytes for the current block.
97 ThreadSafeArenaStats::BlockStats& curr =
98 info->block_histogram[ThreadSafeArenaStats::FindBin(allocated)];
99 curr.bytes_allocated.fetch_add(allocated, std::memory_order_relaxed);
100 curr.num_allocations.fetch_add(1, std::memory_order_relaxed);
101
102 // Update the used and wasted bytes for the previous block.
103 ThreadSafeArenaStats::BlockStats& prev =
104 info->block_histogram[ThreadSafeArenaStats::FindBin(used + wasted)];
105 prev.bytes_used.fetch_add(used, std::memory_order_relaxed);
106 prev.bytes_wasted.fetch_add(wasted, std::memory_order_relaxed);
107
108 if (info->max_block_size.load(std::memory_order_relaxed) < allocated) {
109 info->max_block_size.store(allocated, std::memory_order_relaxed);
110 }
111 const uint64_t tid = 1ULL << (GetCachedTID() % 63);
112 info->thread_ids.fetch_or(tid, std::memory_order_relaxed);
113 }
114
SampleSlow(SamplingState & sampling_state)115 ThreadSafeArenaStats* SampleSlow(SamplingState& sampling_state) {
116 bool first = sampling_state.next_sample < 0;
117 const int64_t next_stride = g_exponential_biased_generator.GetStride(
118 g_arenaz_sample_parameter.load(std::memory_order_relaxed));
119 // Small values of interval are equivalent to just sampling next time.
120 ABSL_ASSERT(next_stride >= 1);
121 sampling_state.next_sample = next_stride;
122 const int64_t old_stride =
123 std::exchange(sampling_state.sample_stride, next_stride);
124
125 // g_arenaz_enabled can be dynamically flipped, we need to set a threshold low
126 // enough that we will start sampling in a reasonable time, so we just use the
127 // default sampling rate.
128 if (!g_arenaz_enabled.load(std::memory_order_relaxed)) return nullptr;
129 // We will only be negative on our first count, so we should just retry in
130 // that case.
131 if (first) {
132 if (PROTOBUF_PREDICT_TRUE(--sampling_state.next_sample > 0)) return nullptr;
133 return SampleSlow(sampling_state);
134 }
135
136 return GlobalThreadSafeArenazSampler().Register(old_stride);
137 }
138
SetThreadSafeArenazConfigListener(ThreadSafeArenazConfigListener l)139 void SetThreadSafeArenazConfigListener(ThreadSafeArenazConfigListener l) {
140 g_arenaz_config_listener.store(l, std::memory_order_release);
141 }
142
IsThreadSafeArenazEnabled()143 bool IsThreadSafeArenazEnabled() {
144 return g_arenaz_enabled.load(std::memory_order_acquire);
145 }
146
SetThreadSafeArenazEnabled(bool enabled)147 void SetThreadSafeArenazEnabled(bool enabled) {
148 SetThreadSafeArenazEnabledInternal(enabled);
149 TriggerThreadSafeArenazConfigListener();
150 }
151
SetThreadSafeArenazEnabledInternal(bool enabled)152 void SetThreadSafeArenazEnabledInternal(bool enabled) {
153 g_arenaz_enabled.store(enabled, std::memory_order_release);
154 }
155
SetThreadSafeArenazSampleParameter(int32_t rate)156 void SetThreadSafeArenazSampleParameter(int32_t rate) {
157 SetThreadSafeArenazSampleParameterInternal(rate);
158 TriggerThreadSafeArenazConfigListener();
159 }
160
SetThreadSafeArenazSampleParameterInternal(int32_t rate)161 void SetThreadSafeArenazSampleParameterInternal(int32_t rate) {
162 if (rate > 0) {
163 g_arenaz_sample_parameter.store(rate, std::memory_order_release);
164 } else {
165 ABSL_RAW_LOG(ERROR, "Invalid thread safe arenaz sample rate: %lld",
166 static_cast<long long>(rate)); // NOLINT(runtime/int)
167 }
168 }
169
ThreadSafeArenazSampleParameter()170 int32_t ThreadSafeArenazSampleParameter() {
171 return g_arenaz_sample_parameter.load(std::memory_order_relaxed);
172 }
173
SetThreadSafeArenazMaxSamples(int32_t max)174 void SetThreadSafeArenazMaxSamples(int32_t max) {
175 SetThreadSafeArenazMaxSamplesInternal(max);
176 TriggerThreadSafeArenazConfigListener();
177 }
178
SetThreadSafeArenazMaxSamplesInternal(int32_t max)179 void SetThreadSafeArenazMaxSamplesInternal(int32_t max) {
180 if (max > 0) {
181 GlobalThreadSafeArenazSampler().SetMaxSamples(max);
182 } else {
183 ABSL_RAW_LOG(ERROR, "Invalid thread safe arenaz max samples: %lld",
184 static_cast<long long>(max)); // NOLINT(runtime/int)
185 }
186 }
187
ThreadSafeArenazMaxSamples()188 size_t ThreadSafeArenazMaxSamples() {
189 return GlobalThreadSafeArenazSampler().GetMaxSamples();
190 }
191
SetThreadSafeArenazGlobalNextSample(int64_t next_sample)192 void SetThreadSafeArenazGlobalNextSample(int64_t next_sample) {
193 if (next_sample >= 0) {
194 global_sampling_state.next_sample = next_sample;
195 global_sampling_state.sample_stride = next_sample;
196 } else {
197 ABSL_RAW_LOG(ERROR, "Invalid thread safe arenaz next sample: %lld",
198 static_cast<long long>(next_sample)); // NOLINT(runtime/int)
199 }
200 }
201
202 #else
SampleSlow(int64_t * next_sample)203 ThreadSafeArenaStats* SampleSlow(int64_t* next_sample) {
204 *next_sample = std::numeric_limits<int64_t>::max();
205 return nullptr;
206 }
207
SetThreadSafeArenazConfigListener(ThreadSafeArenazConfigListener)208 void SetThreadSafeArenazConfigListener(ThreadSafeArenazConfigListener) {}
SetThreadSafeArenazEnabled(bool enabled)209 void SetThreadSafeArenazEnabled(bool enabled) {}
SetThreadSafeArenazEnabledInternal(bool enabled)210 void SetThreadSafeArenazEnabledInternal(bool enabled) {}
IsThreadSafeArenazEnabled()211 bool IsThreadSafeArenazEnabled() { return false; }
SetThreadSafeArenazSampleParameter(int32_t rate)212 void SetThreadSafeArenazSampleParameter(int32_t rate) {}
SetThreadSafeArenazSampleParameterInternal(int32_t rate)213 void SetThreadSafeArenazSampleParameterInternal(int32_t rate) {}
ThreadSafeArenazSampleParameter()214 int32_t ThreadSafeArenazSampleParameter() { return 0; }
SetThreadSafeArenazMaxSamples(int32_t max)215 void SetThreadSafeArenazMaxSamples(int32_t max) {}
SetThreadSafeArenazMaxSamplesInternal(int32_t max)216 void SetThreadSafeArenazMaxSamplesInternal(int32_t max) {}
ThreadSafeArenazMaxSamples()217 size_t ThreadSafeArenazMaxSamples() { return 0; }
SetThreadSafeArenazGlobalNextSample(int64_t next_sample)218 void SetThreadSafeArenazGlobalNextSample(int64_t next_sample) {}
219 #endif // defined(PROTOBUF_ARENAZ_SAMPLE)
220
221 } // namespace internal
222 } // namespace protobuf
223 } // namespace google
224
225 #include "google/protobuf/port_undef.inc"
226