// Copyright 2018 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/sampling_heap_profiler/sampling_heap_profiler.h"

#include <stdlib.h>

#include <cinttypes>

#include "base/allocator/dispatcher/dispatcher.h"
#include "base/allocator/dispatcher/notification_data.h"
#include "base/allocator/dispatcher/subsystem.h"
#include "base/containers/heap_array.h"
#include "base/debug/alias.h"
#include "base/features.h"
#include "base/memory/raw_ptr.h"
#include "base/rand_util.h"
#include "base/sampling_heap_profiler/poisson_allocation_sampler.h"
#include "base/synchronization/waitable_event.h"
#include "base/test/scoped_feature_list.h"
#include "base/threading/simple_thread.h"
#include "build/build_config.h"
#include "partition_alloc/shim/allocator_shim.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace base {

using ScopedSuppressRandomnessForTesting =
    PoissonAllocationSampler::ScopedSuppressRandomnessForTesting;
using base::allocator::dispatcher::AllocationNotificationData;
using base::allocator::dispatcher::AllocationSubsystem;
using base::allocator::dispatcher::FreeNotificationData;

class SamplingHeapProfilerTest : public ::testing::Test {
 public:
  void SetUp() override {
#if BUILDFLAG(IS_APPLE)
    allocator_shim::InitializeAllocatorShim();
#endif
    SamplingHeapProfiler::Init();

    // Ensure the PoissonAllocationSampler starts in the default state.
    ASSERT_FALSE(PoissonAllocationSampler::AreHookedSamplesMuted());
    ASSERT_FALSE(PoissonAllocationSampler::ScopedMuteThreadSamples::IsMuted());
    ASSERT_FALSE(ScopedSuppressRandomnessForTesting::IsSuppressed());

    allocator::dispatcher::Dispatcher::GetInstance().InitializeForTesting(
        PoissonAllocationSampler::Get());
  }

  void TearDown() override {
    allocator::dispatcher::Dispatcher::GetInstance().ResetForTesting();
  }

  size_t GetNextSample(size_t mean_interval) {
    return PoissonAllocationSampler::GetNextSampleInterval(mean_interval);
  }

  static int GetRunningSessionsCount() {
    SamplingHeapProfiler* p = SamplingHeapProfiler::Get();
    AutoLock lock(p->start_stop_mutex_);
    return p->sessions_.size();
  }

  static void RunStartStopLoop(SamplingHeapProfiler* profiler) {
    for (int i = 0; i < 100000; ++i) {
      auto session = profiler->Start(
          base::ByteSize(1024), SamplingHeapProfiler::Priority::kBackground);
      ASSERT_TRUE(session.has_value());
      EXPECT_LE(1, GetRunningSessionsCount());
      profiler->Stop(*session);
    }
  }
};

class SamplesCollector : public PoissonAllocationSampler::SamplesObserver {
 public:
  explicit SamplesCollector(size_t watch_size) : watch_size_(watch_size) {}

  void SampleAdded(void* address,
                   size_t size,
                   size_t,
                   AllocationSubsystem,
                   const char*) override {
    if (sample_added || size != watch_size_) {
      return;
    }
    sample_address_ = address;
    sample_added = true;
  }

  void SampleRemoved(void* address) override {
    if (address == sample_address_) {
      sample_removed = true;
      sample_address_ = nullptr;
    }
  }

  bool sample_added = false;
  bool sample_removed = false;

 private:
  size_t watch_size_;
  raw_ptr<void> sample_address_ = nullptr;
};

TEST_F(SamplingHeapProfilerTest, SampleObserver) {
  ScopedSuppressRandomnessForTesting suppress;
  SamplesCollector collector(10000);
  auto* sampler = PoissonAllocationSampler::Get();
  sampler->SetSamplingInterval(1024);
  sampler->AddSamplesObserver(&collector);
  void* volatile p = malloc(10000);
  free(p);
  sampler->RemoveSamplesObserver(&collector);
  EXPECT_TRUE(collector.sample_added);
  EXPECT_TRUE(collector.sample_removed);
}

TEST_F(SamplingHeapProfilerTest, SampleObserverMuted) {
  ScopedSuppressRandomnessForTesting suppress;
  SamplesCollector collector(10000);
  auto* sampler = PoissonAllocationSampler::Get();
  sampler->SetSamplingInterval(1024);
  sampler->AddSamplesObserver(&collector);
  {
    PoissonAllocationSampler::ScopedMuteThreadSamples muted_scope;
    void* volatile p = malloc(10000);
    free(p);
  }
  sampler->RemoveSamplesObserver(&collector);
  EXPECT_FALSE(collector.sample_added);
  EXPECT_FALSE(collector.sample_removed);
}

TEST_F(SamplingHeapProfilerTest, IntervalRandomizationSanity) {
  ASSERT_FALSE(ScopedSuppressRandomnessForTesting::IsSuppressed());
  constexpr int iterations = 50;
  constexpr size_t target = 10000000;
  int sum = 0;
  for (int i = 0; i < iterations; ++i) {
    int samples = 0;
    for (size_t value = 0; value < target; value += GetNextSample(10000)) {
      ++samples;
    }
    // There are should be ~ target/10000 = 1000 samples.
    sum += samples;
  }
  int mean_samples = sum / iterations;
  EXPECT_NEAR(1000, mean_samples, 100);  // 10% tolerance.
}

#if BUILDFLAG(IS_IOS)
// iOS devices generally have ~4GB of RAM with no swap and therefore need a
// lower allocation limit here.
const int kNumberOfAllocations = 1000;
#else
const int kNumberOfAllocations = 10000;
#endif

NOINLINE void Allocate1() {
  void* p = malloc(400);
  base::debug::Alias(&p);
}

NOINLINE void Allocate2() {
  void* p = malloc(700);
  base::debug::Alias(&p);
}

NOINLINE void Allocate3() {
  void* p = malloc(20480);
  base::debug::Alias(&p);
}

class MyThread1 : public SimpleThread {
 public:
  MyThread1() : SimpleThread("MyThread1") {}
  void Run() override {
    for (int i = 0; i < kNumberOfAllocations; ++i) {
      Allocate1();
    }
  }
};

class MyThread2 : public SimpleThread {
 public:
  MyThread2() : SimpleThread("MyThread2") {}
  void Run() override {
    for (int i = 0; i < kNumberOfAllocations; ++i) {
      Allocate2();
    }
  }
};

void CheckAllocationPattern(void (*allocate_callback)()) {
  ASSERT_FALSE(ScopedSuppressRandomnessForTesting::IsSuppressed());
  auto* profiler = SamplingHeapProfiler::Get();
  base::TimeTicks t0 = base::TimeTicks::Now();
  std::map<size_t, size_t> sums;
  const int iterations = 40;
  for (int i = 0; i < iterations; ++i) {
    auto session = profiler->Start(base::ByteSize(10240),
                                   SamplingHeapProfiler::Priority::kBackground);
    ASSERT_TRUE(session.has_value());
    allocate_callback();
    std::vector<SamplingHeapProfiler::Sample> samples =
        profiler->GetSamples(session);
    profiler->Stop(*session);
    std::map<size_t, size_t> buckets;
    for (auto& sample : samples) {
      buckets[sample.size] += sample.total;
    }
    for (auto& it : buckets) {
      if (it.first != 400 && it.first != 700 && it.first != 20480) {
        continue;
      }
      sums[it.first] += it.second;
      printf("%zu,", it.second);
    }
    printf("\n");
  }

  printf("Time taken %" PRIu64 "ms\n",
         (base::TimeTicks::Now() - t0).InMilliseconds());

  for (auto sum : sums) {
    intptr_t expected = sum.first * kNumberOfAllocations;
    intptr_t actual = sum.second / iterations;
    printf("%zu:\tmean: %zu\trelative error: %.2f%%\n", sum.first, actual,
           100. * (actual - expected) / expected);
  }
}

// Manual tests to check precision of the sampling profiler.
// Yes, they do leak lots of memory.

TEST_F(SamplingHeapProfilerTest, DISABLED_ParallelLargeSmallStats) {
  CheckAllocationPattern([] {
    MyThread1 t1;
    MyThread1 t2;
    t1.Start();
    t2.Start();
    for (int i = 0; i < kNumberOfAllocations; ++i) {
      Allocate3();
    }
    t1.Join();
    t2.Join();
  });
}

TEST_F(SamplingHeapProfilerTest, DISABLED_SequentialLargeSmallStats) {
  CheckAllocationPattern([] {
    for (int i = 0; i < kNumberOfAllocations; ++i) {
      Allocate1();
      Allocate2();
      Allocate3();
    }
  });
}

// Platform TLS: alloc+free[ns]: 22.184  alloc[ns]: 8.910  free[ns]: 13.274
// thread_local: alloc+free[ns]: 18.353  alloc[ns]: 5.021  free[ns]: 13.331
// TODO(crbug.com/486038197) DISABLED_ because this is running on some builders
// despite being marked MANUAL_.
TEST_F(SamplingHeapProfilerTest, DISABLED_MANUAL_SamplerMicroBenchmark) {
  // With the sampling interval of 100KB it happens to record ~ every 450th
  // allocation in the browser process. We model this pattern here.
  constexpr size_t sampling_interval = 100000;
  constexpr size_t allocation_size = sampling_interval / 450;
  SamplesCollector collector(0);
  auto* sampler = PoissonAllocationSampler::Get();
  sampler->SetSamplingInterval(sampling_interval);
  sampler->AddSamplesObserver(&collector);
  int kNumAllocations = 50000000;

  base::TimeTicks t0 = base::TimeTicks::Now();
  for (int i = 1; i <= kNumAllocations; ++i) {
    sampler->OnAllocation(AllocationNotificationData(
        reinterpret_cast<void*>(static_cast<intptr_t>(i)), allocation_size,
        nullptr, AllocationSubsystem::kAllocatorShim));
  }
  base::TimeTicks t1 = base::TimeTicks::Now();
  for (int i = 1; i <= kNumAllocations; ++i) {
    sampler->OnFree(
        FreeNotificationData(reinterpret_cast<void*>(static_cast<intptr_t>(i)),
                             AllocationSubsystem::kAllocatorShim));
  }
  base::TimeTicks t2 = base::TimeTicks::Now();

  printf(
      "alloc+free[ns]: %.3f   alloc[ns]: %.3f   free[ns]: %.3f   "
      "alloc+free[mln/s]: %.1f   total[ms]: %.1f\n",
      (t2 - t0).InNanoseconds() * 1. / kNumAllocations,
      (t1 - t0).InNanoseconds() * 1. / kNumAllocations,
      (t2 - t1).InNanoseconds() * 1. / kNumAllocations,
      kNumAllocations / (t2 - t0).InMicrosecondsF(),
      (t2 - t0).InMillisecondsF());

  sampler->RemoveSamplesObserver(&collector);
}

class StartStopThread : public SimpleThread {
 public:
  explicit StartStopThread(WaitableEvent* event)
      : SimpleThread("MyThread2"), event_(event) {}
  void Run() override {
    auto* profiler = SamplingHeapProfiler::Get();
    event_->Signal();
    SamplingHeapProfilerTest::RunStartStopLoop(profiler);
  }

 private:
  raw_ptr<WaitableEvent> event_;
};

TEST_F(SamplingHeapProfilerTest, StartStop) {
  auto* profiler = SamplingHeapProfiler::Get();
  EXPECT_EQ(0, GetRunningSessionsCount());

  // Start with 1024 (Background)
  auto session1 = profiler->Start(base::ByteSize(1024),
                                  SamplingHeapProfiler::Priority::kBackground);
  ASSERT_TRUE(session1.has_value());
  EXPECT_EQ(1, GetRunningSessionsCount());
  EXPECT_EQ(1024u, base::PoissonAllocationSampler::Get()->SamplingInterval());

  // Start with 512 (Interactive) -> should override background
  auto session2 = profiler->Start(base::ByteSize(512),
                                  SamplingHeapProfiler::Priority::kInteractive);
  ASSERT_TRUE(session2.has_value());
  EXPECT_EQ(2, GetRunningSessionsCount());
  EXPECT_EQ(512u, base::PoissonAllocationSampler::Get()->SamplingInterval());

  // Start with 2048 (Interactive) -> 512 is still min interactive
  auto session3 = profiler->Start(base::ByteSize(2048),
                                  SamplingHeapProfiler::Priority::kInteractive);
  ASSERT_TRUE(session3.has_value());
  EXPECT_EQ(3, GetRunningSessionsCount());
  EXPECT_EQ(512u, base::PoissonAllocationSampler::Get()->SamplingInterval());

  // Stop session2 (512 Interactive) -> min interactive becomes 2048
  profiler->Stop(*session2);
  EXPECT_EQ(2, GetRunningSessionsCount());
  EXPECT_EQ(2048u, base::PoissonAllocationSampler::Get()->SamplingInterval());

  // Stop session3 (2048 Interactive) -> no interactive sessions, fallback to
  // background (1024)
  profiler->Stop(*session3);
  EXPECT_EQ(1, GetRunningSessionsCount());
  EXPECT_EQ(1024u, base::PoissonAllocationSampler::Get()->SamplingInterval());

  // Stop session1 (1024 Background) -> no sessions, should reset to default
  profiler->Stop(*session1);
  EXPECT_EQ(0, GetRunningSessionsCount());
  EXPECT_EQ(PoissonAllocationSampler::kDefaultSamplingIntervalBytes,
            base::PoissonAllocationSampler::Get()->SamplingInterval());
}

// TODO(crbug.com/40711998): Test is crashing on Mac.
#if BUILDFLAG(IS_MAC)
#define MAYBE_ConcurrentStartStop DISABLED_ConcurrentStartStop
#else
#define MAYBE_ConcurrentStartStop ConcurrentStartStop
#endif
TEST_F(SamplingHeapProfilerTest, MAYBE_ConcurrentStartStop) {
  auto* profiler = SamplingHeapProfiler::Get();
  WaitableEvent event;
  StartStopThread thread(&event);
  thread.Start();
  event.Wait();
  RunStartStopLoop(profiler);
  thread.Join();
  EXPECT_EQ(0, GetRunningSessionsCount());
}

TEST_F(SamplingHeapProfilerTest, HookedAllocatorMuted) {
  ScopedSuppressRandomnessForTesting suppress;
  EXPECT_FALSE(PoissonAllocationSampler::AreHookedSamplesMuted());

  auto* sampler = PoissonAllocationSampler::Get();
  sampler->SetSamplingInterval(1024);

  {
    PoissonAllocationSampler::ScopedMuteHookedSamplesForTesting mute_hooks;
    EXPECT_TRUE(PoissonAllocationSampler::AreHookedSamplesMuted());

    SamplesCollector collector(10000);

    // A ScopedMuteHookedSamplesForTesting exists so hooked allocations should
    // be ignored.
    sampler->AddSamplesObserver(&collector);
    void* volatile p = malloc(10000);
    free(p);
    sampler->RemoveSamplesObserver(&collector);
    EXPECT_FALSE(collector.sample_added);
    EXPECT_FALSE(collector.sample_removed);

    // Manual allocations should be captured.
    sampler->AddSamplesObserver(&collector);
    void* const kAddress = reinterpret_cast<void*>(0x1234);
    sampler->OnAllocation(AllocationNotificationData(
        kAddress, 10000, nullptr, AllocationSubsystem::kManualForTesting));
    sampler->OnFree(
        FreeNotificationData(kAddress, AllocationSubsystem::kManualForTesting));
    sampler->RemoveSamplesObserver(&collector);
    EXPECT_TRUE(collector.sample_added);
    EXPECT_TRUE(collector.sample_removed);
  }

  EXPECT_FALSE(PoissonAllocationSampler::AreHookedSamplesMuted());

  // Hooked allocations should be captured again.
  SamplesCollector collector(10000);
  sampler->AddSamplesObserver(&collector);
  void* volatile p = malloc(10000);
  free(p);
  sampler->RemoveSamplesObserver(&collector);
  EXPECT_TRUE(collector.sample_added);
  EXPECT_TRUE(collector.sample_removed);
}

// Counting resident bytes is not implemented on Fuchsia.
#if BUILDFLAG(IS_FUCHSIA)
#define MAYBE_GetSamplesWithResidency DISABLED_GetSamplesWithResidency
#else
#define MAYBE_GetSamplesWithResidency GetSamplesWithResidency
#endif
TEST_F(SamplingHeapProfilerTest, MAYBE_GetSamplesWithResidency) {
  base::test::ScopedFeatureList scoped_feature_list;
  scoped_feature_list.InitAndEnableFeature(
      features::kHeapProfilerIncludeResidency);

  ScopedSuppressRandomnessForTesting suppress;
  auto* profiler = SamplingHeapProfiler::Get();
  auto session = profiler->Start(base::ByteSize(1024),
                                 SamplingHeapProfiler::Priority::kBackground);
  ASSERT_TRUE(session.has_value());

  auto buffer = base::HeapArray<uint8_t>::WithSize(10000);
  std::ranges::fill(buffer, 0u);
  base::debug::Alias(buffer.data());

  std::vector<SamplingHeapProfiler::Sample> samples =
      profiler->GetSamples(session);
  profiler->Stop(*session);
  buffer = base::HeapArray<uint8_t>();

  ASSERT_FALSE(samples.empty());
  bool found = false;
  for (const auto& sample : samples) {
    if (sample.size == 10000) {
      found = true;
      ASSERT_TRUE(sample.resident_total.has_value());
      EXPECT_GT(*sample.resident_total, 0u);
      break;
    }
  }
  EXPECT_TRUE(found);
}

}  // namespace base