// Protocol Buffers - Google's data interchange format

// Copyright 2008 Google Inc. All rights reserved.

//

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file or at

// https://developers.google.com/open-source/licenses/bsd

#ifndef GOOGLE_PROTOBUF_ARENAZ_SAMPLER_H__

#define GOOGLE_PROTOBUF_ARENAZ_SAMPLER_H__

#include <array>

#include <atomic>

#include <cstddef>

#include <cstdint>

#include <utility>

// Must be included last.

#include "google/protobuf/port_def.inc"

namespace google {

namespace protobuf {

namespace internal {

#if defined(PROTOBUF_ARENAZ_SAMPLE)

struct ThreadSafeArenaStats;

void RecordAllocateSlow(ThreadSafeArenaStats* info, size_t used,

size_t allocated, size_t wasted);

// Stores information about a sampled thread safe arena. All mutations to this

// *must* be made through `Record*` functions below. All reads from this *must*

// only occur in the callback to `ThreadSafeArenazSampler::Iterate`.

struct [[nodiscard]] ThreadSafeArenaStats

: public absl::profiling_internal::Sample<ThreadSafeArenaStats> {

// Constructs the object but does not fill in any fields.

ThreadSafeArenaStats();

~ThreadSafeArenaStats();

// Puts the object into a clean state, fills in the logically `const` members,

// blocking for any readers that are currently sampling the object. The

// 'stride' parameter is the number of ThreadSafeArenas that were instantiated

// between this sample and the previous one.

void PrepareForSampling(int64_t stride)

ABSL_EXCLUSIVE_LOCKS_REQUIRED(init_mu);

// These fields are mutated by the various Record* APIs and need to be

// thread-safe.

struct BlockStats {

std::atomic<int> num_allocations;

std::atomic<size_t> bytes_allocated;

std::atomic<size_t> bytes_used;

std::atomic<size_t> bytes_wasted;

void PrepareForSampling();

};

// block_histogram is a kBlockHistogramBins sized histogram. The zeroth bin

// stores info about blocks of size \in [1, 1 << kLogMaxSizeForBinZero]. Bin

// i, where i > 0, stores info for blocks of size \in (max_size_bin (i-1),

// 1 << (kLogMaxSizeForBinZero + i)]. The final bin stores info about blocks

// of size \in [kMaxSizeForPenultimateBin + 1,

// std::numeric_limits<size_t>::max()].

static constexpr size_t kBlockHistogramBins = 15;

static constexpr size_t kLogMaxSizeForBinZero = 7;

static constexpr size_t kMaxSizeForBinZero = (1 << kLogMaxSizeForBinZero);

static constexpr size_t kMaxSizeForPenultimateBin =

1 << (kLogMaxSizeForBinZero + kBlockHistogramBins - 2);

std::array<BlockStats, kBlockHistogramBins> block_histogram;

// Records the largest block allocated for the arena.

std::atomic<size_t> max_block_size;

// Bit `i` is set to 1 indicates that a thread with `tid % 63 = i` accessed

// the underlying arena. We use `% 63` as a rudimentary hash to ensure some

// bit mixing for thread-ids; `% 64` would only grab the low bits and might

// create sampling artifacts.

std::atomic<uint64_t> thread_ids;

// All of the fields below are set by `PrepareForSampling`, they must not

// be mutated in `Record*` functions. They are logically `const` in that

// sense. These are guarded by init_mu, but that is not externalized to

// clients, who can only read them during

// `ThreadSafeArenazSampler::Iterate` which will hold the lock.

static constexpr int kMaxStackDepth = 64;

int32_t depth;

void* stack[kMaxStackDepth];

static void RecordAllocateStats(ThreadSafeArenaStats* info, size_t used,

size_t allocated, size_t wasted) {

if (ABSL_PREDICT_TRUE(info == nullptr)) return;

RecordAllocateSlow(info, used, allocated, wasted);

}

// Returns the bin for the provided size.

static size_t FindBin(size_t bytes);

// Returns the min and max bytes that can be stored in the histogram for

// blocks in the provided bin.

static std::pair<size_t, size_t> MinMaxBlockSizeForBin(size_t bin);

};

struct SamplingState {

// Number of ThreadSafeArenas that should be instantiated before the next

// ThreadSafeArena is sampled. This variable is decremented with each

// instantiation.

int64_t next_sample;

// When we make a sampling decision, we record that distance between from the

// previous sample so we can weight each sample. 'distance' here is the

// number of instantiations of ThreadSafeArena.

int64_t sample_stride;

};

[[nodiscard]] ThreadSafeArenaStats* SampleSlow(SamplingState& sampling_state);

void UnsampleSlow(ThreadSafeArenaStats* info);

class ThreadSafeArenaStatsHandle {

public:

explicit ThreadSafeArenaStatsHandle() = default;

explicit ThreadSafeArenaStatsHandle(ThreadSafeArenaStats* info)

: info_(info) {}

~ThreadSafeArenaStatsHandle() {

if (ABSL_PREDICT_TRUE(info_ == nullptr)) return;

UnsampleSlow(info_);

}

ThreadSafeArenaStatsHandle(ThreadSafeArenaStatsHandle&& other) noexcept

: info_(std::exchange(other.info_, nullptr)) {}

ThreadSafeArenaStatsHandle& operator=(

ThreadSafeArenaStatsHandle&& other) noexcept {

if (ABSL_PREDICT_FALSE(info_ != nullptr)) {

UnsampleSlow(info_);

}

info_ = std::exchange(other.info_, nullptr);

return *this;

}

ThreadSafeArenaStats* MutableStats() { return info_; }

friend void swap(ThreadSafeArenaStatsHandle& lhs,

ThreadSafeArenaStatsHandle& rhs) noexcept {

std::swap(lhs.info_, rhs.info_);

}

friend class ThreadSafeArenaStatsHandlePeer;

private:

ThreadSafeArenaStats* info_ = nullptr;

};

using ThreadSafeArenazSampler =

::absl::profiling_internal::SampleRecorder<ThreadSafeArenaStats>;

extern PROTOBUF_THREAD_LOCAL SamplingState global_sampling_state;

// Returns an RAII sampling handle that manages registration and unregistation

// with the global sampler.

[[nodiscard]] inline ThreadSafeArenaStatsHandle Sample() {

if (ABSL_PREDICT_TRUE(--global_sampling_state.next_sample > 0)) {

return ThreadSafeArenaStatsHandle(nullptr);

}

return ThreadSafeArenaStatsHandle(SampleSlow(global_sampling_state));

}

#else

using SamplingState = int64_t;

struct ThreadSafeArenaStats {

static void RecordAllocateStats(ThreadSafeArenaStats*, size_t /*requested*/,

size_t /*allocated*/, size_t /*wasted*/) {}

};

[[nodiscard]] ThreadSafeArenaStats* SampleSlow(SamplingState& next_sample);

void UnsampleSlow(ThreadSafeArenaStats* info);

class [[nodiscard]] ThreadSafeArenaStatsHandle {

public:

explicit ThreadSafeArenaStatsHandle() = default;

explicit ThreadSafeArenaStatsHandle(ThreadSafeArenaStats*) {}

void RecordReset() {}

ThreadSafeArenaStats* MutableStats() { return nullptr; }

friend void swap(ThreadSafeArenaStatsHandle&, ThreadSafeArenaStatsHandle&) {}

private:

friend class ThreadSafeArenaStatsHandlePeer;

};

class ThreadSafeArenazSampler {

public:

void Unregister(ThreadSafeArenaStats*) {}

void SetMaxSamples(int32_t) {}

};

// Returns an RAII sampling handle that manages registration and unregistation

// with the global sampler.

[[nodiscard]] inline ThreadSafeArenaStatsHandle Sample() {

return ThreadSafeArenaStatsHandle(nullptr);

}

#endif // defined(PROTOBUF_ARENAZ_SAMPLE)

// Returns a global Sampler.

ThreadSafeArenazSampler& GlobalThreadSafeArenazSampler();

using ThreadSafeArenazConfigListener = void (*)();

void SetThreadSafeArenazConfigListener(ThreadSafeArenazConfigListener l);

// Enables or disables sampling for thread safe arenas.

void SetThreadSafeArenazEnabled(bool enabled);

void SetThreadSafeArenazEnabledInternal(bool enabled);

// Returns true if sampling is on, false otherwise.

[[nodiscard]] bool IsThreadSafeArenazEnabled();

// Sets the rate at which thread safe arena will be sampled.

void SetThreadSafeArenazSampleParameter(int32_t rate);

void SetThreadSafeArenazSampleParameterInternal(int32_t rate);

// Returns the rate at which thread safe arena will be sampled.

[[nodiscard]] int32_t ThreadSafeArenazSampleParameter();

// Sets a soft max for the number of samples that will be kept.

void SetThreadSafeArenazMaxSamples(int32_t max);

void SetThreadSafeArenazMaxSamplesInternal(int32_t max);

// Returns the max number of samples that will be kept.

[[nodiscard]] size_t ThreadSafeArenazMaxSamples();

// Sets the current value for when arenas should be next sampled.

void SetThreadSafeArenazGlobalNextSample(int64_t next_sample);

} // namespace internal

} // namespace protobuf

} // namespace google

#include "google/protobuf/port_undef.inc"

#endif // GOOGLE_PROTOBUF_ARENAZ_SAMPLER_H__