// Copyright 2022 The TCMalloc Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "tcmalloc/deallocation_profiler.h"

#include <algorithm>
#include <cmath>  // for std::lround
#include <cstddef>
#include <cstdint>  // for uintptr_t
#include <cstring>
#include <functional>
#include <iterator>
#include <limits>
#include <memory>
#include <new>
#include <optional>
#include <type_traits>
#include <utility>

#include "absl/base/attributes.h"
#include "absl/base/const_init.h"
#include "absl/base/internal/low_level_alloc.h"
#include "absl/base/internal/spinlock.h"
#include "absl/base/internal/sysinfo.h"
#include "absl/base/macros.h"
#include "absl/container/flat_hash_map.h"
#include "absl/debugging/stacktrace.h"  // for GetStackTrace
#include "absl/functional/function_ref.h"
#include "absl/hash/hash.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "tcmalloc/internal/allocation_guard.h"
#include "tcmalloc/internal/cache_topology.h"
#include "tcmalloc/internal/config.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/internal/percpu.h"
#include "tcmalloc/internal/sampled_allocation.h"
#include "tcmalloc/internal_malloc_extension.h"
#include "tcmalloc/malloc_extension.h"
#include "tcmalloc/sampler.h"
#include "tcmalloc/static_vars.h"

GOOGLE_MALLOC_SECTION_BEGIN
namespace tcmalloc {
namespace deallocationz {
namespace {
using ::absl::base_internal::SpinLock;
using tcmalloc_internal::AllocationGuardSpinLockHolder;

// STL adaptor for an arena based allocator which provides the following:
//   static void* Alloc::Allocate(size_t size);
//   static void Alloc::Free(void* ptr, size_t size);
template <typename T, class Alloc>
class AllocAdaptor final {
 public:
  using value_type = T;

  AllocAdaptor() = default;
  AllocAdaptor(const AllocAdaptor&) = default;

  template <class T1>
  using rebind = AllocAdaptor<T1, Alloc>;

  template <class T1>
  explicit AllocAdaptor(const AllocAdaptor<T1, Alloc>&) {}

  T* allocate(size_t n) {
    // Check if n is too big to allocate.
    TC_ASSERT_EQ((n * sizeof(T)) / sizeof(T), n);
    return static_cast<T*>(Alloc::Allocate(n * sizeof(T)));
  }
  void deallocate(T* p, size_t n) { Alloc::Free(p, n * sizeof(T)); }
};

// Stores stack traces and metadata for any allocation or deallocation
// encountered by the profiler.
struct DeallocationSampleRecord {
  tcmalloc_internal::StackTrace stack_trace;
  int cpu_id = -1;
  int vcpu_id = -1;
  int l3_id = -1;
  int numa_id = -1;
  pid_t thread_id = 0;

  template <typename H>
  friend H AbslHashValue(H h, const DeallocationSampleRecord& c) {
    return H::combine(H::combine_contiguous(std::move(h), c.stack_trace.stack,
                                            c.stack_trace.depth),
                      c.stack_trace.depth, c.stack_trace.requested_size,
                      c.stack_trace.requested_alignment,
                      c.stack_trace.allocated_size,
                      c.stack_trace.requested_size_returning,
                      c.stack_trace.access_hint, c.stack_trace.cold_allocated,
                      c.stack_trace.token_id, c.stack_trace.guarded_status,
                      c.stack_trace.allocation_type);
  }

  bool operator==(const DeallocationSampleRecord& other) const {
    if (stack_trace.depth != other.stack_trace.depth ||
        stack_trace.requested_size != other.stack_trace.requested_size ||
        stack_trace.requested_alignment !=
            other.stack_trace.requested_alignment ||
        stack_trace.allocated_size != other.stack_trace.allocated_size ||
        stack_trace.requested_size_returning !=
            other.stack_trace.requested_size_returning ||
        stack_trace.access_hint != other.stack_trace.access_hint ||
        stack_trace.cold_allocated != other.stack_trace.cold_allocated ||
        stack_trace.token_id != other.stack_trace.token_id ||
        stack_trace.guarded_status != other.stack_trace.guarded_status ||
        stack_trace.allocation_type != other.stack_trace.allocation_type) {
      return false;
    }
    return std::equal(stack_trace.stack, stack_trace.stack + stack_trace.depth,
                      other.stack_trace.stack);
  }
};

// Tracks whether an object was allocated/deallocated by the same CPU/thread.
struct CpuThreadMatchingStatus {
  constexpr CpuThreadMatchingStatus(bool physical_cpu_matched,
                                    bool virtual_cpu_matched, bool l3_matched,
                                    bool numa_matched, bool thread_matched)
      : physical_cpu_matched(physical_cpu_matched),
        virtual_cpu_matched(virtual_cpu_matched),
        l3_matched(l3_matched),
        numa_matched(numa_matched),
        thread_matched(thread_matched),
        value((static_cast<int>(physical_cpu_matched) << 4) |
              (static_cast<int>(virtual_cpu_matched) << 3) |
              (static_cast<int>(l3_matched) << 2) |
              (static_cast<int>(numa_matched) << 1) |
              static_cast<int>(thread_matched)) {}
  bool physical_cpu_matched;
  bool virtual_cpu_matched;
  bool l3_matched;
  bool numa_matched;
  bool thread_matched;
  int value;
};

struct RpcMatchingStatus {
  static constexpr int ComputeValue(uint64_t alloc, uint64_t dealloc) {
    if (alloc != 0 && dealloc != 0) {
      return static_cast<int>(alloc == dealloc);
    } else {
      return 2;
    }
  }

  constexpr RpcMatchingStatus(uint64_t alloc, uint64_t dealloc)
      : value(ComputeValue(alloc, dealloc)) {}

  int value;
};

int ComputeIndex(CpuThreadMatchingStatus status, RpcMatchingStatus rpc_status) {
  return status.value * 3 + rpc_status.value;
}

int GetL3Id(int cpu_id) {
  return cpu_id >= 0
             ? tcmalloc_internal::CacheTopology::Instance().GetL3FromCpuId(
                   cpu_id)
             : -1;
}

int GetNumaId(int cpu_id) {
  return cpu_id >= 0
             ? tcmalloc_internal::tc_globals.numa_topology().GetCpuPartition(
                   cpu_id)
             : -1;
}

constexpr std::pair<CpuThreadMatchingStatus, RpcMatchingStatus> kAllCases[] = {
    // clang-format off
    {CpuThreadMatchingStatus(false, false, false, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, false, false, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, false, false, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, false, false, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(false, false, false, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, false, false, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, false, false, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, false, false, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(false, false, false, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, false, false, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, false, false, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, false, false, true, true), RpcMatchingStatus(1, 1)},

    {CpuThreadMatchingStatus(false, false, true, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, false, true, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, false, true, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, false, true, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(false, false, true, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, false, true, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, false, true, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, false, true, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(false, false, true, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, false, true, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, false, true, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, false, true, true, true), RpcMatchingStatus(1, 1)},

    {CpuThreadMatchingStatus(false, true, false, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, true, false, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, true, false, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, true, false, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(false, true, false, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, true, false, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, true, false, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, true, false, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(false, true, false, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, true, false, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, true, false, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, true, false, true, true), RpcMatchingStatus(1, 1)},

    {CpuThreadMatchingStatus(false, true, true, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, true, true, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, true, true, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(false, true, true, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(false, true, true, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, true, true, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, true, true, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(false, true, true, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(false, true, true, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, true, true, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, true, true, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(false, true, true, true, true), RpcMatchingStatus(1, 1)},

    {CpuThreadMatchingStatus(true, false, false, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, false, false, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, false, false, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, false, false, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(true, false, false, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, false, false, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, false, false, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, false, false, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(true, false, false, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, false, false, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, false, false, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, false, false, true, true), RpcMatchingStatus(1, 1)},

    {CpuThreadMatchingStatus(true, false, true, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, false, true, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, false, true, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, false, true, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(true, false, true, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, false, true, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, false, true, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, false, true, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(true, false, true, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, false, true, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, false, true, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, false, true, true, true), RpcMatchingStatus(1, 1)},

    {CpuThreadMatchingStatus(true, true, false, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, true, false, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, true, false, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, true, false, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(true, true, false, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, true, false, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, true, false, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, true, false, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(true, true, false, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, true, false, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, true, false, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, true, false, true, true), RpcMatchingStatus(1, 1)},

    {CpuThreadMatchingStatus(true, true, true, false, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, true, true, false, true), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, true, true, true, false), RpcMatchingStatus(0, 0)},
    {CpuThreadMatchingStatus(true, true, true, true, true), RpcMatchingStatus(0, 0)},

    {CpuThreadMatchingStatus(true, true, true, false, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, true, true, false, true), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, true, true, true, false), RpcMatchingStatus(1, 2)},
    {CpuThreadMatchingStatus(true, true, true, true, true), RpcMatchingStatus(1, 2)},

    {CpuThreadMatchingStatus(true, true, true, false, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, true, true, false, true), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, true, true, true, false), RpcMatchingStatus(1, 1)},
    {CpuThreadMatchingStatus(true, true, true, true, true), RpcMatchingStatus(1, 1)},
    // clang-format on
};
}  // namespace

class DeallocationProfiler {
 private:
  // Arena and allocator used to back STL objects used by DeallocationProfiler
  // Shared between all instances of DeallocationProfiler
  // TODO(b/248332543): Use TCMalloc's own arena allocator instead of defining a
  // new one here. The need for refcount management could be the reason for
  // using a custom allocator in the first place.
  class MyAllocator {
   public:
    static void* Allocate(size_t n) {
      return absl::base_internal::LowLevelAlloc::AllocWithArena(n, arena_);
    }
    static void Free(const void* p, size_t /* n */) {
      absl::base_internal::LowLevelAlloc::Free(const_cast<void*>(p));
    }

    // The lifetime of the arena is managed using a reference count and
    // determined by how long at least one emitted Profile remains alive.
    struct LowLevelArenaReference {
      LowLevelArenaReference() {
        AllocationGuardSpinLockHolder h(arena_lock_);
        if ((refcount_++) == 0) {
          TC_CHECK_EQ(arena_, nullptr);
          arena_ = absl::base_internal::LowLevelAlloc::NewArena(0);
        }
      }

      ~LowLevelArenaReference() {
        AllocationGuardSpinLockHolder h(arena_lock_);
        if ((--refcount_) == 0) {
          TC_CHECK(absl::base_internal::LowLevelAlloc::DeleteArena(arena_));
          arena_ = nullptr;
        }
      }
    };

   private:
    // We need to protect the arena with a mutex and ensure that every thread
    // acquires that mutex before it uses the arena for the first time. Once
    // it has acquired the mutex, it is guaranteed that arena won't change
    // between that point in time and when the thread stops accessing it (as
    // enforced by LowLevelArenaReference below).
    ABSL_CONST_INIT static SpinLock arena_lock_;
    static absl::base_internal::LowLevelAlloc::Arena* arena_;

    // We assume that launching a new deallocation profiler takes too long
    // to cause this to overflow within the sampling period. The reason this
    // is not using std::shared_ptr is that we do not only need to protect the
    // value of the reference count but also the pointer itself (and therefore
    // need a separate mutex either way).
    static uint32_t refcount_;
  };

  // This must be the first member of the class to be initialized. The
  // underlying arena must stay alive as long as the profiler.
  MyAllocator::LowLevelArenaReference arena_ref_;

  // All active profilers are stored in a list.
  DeallocationProfiler* next_;
  DeallocationProfilerList* list_ = nullptr;
  friend class DeallocationProfilerList;

  using AllocsTable = absl::flat_hash_map<
      tcmalloc_internal::AllocHandle, DeallocationSampleRecord,
      absl::Hash<tcmalloc_internal::AllocHandle>,
      std::equal_to<tcmalloc_internal::AllocHandle>,
      AllocAdaptor<std::pair<const tcmalloc_internal::AllocHandle,
                             DeallocationSampleRecord>,
                   MyAllocator>>;

  class DeallocationStackTraceTable
      final : public tcmalloc_internal::ProfileBase {
   public:
    // We define the dtor to ensure it is placed in the desired text section.
    ~DeallocationStackTraceTable() override = default;
    void AddTrace(const DeallocationSampleRecord& alloc_trace,
                  const DeallocationSampleRecord& dealloc_trace);

    void Iterate(
        absl::FunctionRef<void(const Profile::Sample&)> func) const override;

    ProfileType Type() const override {
      return tcmalloc::ProfileType::kLifetimes;
    }

    std::optional<absl::Time> StartTime() const override { return start_time_; }

    absl::Duration Duration() const override {
      return stop_time_ - start_time_;
    }

    void StopAndRecord(const AllocsTable& allocs);

   private:
    // This must be the first member of the class to be initialized. The
    // underlying arena must stay alive as long as the profile.
    MyAllocator::LowLevelArenaReference arena_ref_;

    static constexpr int kNumCases = std::size(kAllCases);

    struct Key {
      DeallocationSampleRecord alloc;
      DeallocationSampleRecord dealloc;

      Key(const DeallocationSampleRecord& alloc,
          const DeallocationSampleRecord& dealloc)
          : alloc(alloc), dealloc(dealloc) {}

      template <typename H>
      friend H AbslHashValue(H h, const Key& c) {
        return H::combine(std::move(h), c.alloc, c.dealloc);
      }

      bool operator==(const Key& other) const {
        return (alloc == other.alloc) && (dealloc == other.dealloc);
      }
    };

    struct Value {
      // for each possible cases, we collect repetition count and avg lifetime
      // we also collect the minimum and maximum lifetimes, as well as the sum
      // of squares (to calculate the standard deviation).
      double counts[kNumCases] = {0.0};
      double mean_life_times_ns[kNumCases] = {0.0};
      double variance_life_times_ns[kNumCases] = {0.0};
      double min_life_times_ns[kNumCases] = {0.0};
      double max_life_times_ns[kNumCases] = {0.0};

      Value() {
        std::fill_n(min_life_times_ns, kNumCases,
                    std::numeric_limits<double>::max());
      }
    };

    absl::flat_hash_map<Key, Value, absl::Hash<Key>, std::equal_to<Key>,
                        AllocAdaptor<std::pair<const Key, Value>, MyAllocator>>
        table_;

    absl::Time start_time_ = absl::Now();
    absl::Time stop_time_;
  };

  // Keep track of allocations that are in flight
  AllocsTable allocs_;

  // Table to store lifetime information collected by this profiler
  std::unique_ptr<DeallocationStackTraceTable> reports_ = nullptr;

 public:
  explicit DeallocationProfiler(DeallocationProfilerList* list) : list_(list) {
    reports_ = std::make_unique<DeallocationStackTraceTable>();
    list_->Add(this);
  }

  ~DeallocationProfiler() {
    if (reports_ != nullptr) {
      Stop();
    }
  }

  tcmalloc::Profile Stop() {
    if (reports_ != nullptr) {
      // We first remove the profiler from the list to avoid racing with
      // potential allocations which may modify the allocs_ table.
      list_->Remove(this);
      reports_->StopAndRecord(allocs_);
      return tcmalloc_internal::ProfileAccessor::MakeProfile(
          std::move(reports_));
    }
    return tcmalloc::Profile();
  }

  void ReportMalloc(const tcmalloc_internal::StackTrace& stack_trace) {
    // store sampled alloc in the hashmap
    DeallocationSampleRecord& allocation =
        allocs_[stack_trace.sampled_alloc_handle];

    allocation.stack_trace = stack_trace;
    // TODO(mmaas): Do we need to worry about b/65384231 anymore?
    allocation.cpu_id = tcmalloc_internal::subtle::percpu::GetRealCpu();
    allocation.vcpu_id = tcmalloc_internal::subtle::percpu::VirtualCpu::get();
    allocation.l3_id = GetL3Id(allocation.cpu_id);
    allocation.numa_id = GetNumaId(allocation.cpu_id);
    allocation.thread_id = absl::base_internal::GetTID();
  }

  void ReportFree(tcmalloc_internal::AllocHandle handle) {
    auto it = allocs_.find(handle);

    // Handle the case that we observed the deallocation but not the allocation
    if (it == allocs_.end()) {
      return;
    }

    DeallocationSampleRecord sample = it->second;
    allocs_.erase(it);

    DeallocationSampleRecord deallocation;
    deallocation.stack_trace = sample.stack_trace;
    deallocation.stack_trace.allocation_time = absl::Now();
    deallocation.cpu_id = tcmalloc_internal::subtle::percpu::GetRealCpu();
    deallocation.vcpu_id = tcmalloc_internal::subtle::percpu::VirtualCpu::get();
    deallocation.l3_id = GetL3Id(deallocation.cpu_id);
    deallocation.numa_id = GetNumaId(deallocation.cpu_id);
    deallocation.thread_id = absl::base_internal::GetTID();
    deallocation.stack_trace.depth = absl::GetStackTrace(
        deallocation.stack_trace.stack, tcmalloc_internal::kMaxStackDepth, 1);

    reports_->AddTrace(sample, deallocation);
  }
};

void DeallocationProfilerList::Add(DeallocationProfiler* profiler) {
  AllocationGuardSpinLockHolder h(profilers_lock_);
  profiler->next_ = first_;
  first_ = profiler;

  // Whenever a new profiler is created, we seed it with live allocations.
  tcmalloc_internal::tc_globals.sampled_allocation_recorder().Iterate(
      [profiler](
          const tcmalloc_internal::SampledAllocation& sampled_allocation) {
        profiler->ReportMalloc(sampled_allocation.sampled_stack);
      });
}

// This list is very short and we're nowhere near a hot path, just walk
void DeallocationProfilerList::Remove(DeallocationProfiler* profiler) {
  AllocationGuardSpinLockHolder h(profilers_lock_);
  DeallocationProfiler** link = &first_;
  DeallocationProfiler* cur = first_;
  while (cur != profiler) {
    TC_CHECK_NE(cur, nullptr);
    link = &cur->next_;
    cur = cur->next_;
  }
  *link = profiler->next_;
}

void DeallocationProfilerList::ReportMalloc(
    const tcmalloc_internal::StackTrace& stack_trace) {
  AllocationGuardSpinLockHolder h(profilers_lock_);
  DeallocationProfiler* cur = first_;
  while (cur != nullptr) {
    cur->ReportMalloc(stack_trace);
    cur = cur->next_;
  }
}

void DeallocationProfilerList::ReportFree(
    tcmalloc_internal::AllocHandle handle) {
  AllocationGuardSpinLockHolder h(profilers_lock_);
  DeallocationProfiler* cur = first_;
  while (cur != nullptr) {
    cur->ReportFree(handle);
    cur = cur->next_;
  }
}

// Initialize static variables
absl::base_internal::LowLevelAlloc::Arena*
    DeallocationProfiler::MyAllocator::arena_ = nullptr;
uint32_t DeallocationProfiler::MyAllocator::refcount_ = 0;
ABSL_CONST_INIT SpinLock DeallocationProfiler::MyAllocator::arena_lock_(
    absl::base_internal::SCHEDULE_KERNEL_ONLY);

void DeallocationProfiler::DeallocationStackTraceTable::StopAndRecord(
    const AllocsTable& allocs) {
  stop_time_ = absl::Now();

  // Insert a dummy DeallocationSampleRecord since the table stores pairs. This
  // allows us to make minimal changes to the rest of the sample processing
  // steps reducing special casing for censored samples. This also allows us to
  // aggregate censored samples just like regular deallocation samples.
  for (const auto& [unused, alloc] : allocs) {
    DeallocationSampleRecord censored;
    censored.stack_trace = alloc.stack_trace;
    censored.stack_trace.allocation_time = stop_time_;
    censored.stack_trace.depth = 0;
    AddTrace(alloc, censored);
  }
}

void DeallocationProfiler::DeallocationStackTraceTable::AddTrace(
    const DeallocationSampleRecord& alloc_trace,
    const DeallocationSampleRecord& dealloc_trace) {
  CpuThreadMatchingStatus status =
      CpuThreadMatchingStatus(alloc_trace.cpu_id == dealloc_trace.cpu_id,
                              alloc_trace.vcpu_id == dealloc_trace.vcpu_id,
                              alloc_trace.l3_id == dealloc_trace.l3_id,
                              alloc_trace.numa_id == dealloc_trace.numa_id,
                              alloc_trace.thread_id == dealloc_trace.thread_id);

  // Initialize a default rpc matched status.
  RpcMatchingStatus rpc_status(/*alloc=*/0, /*dealloc=*/0);

  const int index = ComputeIndex(status, rpc_status);

  DeallocationStackTraceTable::Value& v =
      table_[DeallocationStackTraceTable::Key(alloc_trace, dealloc_trace)];

  const absl::Duration life_time = dealloc_trace.stack_trace.allocation_time -
                                   alloc_trace.stack_trace.allocation_time;
  double life_time_ns = absl::ToDoubleNanoseconds(life_time);

  // Update mean and variance using Welford’s online algorithm.
  TC_ASSERT_LT(index, std::size(v.counts));

  double old_mean_ns = v.mean_life_times_ns[index];
  v.mean_life_times_ns[index] +=
      (life_time_ns - old_mean_ns) / static_cast<double>(v.counts[index] + 1);
  v.variance_life_times_ns[index] +=
      (life_time_ns - v.mean_life_times_ns[index]) *
      (v.mean_life_times_ns[index] - old_mean_ns);

  v.min_life_times_ns[index] =
      std::min(v.min_life_times_ns[index], life_time_ns);
  v.max_life_times_ns[index] =
      std::max(v.max_life_times_ns[index], life_time_ns);
  v.counts[index]++;
}

void DeallocationProfiler::DeallocationStackTraceTable::Iterate(
    absl::FunctionRef<void(const Profile::Sample&)> func) const {
  uint64_t pair_id = 1;

  for (auto& it : table_) {
    const Key& k = it.first;
    const Value& v = it.second;

    // Report total bytes that are a multiple of the object size.
    size_t allocated_size = k.alloc.stack_trace.allocated_size;

    for (const auto& matching_case : kAllCases) {
      const int index = ComputeIndex(matching_case.first, matching_case.second);
      if (v.counts[index] == 0) {
        continue;
      }

      uintptr_t bytes =
          std::lround(v.counts[index] * AllocatedBytes(k.alloc.stack_trace));
      int64_t count;
      if (allocated_size != 0) {
        count = (bytes + allocated_size - 1) / allocated_size;
      } else {
        // Zero-byte allocations without any bytes allocated are serviced by
        // GuardedPageAllocator, and we know there was only one sample.
        count = 1;
      }
      int64_t sum = count * allocated_size;

      // The variance should be >= 0, but it's not impossible that it drops
      // below 0 for numerical reasons. We don't want to crash in this case,
      // so we ensure to return 0 if this happens.
      double stddev_life_time_ns =
          sqrt(std::max(0.0, v.variance_life_times_ns[index] /
                                 static_cast<double>((v.counts[index]))));

      const auto bucketize = internal::LifetimeNsToBucketedDuration;
      Profile::Sample sample;
      sample.sum = sum;
      sample.requested_size = k.alloc.stack_trace.requested_size;
      sample.requested_alignment = k.alloc.stack_trace.requested_alignment;
      sample.allocated_size = allocated_size;
      sample.profile_id = pair_id++;
      sample.requested_size_returning =
          k.alloc.stack_trace.requested_size_returning;
      sample.access_hint =
          static_cast<hot_cold_t>(k.alloc.stack_trace.access_hint);
      sample.access_allocated = k.alloc.stack_trace.cold_allocated
                                    ? Profile::Sample::Access::Cold
                                    : Profile::Sample::Access::Hot;
      sample.token_id = k.alloc.stack_trace.token_id;
      sample.guarded_status = k.alloc.stack_trace.guarded_status;
      sample.type = k.alloc.stack_trace.allocation_type;
      // Set the is_censored flag so that when we create a proto
      // sample later we can treat the *_lifetime accordingly.
      sample.is_censored = (k.dealloc.stack_trace.depth == 0);
      sample.avg_lifetime = bucketize(v.mean_life_times_ns[index]);
      sample.stddev_lifetime = bucketize(stddev_life_time_ns);
      sample.min_lifetime = bucketize(v.min_life_times_ns[index]);
      sample.max_lifetime = bucketize(v.max_life_times_ns[index]);
      // Only set the cpu and thread matched flags if the sample is not
      // censored.
      if (!sample.is_censored) {
        sample.allocator_deallocator_physical_cpu_matched =
            matching_case.first.physical_cpu_matched;
        sample.allocator_deallocator_virtual_cpu_matched =
            matching_case.first.virtual_cpu_matched;
        sample.allocator_deallocator_l3_matched =
            matching_case.first.l3_matched;
        sample.allocator_deallocator_numa_matched =
            matching_case.first.numa_matched;
        sample.allocator_deallocator_thread_matched =
            matching_case.first.thread_matched;
      }

      // first for allocation
      sample.count = count;
      sample.depth = k.alloc.stack_trace.depth;
      std::copy(k.alloc.stack_trace.stack,
                k.alloc.stack_trace.stack + k.alloc.stack_trace.depth,
                sample.stack);
      func(sample);

      // If this is a right-censored allocation (i.e. we did not observe the
      // deallocation) then do not emit a deallocation sample pair.
      if (sample.is_censored) {
        continue;
      }

      // second for deallocation
      static_assert(
          std::is_signed_v<decltype(tcmalloc::Profile::Sample::count)>,
          "Deallocation samples are tagged with negative count values.");
      sample.count = -1 * count;
      sample.depth = k.dealloc.stack_trace.depth;
      std::copy(k.dealloc.stack_trace.stack,
                k.dealloc.stack_trace.stack + k.dealloc.stack_trace.depth,
                sample.stack);
      func(sample);
    }
  }
}

DeallocationSample::DeallocationSample(DeallocationProfilerList* list) {
  profiler_ = std::make_unique<DeallocationProfiler>(list);
}

DeallocationSample::~DeallocationSample() = default;

tcmalloc::Profile DeallocationSample::Stop() && {
  if (profiler_ != nullptr) {
    tcmalloc::Profile profile = profiler_->Stop();
    profiler_.reset();
    return profile;
  }
  return tcmalloc::Profile();
}

namespace internal {

// Lifetimes below 1ns are truncated to 1ns.  Lifetimes between 1ns and 1ms
// are rounded to the next smaller power of 10.  Lifetimes above 1ms are rounded
// down to the nearest millisecond.
absl::Duration LifetimeNsToBucketedDuration(double lifetime_ns) {
  if (lifetime_ns < 1000000.0) {
    if (lifetime_ns <= 1) {
      // Avoid negatives.  We can't allocate in a negative amount of time or
      // even as quickly as a nanosecond (microbenchmarks of
      // allocation/deallocation in a tight loop are several nanoseconds), so
      // results this small indicate probable clock skew or other confounding
      // factors in the data.
      return absl::Nanoseconds(1);
    }

    for (uint64_t cutoff_ns = 10; cutoff_ns <= 1000000; cutoff_ns *= 10) {
      if (lifetime_ns < cutoff_ns) {
        return absl::Nanoseconds(cutoff_ns / 10);
      }
    }
  }

  // Round down to nearest millisecond.
  return absl::Nanoseconds(static_cast<uint64_t>(lifetime_ns / 1000000.0) *
                           1000000L);
}

}  // namespace internal
}  // namespace deallocationz
}  // namespace tcmalloc
GOOGLE_MALLOC_SECTION_END
