// 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.

#ifndef TCMALLOC_ALLOCATION_SAMPLING_H_
#define TCMALLOC_ALLOCATION_SAMPLING_H_

#include <cstddef>
#include <memory>
#include <optional>

#include "absl/base/attributes.h"
#include "absl/debugging/stacktrace.h"
#include "tcmalloc/error_reporting.h"
#include "tcmalloc/internal/config.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/internal/percpu.h"
#include "tcmalloc/malloc_extension.h"
#include "tcmalloc/malloc_hook.h"
#include "tcmalloc/malloc_hook_invoke.h"
#include "tcmalloc/pagemap.h"
#include "tcmalloc/parameters.h"
#include "tcmalloc/sampler.h"
#include "tcmalloc/span.h"
#include "tcmalloc/tcmalloc_policy.h"

GOOGLE_MALLOC_SECTION_BEGIN
namespace tcmalloc::tcmalloc_internal {

class Static;

std::unique_ptr<const ProfileBase> DumpHeapProfile(Static& state);
#if !TCMALLOC_INTERNAL_PERCPU_USE_RSEQ
// For RSEQ enabled builds, we declare the sampler in percpu.h so that we can
// reference its address in percpu_tcmalloc.h without creating a circular
// dependency with the Sampler definition. Otherwise, we declare it here to
// avoid relying on the percpu code.
extern "C" ABSL_CONST_INIT thread_local Sampler tcmalloc_sampler
    ABSL_ATTRIBUTE_INITIAL_EXEC;

#endif

ABSL_CONST_INIT ABSL_ATTRIBUTE_WEAK thread_local Sampler tcmalloc_sampler
    ABSL_ATTRIBUTE_INITIAL_EXEC;

inline Sampler& GetThreadSampler() {
  static_assert(sizeof(Sampler) == TCMALLOC_SAMPLER_SIZE,
                "update TCMALLOC_SAMPLER_SIZE");
  static_assert(alignof(Sampler) == TCMALLOC_SAMPLER_ALIGN,
                "update TCMALLOC_SAMPLER_ALIGN");
  static_assert(Sampler::HotDataOffset() == TCMALLOC_SAMPLER_HOT_OFFSET,
                "update TCMALLOC_SAMPLER_HOT_OFFSET");
#if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ
  return *static_cast<Sampler*>(subtle::percpu::GetThreadSamplerAddress());
#else
  return tcmalloc_sampler;
#endif
}

// Performs sampling for already occurred allocation of object.
//
// For small object sizes, we allocate a new object in a sampled span.
//
// For large objects (i.e. allocated with do_malloc_pages) they are
// also fully reused and their span is marked as sampled.
//
// Note that do_free_with_size assumes sampled objects have
// page-aligned addresses. Please change both functions if need to
// invalidate the assumption.
//
// Note that size_class might not match requested_size in case of
// memalign. I.e. when larger than requested allocation is done to
// satisfy alignment constraint.
//
// In case of out-of-memory condition when allocating span or
// stacktrace struct, this function simply cheats and returns original
// object. As if no sampling was requested.
template <typename Policy>
ABSL_ATTRIBUTE_NOINLINE sized_ptr_t SampleifyAllocation(
    Static& state, Policy policy, size_t requested_size, size_t weight,
    size_t size_class, Span* absl_nullable span) {
  TC_CHECK_EQ(size_class != 0, span == nullptr);

  StackTrace stack_trace;
  stack_trace.requested_size = requested_size;
  // Grab the stack trace outside the heap lock.
  stack_trace.depth = absl::GetStackTrace(stack_trace.stack, kMaxStackDepth, 0);

  if (policy.has_explicit_alignment()) {
    stack_trace.requested_alignment = policy.align();
  } else {
    stack_trace.requested_alignment = std::nullopt;
  }

  stack_trace.requested_size_returning = policy.size_returning();
  stack_trace.access_hint = static_cast<uint8_t>(policy.access());
  stack_trace.weight = weight;
  stack_trace.token_id = policy.token_id();

  // How many allocations does this sample represent, given the sampling
  // frequency (weight) and its size.
  const double allocation_estimate =
      static_cast<double>(weight) / (requested_size + 1);

  GuardedAllocWithStatus alloc_with_status{
      nullptr, Profile::Sample::GuardedStatus::NotAttempted};

  const MemoryTag tag =
      Parameters::heap_partitioning_mode() == HeapPartitioningMode::kFull &&
              policy.partition() == 1
          ? MemoryTag::kSampledP1
          : MemoryTag::kSampled;
  size_t capacity = 0;
  if (size_class != 0) {
    state.per_size_class_counts()[size_class].Add(allocation_estimate);

    stack_trace.allocated_size = state.sizemap().class_to_size(size_class);
    stack_trace.cold_allocated = IsExpandedSizeClass(size_class);

    Length num_pages = BytesToLengthCeil(stack_trace.allocated_size);
    std::align_val_t sample_alignment = policy.align();
    alloc_with_status = state.guardedpage_allocator().TrySample(
        requested_size, sample_alignment, num_pages, stack_trace);
    if (alloc_with_status.status == Profile::Sample::GuardedStatus::Guarded) {
      TC_ASSERT(!IsNormalMemory(alloc_with_status.alloc));
      const PageId p = PageIdContaining(alloc_with_status.alloc);
#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING
      PageHeapSpinLockHolder l;
#endif  // TCMALLOC_INTERNAL_LEGACY_LOCKING
      span = Span::New(Range(p, num_pages));
      state.pagemap().Set(p, span);
      // If we report capacity back from a size returning allocation, we can not
      // report the stack_trace.allocated_size, as we guard the size to
      // 'requested_size', and we maintain the invariant that GetAllocatedSize()
      // must match the returned size from size returning allocations. So in
      // that case, we report the requested size for both capacity and
      // GetAllocatedSize().
      if (policy.size_returning()) {
        stack_trace.allocated_size = requested_size;
      }
      capacity = requested_size;
    } else if ((span = state.page_allocator().New(
                    num_pages, {1, AccessDensityPrediction::kSparse}, tag))) {
      capacity = stack_trace.allocated_size;
    } else {
      return {};
    }
  } else {
    // Set stack_trace.allocated_size to the exact size for a page allocation.
    // NOTE: if we introduce gwp-asan sampling / guarded allocations
    // for page allocations, then we need to revisit do_malloc_pages as
    // the current assumption is that only class sized allocs are sampled
    // for gwp-asan.
    stack_trace.allocated_size = span->bytes_in_span();
    stack_trace.cold_allocated =
        GetMemoryTag(span->start_address()) == MemoryTag::kCold;
    capacity = stack_trace.allocated_size;
  }

  // A span must be provided or created by this point.
  TC_ASSERT_NE(span, nullptr);

  // TODO(b/414876446): Add entropy to the handles generated.
  stack_trace.sampled_alloc_handle =
      AllocHandle(state.sampled_alloc_handle_generator.fetch_add(
                      1, std::memory_order_relaxed) +
                  1);
  stack_trace.span_start_address = span->start_address();
  stack_trace.allocation_time = absl::Now();
  stack_trace.guarded_status = alloc_with_status.status;
  stack_trace.allocation_type = policy.allocation_type();

  // Adjust our estimate of internal fragmentation.
  TC_ASSERT_LE(requested_size, stack_trace.allocated_size);
  if (requested_size < stack_trace.allocated_size) {
    state.sampled_internal_fragmentation_.Add(
        allocation_estimate * (stack_trace.allocated_size - requested_size));
  }

  MallocHook::SampledAlloc sampled_alloc = {
      .handle = stack_trace.sampled_alloc_handle,
      .requested_size = stack_trace.requested_size,
      .requested_alignment = stack_trace.requested_alignment,
      .allocated_size = stack_trace.allocated_size,
      .weight = allocation_estimate,
      .stack = absl::MakeSpan(stack_trace.stack, stack_trace.depth),
      .allocation_time = stack_trace.allocation_time,
      .ptr = (alloc_with_status.alloc != nullptr)
                 ? alloc_with_status.alloc
                 : stack_trace.span_start_address,
      .access_hint = stack_trace.access_hint,
      .access_allocated = stack_trace.cold_allocated ? MallocHook::Access::Cold
                                                     : MallocHook::Access::Hot,
  };
  MallocHook::InvokeSampledNewHook(sampled_alloc);

  state.allocation_samples.ReportMalloc(stack_trace);

  state.deallocation_samples.ReportMalloc(stack_trace);

  // The SampledAllocation object is visible to readers after this. Readers only
  // care about its various metadata (e.g. stack trace, weight) to generate the
  // heap profile, and won't need any information from Span::Sample() next.
  SampledAllocation* sampled_allocation =
      state.sampled_allocation_recorder().Register(std::move(stack_trace));
  // No pageheap_lock required. The span is freshly allocated and no one else
  // can access it. It is visible after we return from this allocation path.
  span->Sample(sampled_allocation);

  state.peak_heap_tracker().MaybeSaveSample();

  TC_ASSERT_EQ(state.pagemap().sizeclass(span->first_page()), 0);
  return {(alloc_with_status.alloc != nullptr) ? alloc_with_status.alloc
                                               : span->start_address(),
          capacity};
}

void MaybeUnsampleAllocation(Static& state, void* absl_nonnull ptr,
                             std::optional<size_t> size, Span& span,
                             AllocationType type);

template <typename Policy>
static sized_ptr_t SampleLargeAllocation(Static& state, Policy policy,
                                         size_t requested_size, size_t weight,
                                         Span* span) {
  return SampleifyAllocation(state, policy, requested_size, weight, 0, span);
}

template <typename Policy>
static sized_ptr_t SampleSmallAllocation(Static& state, Policy policy,
                                         size_t requested_size, size_t weight,
                                         size_t size_class) {
  return SampleifyAllocation(state, policy, requested_size, weight, size_class,
                             nullptr);
}

// Rewrite type so that the allocation type falls into one of the categories we
// use for deallocations (new or malloc, not aligned new).
static inline AllocationType SimplifyType(AllocationType type) {
  switch (type) {
    case AllocationType::New:
    case AllocationType::Malloc:
      return type;
    case AllocationType::AlignedMalloc:
      return AllocationType::Malloc;
  }

  ABSL_UNREACHABLE();
}

template <typename Policy>
void MaybeUnsampleAllocation(Static& state, Policy policy,
                             void* absl_nonnull ptr, std::optional<size_t> size,
                             Span& span) {
  // No pageheap_lock required. The sampled span should be unmarked and have its
  // state cleared only once. External synchronization when freeing is required;
  // otherwise, concurrent writes here would likely report a double-free.
  SampledAllocation* sampled_allocation = span.Unsample();
  if (sampled_allocation == nullptr) {
    if (ABSL_PREDICT_TRUE(size.has_value())) {
      const size_t maximum_size = span.bytes_in_span();
      const size_t minimum_size = maximum_size - (kPageSize - 1u);

      // *size should fall in [minimum_size, maximum_size], but a size of 0
      // bytes is rounded up to 1 page if overaligned, we must allow 0 bytes if
      // the span is 1 page.
      if (ABSL_PREDICT_FALSE(*size < minimum_size || *size > maximum_size) &&
          ABSL_PREDICT_TRUE(
              !(*size == 0 && maximum_size == kPageSize &&
                static_cast<size_t>(policy.align()) > kPageSize))) {
        // While we don't have precise allocation-time information because
        // this span was not sampled, the deallocated object's purported size
        // exceeds the span it is on.  This is impossible and indicates
        // corruption.
        ReportMismatchedDelete(state, ptr, *size, minimum_size, maximum_size);
      }
    }

    return;
  }

  TC_ASSERT_EQ(state.pagemap().sizeclass(PageIdContainingTagged(ptr)), 0);

  const size_t weight = sampled_allocation->sampled_stack.weight;
  const size_t requested_size =
      sampled_allocation->sampled_stack.requested_size;
  const size_t allocated_size =
      sampled_allocation->sampled_stack.allocated_size;
  if (size.has_value()) {
    if (sampled_allocation->sampled_stack.requested_size_returning) {
      if (ABSL_PREDICT_FALSE(
              !(requested_size <= *size && *size <= allocated_size))) {
        ReportMismatchedDelete(state, ptr, *sampled_allocation, *size,
                               requested_size, allocated_size);
      }
    } else if (ABSL_PREDICT_FALSE(size != requested_size)) {
      ReportMismatchedDelete(state, ptr, *sampled_allocation, *size,
                             requested_size, std::nullopt);
    }
  }

  if (auto dealloc_type = SimplifyType(policy.allocation_type()),
      alloc_type =
          SimplifyType(sampled_allocation->sampled_stack.allocation_type);
      ABSL_PREDICT_FALSE(dealloc_type != alloc_type)) {
    ReportMismatchedFree(
        state, ptr, sampled_allocation->sampled_stack.allocation_type,
        policy.allocation_type(),
        absl::MakeSpan(sampled_allocation->sampled_stack.stack,
                       sampled_allocation->sampled_stack.depth));
  }

  // Check pointer for misalignment.
  //
  // TODO(ckennelly): Eliminate redundant guarded check with
  // InvokeHooksAndFreePages.
  if (ABSL_PREDICT_FALSE(ptr != span.start_address()) &&
      sampled_allocation->sampled_stack.guarded_status !=
          Profile::Sample::GuardedStatus::Guarded) {
    ReportCorruptedFree(
        tc_globals, static_cast<std::align_val_t>(kPageSize), ptr,
        absl::MakeSpan(sampled_allocation->sampled_stack.stack,
                       sampled_allocation->sampled_stack.depth));
  }

  if ((size.has_value() || policy.allocation_type() == AllocationType::New)) {
    const bool type_mismatch =
        policy.allocation_type() !=
        sampled_allocation->sampled_stack.allocation_type;
    const std::optional<std::align_val_t> deallocated_alignment =
        policy.has_explicit_alignment()
            ? std::make_optional<std::align_val_t>(policy.align())
            : std::nullopt;
    const bool alignment_mismatch =
        deallocated_alignment !=
        sampled_allocation->sampled_stack.requested_alignment;
    if (ABSL_PREDICT_FALSE(type_mismatch || alignment_mismatch)) {
      ReportMismatchedFree(
          state, ptr, sampled_allocation->sampled_stack.requested_alignment,
          deallocated_alignment,
          absl::MakeSpan(sampled_allocation->sampled_stack.stack,
                         sampled_allocation->sampled_stack.depth));
    }
  }

  // How many allocations does this sample represent, given the sampling
  // frequency (weight) and its size.
  const double allocation_estimate =
      static_cast<double>(weight) / (requested_size + 1);
  AllocHandle sampled_alloc_handle =
      sampled_allocation->sampled_stack.sampled_alloc_handle;
  MallocHook::SampledAlloc sampled_alloc = {
      .handle = sampled_alloc_handle,
      .requested_size = requested_size,
      .requested_alignment =
          sampled_allocation->sampled_stack.requested_alignment,
      .allocated_size = allocated_size,
      .weight = allocation_estimate,
      .stack = absl::MakeSpan(sampled_allocation->sampled_stack.stack,
                              sampled_allocation->sampled_stack.depth),
      .allocation_time = sampled_allocation->sampled_stack.allocation_time,
      .ptr = ptr,
      .access_hint = sampled_allocation->sampled_stack.access_hint,
      .access_allocated = sampled_allocation->sampled_stack.cold_allocated
                              ? MallocHook::Access::Cold
                              : MallocHook::Access::Hot,
  };
  state.sampled_allocation_recorder().Unregister(sampled_allocation);

  // Adjust our estimate of internal fragmentation.
  TC_ASSERT_LE(requested_size, allocated_size);
  if (requested_size < allocated_size) {
    const size_t sampled_fragmentation =
        allocation_estimate * (allocated_size - requested_size);

    // Check against wraparound
    TC_ASSERT_GE(state.sampled_internal_fragmentation_.value(),
                 sampled_fragmentation);
    state.sampled_internal_fragmentation_.Add(-sampled_fragmentation);
  }
  MallocHook::InvokeSampledDeleteHook(sampled_alloc);

  state.deallocation_samples.ReportFree(sampled_alloc_handle);
}

}  // namespace tcmalloc::tcmalloc_internal
GOOGLE_MALLOC_SECTION_END

#endif  // TCMALLOC_ALLOCATION_SAMPLING_H_
