// 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 "components/heap_profiling/in_process/heap_profiler_controller.h"

#include <array>
#include <atomic>
#include <iomanip>
#include <memory>
#include <optional>
#include <ostream>
#include <string>
#include <string_view>
#include <tuple>
#include <utility>
#include <vector>

#include "base/allocator/dispatcher/notification_data.h"
#include "base/allocator/dispatcher/subsystem.h"
#include "base/auto_reset.h"
#include "base/barrier_closure.h"
#include "base/command_line.h"
#include "base/functional/bind.h"
#include "base/functional/callback.h"
#include "base/functional/callback_helpers.h"
#include "base/memory/scoped_refptr.h"
#include "base/metrics/field_trial_params.h"
#include "base/metrics/metrics_hashes.h"
#include "base/process/launch.h"
#include "base/process/process.h"
#include "base/sampling_heap_profiler/poisson_allocation_sampler.h"
#include "base/sampling_heap_profiler/sampling_heap_churn_profiler.h"
#include "base/sampling_heap_profiler/sampling_heap_profiler.h"
#include "base/strings/string_number_conversions.h"
#include "base/task/bind_post_task.h"
#include "base/task/sequenced_task_runner.h"
#include "base/task/single_thread_task_runner.h"
#include "base/test/bind.h"
#include "base/test/metrics/histogram_tester.h"
#include "base/test/multiprocess_test.h"
#include "base/test/scoped_command_line.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/task_environment.h"
#include "base/test/test_timeouts.h"
#include "base/time/time.h"
#include "base/types/optional_util.h"
#include "build/build_config.h"
#include "components/heap_profiling/in_process/browser_process_snapshot_controller.h"
#include "components/heap_profiling/in_process/child_process_snapshot_controller.h"
#include "components/heap_profiling/in_process/heap_profiler_parameters.h"
#include "components/heap_profiling/in_process/mojom/snapshot_controller.mojom.h"
#include "components/heap_profiling/in_process/mojom/test_connector.mojom.h"
#include "components/heap_profiling/in_process/switches.h"
#include "components/metrics/call_stacks/call_stack_profile_builder.h"
#include "components/metrics/public/mojom/call_stack_profile_collector.mojom.h"
#include "components/sampling_profiler/process_type.h"
#include "components/version_info/channel.h"
#include "mojo/core/embedder/scoped_ipc_support.h"
#include "mojo/public/cpp/base/proto_wrapper.h"
#include "mojo/public/cpp/bindings/pending_receiver.h"
#include "mojo/public/cpp/bindings/pending_remote.h"
#include "mojo/public/cpp/bindings/receiver.h"
#include "mojo/public/cpp/bindings/remote_set.h"
#include "mojo/public/cpp/bindings/unique_receiver_set.h"
#include "mojo/public/cpp/platform/platform_channel.h"
#include "mojo/public/cpp/system/invitation.h"
#include "mojo/public/cpp/system/message_pipe.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "testing/multiprocess_func_list.h"
#include "third_party/metrics_proto/call_stack_profile.pb.h"
#include "third_party/metrics_proto/execution_context.pb.h"
#include "third_party/metrics_proto/sampled_profile.pb.h"

// Test printers. These need to be in the same namespace as the types to print
// so GTest can find them.

namespace metrics {

void PrintTo(const SampledProfile& profile, std::ostream* os) {
  *os << "process:" << profile.process() << ",samples";
  if (profile.call_stack_profile().stack_sample_size() == 0) {
    *os << ":none";
  } else {
    for (const auto& sample : profile.call_stack_profile().stack_sample()) {
      *os << ":" << sample.count() << "/" << sample.weight();
    }
  }
  *os << ",metadata";
  if (profile.call_stack_profile().profile_metadata_size() == 0) {
    *os << ":none";
  } else {
    const auto& name_hashes = profile.call_stack_profile().metadata_name_hash();
    for (const auto& metadata :
         profile.call_stack_profile().profile_metadata()) {
      if (metadata.name_hash_index() < name_hashes.size()) {
        *os << ":" << std::hex << "0x"
            << name_hashes.at(metadata.name_hash_index()) << std::dec;
      } else {
        *os << ":unknown";
      }
      *os << "=" << metadata.value();
    }
  }
}

}  // namespace metrics

namespace sampling_profiler {

void PrintTo(const ProfilerProcessType& process_type, std::ostream* os) {
  switch (process_type) {
    case ProfilerProcessType::kUnknown:
      *os << "unknown";
      return;
    case ProfilerProcessType::kBrowser:
      *os << "browser";
      return;
    case ProfilerProcessType::kGpu:
      *os << "gpu";
      return;
    case ProfilerProcessType::kNetworkService:
      *os << "network";
      return;
    case ProfilerProcessType::kRenderer:
      *os << "renderer";
      return;
    case ProfilerProcessType::kUtility:
      *os << "utility";
      return;
    default:
      *os << "unsupported (" << static_cast<int>(process_type) << ")";
      return;
  }
}

}  // namespace sampling_profiler

namespace heap_profiling {

namespace {

#if BUILDFLAG(IS_IOS) || BUILDFLAG(IS_ANDROID)
#define ENABLE_MULTIPROCESS_TESTS 0
#else
#define ENABLE_MULTIPROCESS_TESTS 1
#endif

using base::allocator::dispatcher::AllocationNotificationData;
using base::allocator::dispatcher::AllocationSubsystem;
using base::allocator::dispatcher::FreeNotificationData;
using base::test::FeatureRef;
using base::test::FeatureRefAndParams;
using sampling_profiler::ProfilerProcessType;
using ::testing::_;
using ::testing::AllOf;
using ::testing::Combine;
using ::testing::Conditional;
using ::testing::ElementsAre;
using ::testing::Ge;
using ::testing::IsEmpty;
using ::testing::Lt;
using ::testing::Optional;
using ::testing::Property;
using ::testing::ResultOf;
using ::testing::UnorderedElementsAre;
using ::testing::Values;
using ::testing::ValuesIn;

using ProfileCollectorCallback =
    base::RepeatingCallback<void(base::TimeTicks, metrics::SampledProfile)>;
using ScopedMuteHookedSamplesForTesting =
    base::PoissonAllocationSampler::ScopedMuteHookedSamplesForTesting;
using ScopedSuppressRandomnessForTesting =
    base::PoissonAllocationSampler::ScopedSuppressRandomnessForTesting;

constexpr size_t kSamplingRate = 1024;
constexpr size_t kAllocationSize = 42 * kSamplingRate;

// A fake CallStackProfileCollector that deserializes profiles it receives from
// a fake child process, and passes them to the same callback that receives
// profiles from the fake browser process.
class TestCallStackProfileCollector final
    : public metrics::mojom::CallStackProfileCollector {
 public:
  explicit TestCallStackProfileCollector(
      ProfileCollectorCallback collector_callback)
      : collector_callback_(std::move(collector_callback)) {}

  TestCallStackProfileCollector(const TestCallStackProfileCollector&) = delete;
  TestCallStackProfileCollector& operator=(
      const TestCallStackProfileCollector&) = delete;

  ~TestCallStackProfileCollector() final = default;

  // metrics::mojom::CallStackProfileCollector
  void Collect(base::TimeTicks start_timestamp,
               metrics::mojom::TriggerEvent trigger_event,
               metrics::mojom::SampledProfilePtr profile) final {
    metrics::SampledProfile sampled_profile;
    ASSERT_TRUE(profile);
    ASSERT_TRUE(base::OptionalUnwrapTo(
        profile->contents.As<metrics::SampledProfile>(), sampled_profile));
    EXPECT_TRUE(
        (trigger_event ==
             metrics::mojom::TriggerEvent::kPeriodicHeapCollection &&
         sampled_profile.trigger_event() ==
             metrics::SampledProfile::PERIODIC_HEAP_COLLECTION) ||
        (trigger_event ==
             metrics::mojom::TriggerEvent::kPeriodicHeapChurnCollection &&
         sampled_profile.trigger_event() ==
             metrics::SampledProfile::PERIODIC_HEAP_CHURN_COLLECTION));
    collector_callback_.Run(start_timestamp, std::move(sampled_profile));
  }

 private:
  ProfileCollectorCallback collector_callback_;
};

// A scoped holder for callbacks to pass to
// HeapProfilerControllerTest::StartHeapProfiling(), and a BarrierClosure that
// will quit a run loop. The ScopedCallbacks object must stay in scope until the
// run loop finishes.
class ScopedCallbacks {
 public:
  // Creates a BarrierClosure that will invoke the given `quit_closure` after
  // the expected callbacks are invoked:
  //
  // If `expect_take_snapshot` is true, HeapProfilerController::TakeSnapshot()
  // should be called, so first_snapshot_callback() returns a callback that's
  // expected to be invoked. If not, first_snapshot_callback() returns a
  // callback that's expected not to run.
  //
  // If `expected_sampled_profiles` > 0, HeapProfilerController::TakeSnapshot()
  // should find a sample to pass to CallStackProfileBuilder, so
  // collector_callback() returns a callback that's expected to be invoked. It
  // will also run the given `profile_collector_callback`. If not,
  // collector_callback() returns a callback that's expected not to run, and the
  // given `profile_collector_callback` is ignored.
  //
  // If `use_other_process_callback` is true, the test will also fake a snapshot
  // in another process, so other_process_callback() will return a callback to
  // invoke for this.
  ScopedCallbacks(bool expect_take_snapshot,
                  size_t expected_sampled_profiles,
                  bool use_other_process_callback,
                  ProfileCollectorCallback profile_collector_callback,
                  base::OnceClosure quit_closure) {
    size_t num_callbacks = 0;
    if (expect_take_snapshot) {
      num_callbacks += 1;
    }
    num_callbacks += expected_sampled_profiles;
    if (use_other_process_callback) {
      // The test should invoke other_process_snapshot_callback() to simulate a
      // snapshot in another process.
      num_callbacks += 1;
    }
    barrier_closure_ =
        base::BarrierClosure(num_callbacks, std::move(quit_closure));

    // Each callback should invoke `barrier_closure_` once. If they're called
    // too often, log a test failure on the first extra call only to avoid log
    // spam. These lambdas need to take a copy of the method arguments since
    // they outlive the method` scope.
    first_snapshot_callback_ =
        base::BindLambdaForTesting([this, expect_take_snapshot] {
          if (!expect_take_snapshot) {
            FAIL() << "TakeSnapshot called unexpectedly.";
          }
          first_snapshot_count_++;
          if (first_snapshot_count_ == 1) {
            barrier_closure_.Run();
            return;
          }
          if (first_snapshot_count_ == 2) {
            FAIL() << "TakeSnapshot callback invoked too many times.";
          }
        });

    collector_callback_ = base::BindLambdaForTesting(
        [this, expected_sampled_profiles,
         callback = std::move(profile_collector_callback)](
            base::TimeTicks time_ticks, metrics::SampledProfile profile) {
          if (expected_sampled_profiles == 0) {
            FAIL() << "ProfileCollectorCallback called unexpectedly.";
          }
          collector_count_++;
          if (collector_count_ <= expected_sampled_profiles) {
            std::move(callback).Run(time_ticks, profile);
            barrier_closure_.Run();
            return;
          }
          if (collector_count_ == expected_sampled_profiles + 1) {
            FAIL() << "ProfileCollectorCallback invoked too many times.";
          }
        });

    other_process_callback_ =
        base::BindLambdaForTesting([this, use_other_process_callback] {
          if (!use_other_process_callback) {
            FAIL() << "Other process callback invoked unexpectedly.";
          }
          other_process_count_++;
          if (other_process_count_ == 1) {
            barrier_closure_.Run();
            return;
          }
          if (other_process_count_ == 2) {
            FAIL() << "Other process callback invoked too many times.";
          }
        });
  }

  ~ScopedCallbacks() = default;

  // Move-only.
  ScopedCallbacks(const ScopedCallbacks&) = delete;
  ScopedCallbacks& operator=(const ScopedCallbacks&) = delete;
  ScopedCallbacks(ScopedCallbacks&&) = default;
  ScopedCallbacks& operator=(ScopedCallbacks&&) = default;

  base::OnceClosure first_snapshot_callback() {
    return std::move(first_snapshot_callback_);
  }

  ProfileCollectorCallback collector_callback() {
    // Return by copy since this is a RepeatingCallback.
    return collector_callback_;
  }

  base::OnceClosure other_process_callback() {
    return std::move(other_process_callback_);
  }

 private:
  base::RepeatingClosure barrier_closure_;

  base::OnceClosure first_snapshot_callback_;
  ProfileCollectorCallback collector_callback_;
  base::OnceClosure other_process_callback_;

  size_t first_snapshot_count_ = 0;
  size_t collector_count_ = 0;
  size_t other_process_count_ = 0;
};

class ProfilerSetUpMixin {
 public:
  ProfilerSetUpMixin(const std::vector<FeatureRefAndParams>& enabled_features,
                     const std::vector<FeatureRef>& disabled_features) {
    // ScopedFeatureList must be initialized in the constructor, before any
    // threads are started.
    feature_list_.InitWithFeaturesAndParameters(enabled_features,
                                                disabled_features);
  }

  ~ProfilerSetUpMixin() = default;

  base::test::TaskEnvironment& task_env() { return task_environment_; }

 private:
  // Initialize `mute_hooks_` before `task_environment_` so that memory
  // allocations aren't sampled while TaskEnvironment creates a thread. The
  // sampling is crashing in the hooked FreeFunc on some test bots.
  ScopedMuteHookedSamplesForTesting mute_hooks_ =
      base::SamplingHeapProfiler::Get()->MuteHookedSamplesForTesting();
  ScopedSuppressRandomnessForTesting suppress_randomness_;

  // Create `feature_list_` before `task_environment_` and destroy it after to
  // avoid a race in destruction.
  base::test::ScopedFeatureList feature_list_;

  base::test::TaskEnvironment task_environment_{
      base::test::TaskEnvironment::TimeSource::MOCK_TIME,
      base::test::TaskEnvironment::MainThreadType::IO};
};

#if ENABLE_MULTIPROCESS_TESTS

constexpr char kTestChildTypeSwitch[] = "heap-profiler-test-child-type";
constexpr char kTestNumAllocationsSwitch[] =
    "heap-profiler-test-num-allocations";
constexpr char kTestExpectChildProfileSwitch[] =
    "heap-profiler-expect-child-profile";

// Runs the heap profiler in a multiprocess test child. This is used instead of
// HeapProfilerControllerTest::CreateHeapProfiler() in tests that create real
// child processes. (Most tests run only in the test main process and pretend
// that it's the Chrome browser process or a Chrome child process.)
class MultiprocessTestChild final : public mojom::TestConnector,
                                    public ProfilerSetUpMixin {
 public:
  MultiprocessTestChild(
      const std::vector<FeatureRefAndParams>& enabled_features,
      const std::vector<FeatureRef>& disabled_features)
      : ProfilerSetUpMixin(enabled_features, disabled_features),
        quit_closure_(task_env().QuitClosure()) {}

  ~MultiprocessTestChild() final = default;

  MultiprocessTestChild(const MultiprocessTestChild&) = delete;
  MultiprocessTestChild& operator=(const MultiprocessTestChild&) = delete;

  void RunTestInChild() {
    base::HistogramTester histogram_tester;

    // Get the process type and number of allocations to simulate.
    const base::CommandLine* command_line =
        base::CommandLine::ForCurrentProcess();
    ASSERT_TRUE(command_line);
    int process_type = 0;
    ASSERT_TRUE(base::StringToInt(
        command_line->GetSwitchValueASCII(kTestChildTypeSwitch),
        &process_type));
    const ProfilerProcessType profiler_process_type =
        static_cast<ProfilerProcessType>(process_type);
    int num_allocations = 0;
    ASSERT_TRUE(base::StringToInt(
        command_line->GetSwitchValueASCII(kTestNumAllocationsSwitch),
        &num_allocations));

    // Set up mojo support and attach to the parent's pipe.
    mojo::core::ScopedIPCSupport enable_mojo(
        base::SingleThreadTaskRunner::GetCurrentDefault(),
        mojo::core::ScopedIPCSupport::ShutdownPolicy::CLEAN);
    mojo::IncomingInvitation invitation = mojo::IncomingInvitation::Accept(
        mojo::PlatformChannel::RecoverPassedEndpointFromCommandLine(
            *command_line));

    // Handle the TestConnector::Connect() message that will connect the
    // SnapshotController and CallStackProfileCollector interfaces.
    mojo::PendingReceiver<mojom::TestConnector> pending_receiver(
        invitation.ExtractMessagePipe(0));
    mojo::Receiver<mojom::TestConnector> receiver(this,
                                                  std::move(pending_receiver));

    // Start the heap profiler and wait for TakeSnapshot() messages from the
    // parent.
    HeapProfilerController controller(version_info::Channel::STABLE,
                                      profiler_process_type);
    controller.SuppressRandomnessForTesting();
    ASSERT_TRUE(controller.IsEnabled());
    controller.StartIfEnabled();

    // Make a fixed number of allocations at different addresses to include in
    // snapshots. No need to free since the process will exit after the test.
    for (int i = 0; i < num_allocations; ++i) {
      base::PoissonAllocationSampler::Get()->OnAllocation(
          AllocationNotificationData(reinterpret_cast<void*>(0x1337 + i),
                                     kAllocationSize, nullptr,
                                     AllocationSubsystem::kManualForTesting));
    }

    // Loop until the TestConnector::Disconnect() message.
    task_env().RunUntilQuit();

    // Profiler stats should be logged whether or not a snapshot was taken, as
    // long as the child is profiled at all.
    size_t expected_histogram_count =
        command_line->HasSwitch(kTestExpectChildProfileSwitch) ? 1 : 0;
    switch (profiler_process_type) {
      case ProfilerProcessType::kGpu:
        histogram_tester.ExpectTotalCount(
            "HeapProfiling.InProcess.SamplesPerSnapshot.GPU",
            expected_histogram_count);
        break;
      case ProfilerProcessType::kNetworkService:
        histogram_tester.ExpectTotalCount(
            "HeapProfiling.InProcess.SamplesPerSnapshot.Network",
            expected_histogram_count);
        break;
      case ProfilerProcessType::kRenderer:
        histogram_tester.ExpectTotalCount(
            "HeapProfiling.InProcess.SamplesPerSnapshot.Renderer",
            expected_histogram_count);
        break;
      case ProfilerProcessType::kUtility:
        histogram_tester.ExpectTotalCount(
            "HeapProfiling.InProcess.SamplesPerSnapshot.Utility",
            expected_histogram_count);
        break;
      default:
        FAIL() << "Unexpected process type " << process_type;
    }
    histogram_tester.ExpectTotalCount(
        "HeapProfiling.InProcess.SamplesPerSnapshot", expected_histogram_count);
  }

  // mojom::TestConnector:

  void ConnectSnapshotController(
      mojo::PendingReceiver<mojom::SnapshotController> controller,
      base::OnceClosure done_callback) final {
    ChildProcessSnapshotController::CreateSelfOwnedReceiver(
        std::move(controller));
    std::move(done_callback).Run();
  }

  void ConnectProfileCollector(
      mojo::PendingRemote<metrics::mojom::CallStackProfileCollector> collector,
      base::OnceClosure done_callback) final {
    metrics::CallStackProfileBuilder::SetParentProfileCollectorForChildProcess(
        std::move(collector));
    std::move(done_callback).Run();
  }

  void Disconnect() final { std::move(quit_closure_).Run(); }

 private:
  base::OnceClosure quit_closure_;
};

// Manages a set of multiprocess test children and mojo connections to them.
// Created in test cases in the parent process.
class MultiprocessTestParent {
 public:
  MultiprocessTestParent() = default;

  MultiprocessTestParent(const MultiprocessTestParent&) = delete;
  MultiprocessTestParent& operator=(const MultiprocessTestParent&) = delete;

  ~MultiprocessTestParent() {
    // Tell all children to stop profiling and wait for them to exit.
    for (auto& connector : test_connectors_) {
      connector->Disconnect();
    }
    for (const auto& process : child_processes_) {
      int exit_code = 0;
      EXPECT_TRUE(base::WaitForMultiprocessTestChildExit(
          process, TestTimeouts::action_timeout(), &exit_code));
      EXPECT_EQ(exit_code, 0);
    }
  }

  // Waits until `num_children` are connected, then starts profiling the parent
  // process with `controller`.
  void StartHeapProfilingWhenChildrenConnected(
      size_t num_children,
      HeapProfilerController* controller) {
    // StartIfEnabled() needs to run on the current sequence no matter what
    // thread mojo calls `on_child_connected_closure_` from.
    on_child_connected_closure_ = base::BarrierClosure(
        num_children, base::BindPostTaskToCurrentDefault(
                          base::BindLambdaForTesting([this, controller] {
                            // Make sure all children connected successfully.
                            ASSERT_EQ(test_connectors_.size(),
                                      child_processes_.size());
                            EXPECT_TRUE(controller->StartIfEnabled());
                          })));
  }

  // Called from HeapProfilerController::AppendCommandLineSwitchForChildProcess
  // with `connector_id` and `receiver`, plus a `remote` added by the test.
  // `connector_id` is the id of a mojo TestConnector interface for the process.
  // In production this parameter is the child process id.
  void BindTestConnector(
      int connector_id,
      mojo::PendingReceiver<mojom::SnapshotController> receiver,
      mojo::PendingRemote<metrics::mojom::CallStackProfileCollector> remote) {
    mojom::TestConnector* connector = test_connectors_.Get(
        mojo::RemoteSetElementId::FromUnsafeValue(connector_id));
    ASSERT_TRUE(connector);

    // BrowserProcessSnapshotController holds the remote end of the
    // mojom::SnapshotController. Pass the other end to the test child if it
    // should be profiled, otherwise drop it. This verifies that an unbound
    // SnapshotController is handled correctly.
    base::OnceClosure bind_receiver_closure;
    if (should_profile_next_launch_) {
      // Unretained is safe since `test_connectors_` won't be destroyed until
      // the reply callback runs.
      bind_receiver_closure =
          base::BindOnce(&mojom::TestConnector::ConnectSnapshotController,
                         base::Unretained(connector), std::move(receiver),
                         on_child_connected_closure_);
    } else {
      bind_receiver_closure = on_child_connected_closure_;
    }

    // The test fixture holds the receiver end of the CallStackProfileCollector.
    // Pass the other end to the test child. The response will eventually
    // trigger `on_child_connected_closure_`.
    connector->ConnectProfileCollector(std::move(remote),
                                       std::move(bind_receiver_closure));
  }

  // Launches a multiprocess test child and registers it with `controller`.
  // The child will simulate a process of type `process_type` and make
  // `num_allocations` memory allocations to report in heap snapshots. If
  // `should_profile` is false, simulate the embedder refusing to profile the
  // child process.
  void LaunchTestChild(HeapProfilerController* controller,
                       ProfilerProcessType process_type,
                       int num_allocations,
                       bool should_profile) {
    // `should_profile` will apply during next call to BindTestConnector().
    base::AutoReset should_profile_next_launch(&should_profile_next_launch_,
                                               should_profile);

    base::LaunchOptions launch_options;
    base::CommandLine child_command_line =
        base::GetMultiProcessTestChildBaseCommandLine();
    child_command_line.AppendSwitchASCII(
        kTestChildTypeSwitch,
        base::NumberToString(static_cast<int>(process_type)));
    child_command_line.AppendSwitchASCII(kTestNumAllocationsSwitch,
                                         base::NumberToString(num_allocations));
    if (should_profile) {
      child_command_line.AppendSwitch(kTestExpectChildProfileSwitch);
    }

    // Attach a mojo channel to the child.
    mojo::PlatformChannel channel;
    channel.PrepareToPassRemoteEndpoint(&launch_options, &child_command_line);
    mojo::OutgoingInvitation invitation;
    mojo::PendingRemote<mojom::TestConnector> pending_connector(
        invitation.AttachMessagePipe(0), 0);
    mojo::RemoteSetElementId connector_id =
        test_connectors_.Add(std::move(pending_connector));

    // In production this only connects the parent end of the SnapshotController
    // since content::ChildProcessHost brokers the interface with the child. For
    // the test, smuggle the id of a TestConnector to broker the interface by
    // pretending it's the child process id.
    controller->AppendCommandLineSwitchForChildProcess(
        &child_command_line, process_type, connector_id.GetUnsafeValue());

    base::Process child_process = base::SpawnMultiProcessTestChild(
        "HeapProfilerControllerChildMain", child_command_line, launch_options);
    ASSERT_TRUE(child_process.IsValid());

    // Finish connecting the mojo channel. This passes the other end of the
    // TestConnector message pipe to the child.
    channel.RemoteProcessLaunchAttempted();
    mojo::OutgoingInvitation::Send(std::move(invitation),
                                   child_process.Handle(),
                                   channel.TakeLocalEndpoint());

    child_processes_.push_back(std::move(child_process));
  }

 private:
  // All child processes started by the test. If a child dies the process will
  // become invalid but remain in this list.
  std::vector<base::Process> child_processes_;

  // Test interface for controlling each child process. If a child dies the
  // interface will be disconnected and removed from this set.
  mojo::RemoteSet<mojom::TestConnector> test_connectors_;

  // Closure to call whenever a child process is finished connecting.
  base::RepeatingClosure on_child_connected_closure_;

  // While this is false, calls to BindTestConnector should skip binding the
  // mojom::SnapshotController, to verify that BrowserProcessSnapshotController
  // can handle an embedder that doesn't bind the controller to some processes.
  bool should_profile_next_launch_ = true;
};

#endif  // ENABLE_MULTIPROCESS_TESTS

class MockSnapshotController : public mojom::SnapshotController {
 public:
  MOCK_METHOD(void, TakeSnapshot, (uint32_t, uint32_t), (override));
  MOCK_METHOD(void, LogMetricsWithoutSnapshot, (), (override));
};

// Configurations of the HeapProfiler* features to test.
// The default parameters collect samples from stable and nonstable channels in
// the browser process only.
struct FeatureTestParams {
  struct ChannelParams {
    double probability = 1.0;
    bool expect_browser_sample = true;
    bool expect_child_sample = false;
  };
  // Whether HeapProfilerReporting is enabled.
  bool feature_enabled = true;
  // Parameters for different channels. Test suites that only cover one channel
  // will use the `stable` params.
  ChannelParams stable;
  ChannelParams nonstable;
  // Probabilities for snapshotting child processes.
  int gpu_snapshot_prob = 100;
  int network_snapshot_prob = 100;
  int renderer_snapshot_prob = 100;
  int utility_snapshot_prob = 100;

  base::FieldTrialParams ToFieldTrialParams() const;

  std::vector<FeatureRefAndParams> GetEnabledFeatures() const;
  std::vector<FeatureRef> GetDisabledFeatures() const;
};

// Converts the test params to field trial parameters for the
// HeapProfilerReporting feature.
base::FieldTrialParams FeatureTestParams::ToFieldTrialParams() const {
  base::FieldTrialParams field_trial_params;

  // Global parameters.
  field_trial_params["stable-probability"] =
      base::NumberToString(stable.probability);
  field_trial_params["nonstable-probability"] =
      base::NumberToString(nonstable.probability);

  // Per-process parameters.
  field_trial_params["browser-sampling-rate-bytes"] =
      base::NumberToString(kSamplingRate);
  field_trial_params["gpu-sampling-rate-bytes"] =
      base::NumberToString(kSamplingRate);
  field_trial_params["gpu-prob-pct"] = base::NumberToString(gpu_snapshot_prob);
  field_trial_params["network-sampling-rate-bytes"] =
      base::NumberToString(kSamplingRate);
  field_trial_params["network-prob-pct"] =
      base::NumberToString(network_snapshot_prob);
  field_trial_params["renderer-sampling-rate-bytes"] =
      base::NumberToString(kSamplingRate);
  field_trial_params["renderer-prob-pct"] =
      base::NumberToString(renderer_snapshot_prob);
  field_trial_params["utility-sampling-rate-bytes"] =
      base::NumberToString(kSamplingRate);
  field_trial_params["utility-prob-pct"] =
      base::NumberToString(utility_snapshot_prob);

  return field_trial_params;
}

std::vector<FeatureRefAndParams> FeatureTestParams::GetEnabledFeatures() const {
  std::vector<FeatureRefAndParams> enabled_features;
  if (feature_enabled) {
    enabled_features.emplace_back(kHeapProfilerReporting, ToFieldTrialParams());
    enabled_features.emplace_back(
        kHeapProfilerChurnReporting,
        base::FieldTrialParams{{"subsampling-chance", "1.0"}});
  }
  return enabled_features;
}

std::vector<FeatureRef> FeatureTestParams::GetDisabledFeatures() const {
  std::vector<FeatureRef> disabled_features;
  if (!feature_enabled) {
    disabled_features.emplace_back(kHeapProfilerReporting);
    disabled_features.emplace_back(kHeapProfilerChurnReporting);
  }
  return disabled_features;
}

// Formats the test params for error messages.
std::ostream& operator<<(std::ostream& os, const FeatureTestParams& params) {
  os << "enabled: " << params.feature_enabled << std::endl;
  os << "field_trial_params: {" << std::endl;
  for (const auto& field_trial_param : params.ToFieldTrialParams()) {
    os << field_trial_param.first << ": " << field_trial_param.second
       << std::endl;
  }
  os << "}" << std::endl;
  os << "expect_samples: {" << std::endl;
  os << "stable/browser: " << params.stable.expect_browser_sample << std::endl;
  os << "stable/child: " << params.stable.expect_child_sample << std::endl;
  os << "nonstable/browser: " << params.stable.expect_browser_sample
     << std::endl;
  os << "nonstable/child: " << params.stable.expect_child_sample << std::endl;
  os << "}";
  return os;
}

// Generic parameterized test suite. Subsets of the tests will alias this to
// create test suites with different parameter lists. The ProfilerProcessType
// parameter is a type of child process to launch, or kUnknown for tests that
// don't launch child processes.
class HeapProfilerControllerTest
    : public ::testing::TestWithParam<
          std::tuple<FeatureTestParams, ProfilerProcessType>>,
      public ProfilerSetUpMixin {
 public:
  // Sets `sample_received_` to true if any sample is received. This will work
  // even without stack unwinding since it doesn't check the contents of the
  // sample. This must be public so that BindRepeating can access it from
  // subclasses.
  void RecordSampleReceived(base::TimeTicks,
                            metrics::SampledProfile sampled_profile) {
    EXPECT_EQ(sampled_profile.trigger_event(),
              metrics::SampledProfile::PERIODIC_HEAP_COLLECTION);
    EXPECT_EQ(sampled_profile.process(), expected_process_);
    // The mock clock should not have advanced since the sample was recorded, so
    // the collection time can be compared exactly.
    const base::TimeDelta expected_time_offset =
        task_env().NowTicks() - profiler_creation_time_;
    EXPECT_EQ(sampled_profile.call_stack_profile().profile_time_offset_ms(),
              expected_time_offset.InMilliseconds());
    sample_received_ = true;
  }

 protected:
  HeapProfilerControllerTest()
      : ProfilerSetUpMixin(feature_params().GetEnabledFeatures(),
                           feature_params().GetDisabledFeatures()) {}

  ~HeapProfilerControllerTest() override {
    // Remove any collectors that were set in StartHeapProfiling.
    metrics::CallStackProfileBuilder::SetBrowserProcessReceiverCallback(
        base::DoNothing());
    metrics::CallStackProfileBuilder::
        ResetChildCallStackProfileCollectorForTesting();
  }

  const FeatureTestParams& feature_params() const {
    return std::get<0>(GetParam());
  }

  const ProfilerProcessType& child_process_type() const {
    return std::get<1>(GetParam());
  }

  // Creates a HeapProfilerController to mock profiling a process of type
  // `process_type` on `channel`. The test should pass `expect_enabled` as true
  // if heap profiling should be enabled in this test setup.
  //
  // `first_snapshot_callback` will be invoked the first time
  // HeapProfilerController::TakeSnapshot() is called, even if it doesn't
  // collect a profile. `collector_callback` will be invoked whenever
  // TakeSnapshot() passes a profile to CallStackProfileBuilder.
  //
  // The test must call StartIfEnabled() after this to start profiling.
  void CreateHeapProfiler(
      version_info::Channel channel,
      ProfilerProcessType process_type,
      bool expect_enabled,
      base::OnceClosure first_snapshot_callback = base::DoNothing(),
      ProfileCollectorCallback collector_callback = base::DoNothing()) {
    ASSERT_FALSE(controller_) << "CreateHeapProfiler called twice";
    if (process_type == ProfilerProcessType::kBrowser) {
      expected_process_ = metrics::Process::BROWSER_PROCESS;
      metrics::CallStackProfileBuilder::SetBrowserProcessReceiverCallback(
          std::move(collector_callback));
    } else {
      switch (process_type) {
        case ProfilerProcessType::kGpu:
          expected_process_ = metrics::Process::GPU_PROCESS;
          break;
        case ProfilerProcessType::kNetworkService:
          expected_process_ = metrics::Process::NETWORK_SERVICE_PROCESS;
          break;
        case ProfilerProcessType::kRenderer:
          expected_process_ = metrics::Process::RENDERER_PROCESS;
          break;
        case ProfilerProcessType::kUtility:
          expected_process_ = metrics::Process::UTILITY_PROCESS;
          break;
        default:
          // Connect up the profile collector even though we expect the heap
          // profiler not to start, so that the test environment is complete.
          expected_process_ = metrics::Process::UNKNOWN_PROCESS;
          break;
      }
      metrics::CallStackProfileBuilder::
          SetParentProfileCollectorForChildProcess(
              AddTestProfileCollector(std::move(collector_callback)));
    }

    ASSERT_FALSE(HeapProfilerController::GetInstance());
    profiler_creation_time_ = task_env().NowTicks();
    controller_ =
        std::make_unique<HeapProfilerController>(channel, process_type);
    controller_->SuppressRandomnessForTesting();
    controller_->SetFirstSnapshotCallbackForTesting(
        std::move(first_snapshot_callback));

    EXPECT_EQ(HeapProfilerController::GetInstance(), controller_.get());
    EXPECT_EQ(controller_->IsEnabled(), expect_enabled);
  }

  // Creates a HeapProfilerController with CreateHeapProfiler() and starts
  // profiling.
  void StartHeapProfiling(
      version_info::Channel channel,
      ProfilerProcessType process_type,
      bool expect_enabled,
      base::OnceClosure first_snapshot_callback = base::DoNothing(),
      ProfileCollectorCallback collector_callback = base::DoNothing()) {
    CreateHeapProfiler(channel, process_type, expect_enabled,
                       std::move(first_snapshot_callback),
                       std::move(collector_callback));
    EXPECT_EQ(controller_->StartIfEnabled(), expect_enabled);
  }

  void AddOneSampleAndWait() {
    // Do nothing if the test has already failed, to avoid timeouts.
    ASSERT_FALSE(HasFailure());
    auto* sampler = base::PoissonAllocationSampler::Get();
    sampler->OnAllocation(AllocationNotificationData(
        reinterpret_cast<void*>(0x1337), kAllocationSize, nullptr,
        AllocationSubsystem::kManualForTesting));
    task_env().RunUntilQuit();
    // Free the allocation so that other tests can re-use the address.
    sampler->OnFree(
        FreeNotificationData(reinterpret_cast<void*>(0x1337),
                             AllocationSubsystem::kManualForTesting));
  }

  // Creates a TestCallStackProfileCollector that accepts callstacks from the
  // and passes them to `collector_callback`. Returns a remote for the profiler
  // to pass the callstacks to.
  mojo::PendingRemote<metrics::mojom::CallStackProfileCollector>
  AddTestProfileCollector(ProfileCollectorCallback collector_callback) {
    mojo::PendingRemote<metrics::mojom::CallStackProfileCollector> remote;
    profile_collector_receivers_.Add(
        std::make_unique<TestCallStackProfileCollector>(
            std::move(collector_callback)),
        remote.InitWithNewPipeAndPassReceiver());
    return remote;
  }

  ScopedCallbacks CreateScopedCallbacks(
      bool expect_take_snapshot,
      bool expect_sampled_profile,
      bool use_other_process_callback = false) {
    return ScopedCallbacks(
        expect_take_snapshot, expect_sampled_profile ? 1 : 0,
        use_other_process_callback,
        base::BindRepeating(&HeapProfilerControllerTest::RecordSampleReceived,
                            base::Unretained(this)),
        task_env().QuitClosure());
  }

  std::unique_ptr<HeapProfilerController> controller_;
  base::HistogramTester histogram_tester_;

  // The creation time of the HeapProfilerController, saved so that
  // RecordSampleReceived() can test that SampledProfile::ms_after_login() in
  // that sample is a delta from the creation time.
  base::TimeTicks profiler_creation_time_;

  // Expected process type in a sample.
  metrics::Process expected_process_ = metrics::Process::UNKNOWN_PROCESS;

  // `sample_received_` is read from the main thread and written from a
  // background thread, but does not need to be atomic because the write happens
  // during a scheduled sample and the read happens well after that.
  bool sample_received_ = false;

  // Receivers for callstack profiles. Each element of the set is a
  // TestCallStackProfileCollecter and associated mojo::Receiver.
  mojo::UniqueReceiverSet<metrics::mojom::CallStackProfileCollector>
      profile_collector_receivers_;
};

// Basic tests only use the default feature params.
INSTANTIATE_TEST_SUITE_P(All,
                         HeapProfilerControllerTest,
                         Combine(Values(FeatureTestParams{}),
                                 // No child process.
                                 Values(ProfilerProcessType::kUnknown)));

// Sampling profiler is not capable of unwinding stack on Android under tests.
#if !BUILDFLAG(IS_ANDROID)
TEST_P(HeapProfilerControllerTest, ProfileCollectionsScheduler) {
  constexpr int kSnapshotsToCollect = 3;

  std::atomic<int> profile_count{0};
  auto check_profile = [&](base::TimeTicks time,
                           metrics::SampledProfile profile) {
    EXPECT_EQ(metrics::SampledProfile::PERIODIC_HEAP_COLLECTION,
              profile.trigger_event());
    EXPECT_LT(0, profile.call_stack_profile().stack_sample_size());

    bool found = false;
    for (const metrics::CallStackProfile::StackSample& sample :
         profile.call_stack_profile().stack_sample()) {
      // Check that the samples being reported are the allocations created
      // below. The sampler calculates the average weight of each sample, and
      // sometimes rounds up to more than kAllocationSize.
      if (sample.has_weight() &&
          static_cast<size_t>(sample.weight()) >= kAllocationSize) {
        found = true;
        break;
      }
    }
    EXPECT_TRUE(found);
    ++profile_count;
  };

  StartHeapProfiling(version_info::Channel::STABLE,
                     ProfilerProcessType::kBrowser,
                     /*expect_enabled=*/true,
                     /*first_snapshot_callback=*/base::DoNothing(),
                     base::BindLambdaForTesting(check_profile));

  auto* sampler = base::PoissonAllocationSampler::Get();
  sampler->OnAllocation(AllocationNotificationData(
      reinterpret_cast<void*>(0x1337), kAllocationSize, nullptr,
      AllocationSubsystem::kManualForTesting));
  sampler->OnAllocation(AllocationNotificationData(
      reinterpret_cast<void*>(0x7331), kAllocationSize, nullptr,
      AllocationSubsystem::kManualForTesting));

  // The profiler should continue to collect snapshots as long as this memory is
  // allocated. If not the test will time out.
  while (profile_count < kSnapshotsToCollect) {
    task_env().FastForwardBy(kCollectionInterval.Get());
  }

  // Free all recorded memory so the address list is empty for the next test.
  sampler->OnFree(FreeNotificationData(reinterpret_cast<void*>(0x1337),
                                       AllocationSubsystem::kManualForTesting));
  sampler->OnFree(FreeNotificationData(reinterpret_cast<void*>(0x7331),
                                       AllocationSubsystem::kManualForTesting));
}
#endif

TEST_P(HeapProfilerControllerTest, UnhandledProcess) {
  // Starting the heap profiler in an unhandled process type should safely do
  // nothing.
  StartHeapProfiling(version_info::Channel::STABLE,
                     ProfilerProcessType::kUnknown,
                     /*expect_enabled=*/false);
  // The Enabled summary histogram should not be logged for unsupported
  // processes, because they're not included in the per-process histograms that
  // are aggregated into it.
  histogram_tester_.ExpectTotalCount("HeapProfiling.InProcess.Enabled", 0);
}

TEST_P(HeapProfilerControllerTest, EmptyProfile) {
  // Should save an empty profile even though no memory is allocated.
  ScopedCallbacks callbacks = CreateScopedCallbacks(
      /*expect_take_snapshot=*/true, /*expect_sampled_profile=*/true);
  StartHeapProfiling(
      version_info::Channel::STABLE, ProfilerProcessType::kBrowser,
      /*expect_enabled=*/true, callbacks.first_snapshot_callback(),
      callbacks.collector_callback());
  task_env().RunUntilQuit();
  EXPECT_TRUE(sample_received_);
}

TEST_P(HeapProfilerControllerTest, HeapChurnProfile) {
  int heap_profiles_received = 0;
  int churn_profiles_received = 0;
  auto check_profile = [&](base::TimeTicks time,
                           metrics::SampledProfile profile) {
    if (profile.trigger_event() ==
        metrics::SampledProfile::PERIODIC_HEAP_COLLECTION) {
      heap_profiles_received++;
    } else if (profile.trigger_event() ==
               metrics::SampledProfile::PERIODIC_HEAP_CHURN_COLLECTION) {
      churn_profiles_received++;
    }
  };

  StartHeapProfiling(version_info::Channel::STABLE,
                     ProfilerProcessType::kBrowser,
                     /*expect_enabled=*/true,
                     /*first_snapshot_callback=*/base::DoNothing(),
                     base::BindLambdaForTesting(check_profile));
  base::SamplingHeapProfiler::Get()->churn_profiler().SetSubsamplingChance(1.0);

  auto* sampler = base::PoissonAllocationSampler::Get();
  void* const kAddress = reinterpret_cast<void*>(0x1234);
  sampler->OnAllocation(
      AllocationNotificationData(kAddress, kAllocationSize, nullptr,
                                 AllocationSubsystem::kManualForTesting));
  sampler->OnFree(
      FreeNotificationData(kAddress, AllocationSubsystem::kManualForTesting));

  task_env().FastForwardBy(kCollectionInterval.Get());
  EXPECT_EQ(heap_profiles_received, 1);
  EXPECT_EQ(churn_profiles_received, 1);
}

TEST_P(HeapProfilerControllerTest, SamplingIntervalVariance) {
  ScopedCallbacks callbacks = CreateScopedCallbacks(
      /*expect_take_snapshot=*/true, /*expect_sampled_profile=*/true);
  StartHeapProfiling(
      version_info::Channel::STABLE, ProfilerProcessType::kBrowser,
      /*expect_enabled=*/true, callbacks.first_snapshot_callback(),
      callbacks.collector_callback());

  // At this point, the profiler is started. The requested sampling rate is
  // kSamplingRate (1024). The actual sampling rate should also be 1024.
  // We change the actual sampling rate to simulate another user changing it.
  base::PoissonAllocationSampler::Get()->SetSamplingInterval(kSamplingRate * 2);

  task_env().RunUntilQuit();

  // The ratio should be (kSamplingRate * 2) / kSamplingRate = 2, which is 200%.
  histogram_tester_.ExpectUniqueSample(
      "HeapProfiling.InProcess.SamplingIntervalVariance.Browser", 200, 1);
  histogram_tester_.ExpectUniqueSample(
      "HeapProfiling.InProcess.SamplingIntervalVariance", 200, 1);

  // Restore the sampling interval so other tests aren't affected.
  base::PoissonAllocationSampler::Get()->SetSamplingInterval(kSamplingRate);
}

// Test the feature on various channels in the browser process.
constexpr FeatureTestParams kChannelConfigs[] = {
    // Disabled.
    {
        .feature_enabled = false,
        .stable = {.expect_browser_sample = false},
        .nonstable = {.expect_browser_sample = false},
    },
    // Enabled, but with probability 0 on all channels.
    {
        .feature_enabled = true,
        .stable = {.probability = 0.0, .expect_browser_sample = false},
        .nonstable = {.probability = 0.0, .expect_browser_sample = false},
    },
    // Enabled on all channels.
    {
        .feature_enabled = true,
        .stable = {.probability = 1.0, .expect_browser_sample = true},
        .nonstable = {.probability = 1.0, .expect_browser_sample = true},
    },
    // Enabled on stable channel only.
    {
        .feature_enabled = true,
        .stable = {.probability = 1.0, .expect_browser_sample = true},
        .nonstable = {.probability = 0.0, .expect_browser_sample = false},
    },
    // Enabled on non-stable channels only.
    {
        .feature_enabled = true,
        .stable = {.probability = 0.0, .expect_browser_sample = false},
        .nonstable = {.probability = 1.0, .expect_browser_sample = true},
    },
};

using HeapProfilerControllerChannelTest = HeapProfilerControllerTest;

INSTANTIATE_TEST_SUITE_P(All,
                         HeapProfilerControllerChannelTest,
                         Combine(ValuesIn(kChannelConfigs),
                                 // No child process.
                                 Values(ProfilerProcessType::kUnknown)));

TEST_P(HeapProfilerControllerChannelTest, StableChannel) {
  const bool profiling_enabled = feature_params().feature_enabled &&
                                 feature_params().stable.probability > 0.0;
  ScopedCallbacks callbacks = CreateScopedCallbacks(
      /*expect_take_snapshot=*/profiling_enabled,
      feature_params().stable.expect_browser_sample);
  StartHeapProfiling(version_info::Channel::STABLE,
                     ProfilerProcessType::kBrowser, profiling_enabled,
                     callbacks.first_snapshot_callback(),
                     callbacks.collector_callback());
  histogram_tester_.ExpectUniqueSample(
      "HeapProfiling.InProcess.Enabled.Browser", profiling_enabled, 1);
  histogram_tester_.ExpectUniqueSample("HeapProfiling.InProcess.Enabled",
                                       profiling_enabled, 1);
  AddOneSampleAndWait();
  EXPECT_EQ(sample_received_, feature_params().stable.expect_browser_sample);
}

TEST_P(HeapProfilerControllerChannelTest, CanaryChannel) {
  const bool profiling_enabled = feature_params().feature_enabled &&
                                 feature_params().nonstable.probability > 0.0;
  ScopedCallbacks callbacks = CreateScopedCallbacks(
      /*expect_take_snapshot=*/profiling_enabled,
      feature_params().nonstable.expect_browser_sample);
  StartHeapProfiling(version_info::Channel::CANARY,
                     ProfilerProcessType::kBrowser, profiling_enabled,
                     callbacks.first_snapshot_callback(),
                     callbacks.collector_callback());
  histogram_tester_.ExpectUniqueSample(
      "HeapProfiling.InProcess.Enabled.Browser", profiling_enabled, 1);
  histogram_tester_.ExpectUniqueSample("HeapProfiling.InProcess.Enabled",
                                       profiling_enabled, 1);
  AddOneSampleAndWait();
  EXPECT_EQ(sample_received_, feature_params().nonstable.expect_browser_sample);
}

TEST_P(HeapProfilerControllerChannelTest, UnknownChannel) {
  // An unknown channel should be treated like stable, in case a large
  // population doesn't have the channel set.
  const bool profiling_enabled = feature_params().feature_enabled &&
                                 feature_params().stable.probability > 0.0;
  ScopedCallbacks callbacks = CreateScopedCallbacks(
      /*expect_take_snapshot=*/profiling_enabled,
      feature_params().stable.expect_browser_sample);
  StartHeapProfiling(version_info::Channel::UNKNOWN,
                     ProfilerProcessType::kBrowser, profiling_enabled,
                     callbacks.first_snapshot_callback(),
                     callbacks.collector_callback());
  histogram_tester_.ExpectUniqueSample(
      "HeapProfiling.InProcess.Enabled.Browser", profiling_enabled, 1);
  histogram_tester_.ExpectUniqueSample("HeapProfiling.InProcess.Enabled",
                                       profiling_enabled, 1);
  AddOneSampleAndWait();
  EXPECT_EQ(sample_received_, feature_params().stable.expect_browser_sample);
}

// Test the feature in various processes on the stable channel.
constexpr FeatureTestParams kProcessConfigs[] = {
    // Enabled in parent process only.
    {
        .stable = {.probability = 1.0,
                   .expect_browser_sample = true,
                   .expect_child_sample = false},
        .gpu_snapshot_prob = 0,
        .network_snapshot_prob = 0,
        .renderer_snapshot_prob = 0,
        .utility_snapshot_prob = 0,
    },
    // Enabled in child process only.
    // Central control only samples child processes when the browser process is
    // sampled, so no samples are expected even though sampling is supported in
    // the child process.
    {
        .stable = {.probability = 0.0,
                   .expect_browser_sample = false,
                   .expect_child_sample = false},
        .gpu_snapshot_prob = 100,
        .network_snapshot_prob = 100,
        .renderer_snapshot_prob = 100,
        .utility_snapshot_prob = 100,
    },
    // Enabled in parent and child processes.
    {
        .stable = {.probability = 1.0,
                   .expect_browser_sample = true,
                   .expect_child_sample = true},
        .gpu_snapshot_prob = 100,
        .network_snapshot_prob = 100,
        .renderer_snapshot_prob = 100,
        .utility_snapshot_prob = 100,
    },
};

using HeapProfilerControllerProcessTest = HeapProfilerControllerTest;

INSTANTIATE_TEST_SUITE_P(All,
                         HeapProfilerControllerProcessTest,
                         Combine(ValuesIn(kProcessConfigs),
                                 Values(ProfilerProcessType::kGpu,
                                        ProfilerProcessType::kNetworkService,
                                        ProfilerProcessType::kRenderer,
                                        ProfilerProcessType::kUtility,
                                        // Include unsupported process types.
                                        ProfilerProcessType::kUnknown)));

TEST_P(HeapProfilerControllerProcessTest, BrowserProcess) {
  const bool profiling_enabled = feature_params().stable.expect_browser_sample;
  ScopedCallbacks callbacks = CreateScopedCallbacks(
      /*expect_take_snapshot=*/profiling_enabled,
      /*expect_sampled_profile=*/profiling_enabled,
      /*use_other_process_callback=*/
      feature_params().stable.expect_child_sample &&
          child_process_type() != ProfilerProcessType::kUnknown);

  // Mock the child end of the SnapshotController mojo pipe.
  MockSnapshotController mock_child_snapshot_controller;
  mojo::Receiver<mojom::SnapshotController> mock_receiver(
      &mock_child_snapshot_controller);

  StartHeapProfiling(version_info::Channel::STABLE,
                     ProfilerProcessType::kBrowser, profiling_enabled,
                     callbacks.first_snapshot_callback(),
                     callbacks.collector_callback());
  histogram_tester_.ExpectUniqueSample(
      "HeapProfiling.InProcess.Enabled.Browser", profiling_enabled, 1);
  histogram_tester_.ExpectUniqueSample("HeapProfiling.InProcess.Enabled",
                                       profiling_enabled, 1);

  constexpr int kTestChildProcessId = 1;
  if (profiling_enabled) {
    ASSERT_TRUE(controller_->GetBrowserProcessSnapshotController());

    // This callback should be invoked from
    // AppendCommandLineSwitchForChildProcess to bind the child end of the mojo
    // pipe.
    controller_->GetBrowserProcessSnapshotController()
        ->SetBindRemoteForChildProcessCallback(base::BindLambdaForTesting(
            [&](int child_process_id,
                mojo::PendingReceiver<mojom::SnapshotController> receiver) {
              // Should not be called if profiling is unsupported in the child.
              ASSERT_TRUE(feature_params().stable.expect_child_sample);
              ASSERT_NE(child_process_type(), ProfilerProcessType::kUnknown);
              EXPECT_EQ(child_process_id, kTestChildProcessId);
              mock_receiver.Bind(std::move(receiver));
            }));
  } else {
    EXPECT_FALSE(controller_->GetBrowserProcessSnapshotController());
  }

  // Simulate a child process launch. If profiling is enabled in both browser
  // and child processes, this will bind the browser end of the mojo pipe to the
  // BrowserProcessSnapshotController and use the above callback to bind the
  // child end to `mock_child_snapshot_controller`.
  base::CommandLine child_command_line(base::CommandLine::NO_PROGRAM);
  controller_->AppendCommandLineSwitchForChildProcess(
      &child_command_line, child_process_type(), kTestChildProcessId);

  if (feature_params().stable.expect_child_sample &&
      child_process_type() != ProfilerProcessType::kUnknown) {
    EXPECT_CALL(mock_child_snapshot_controller, TakeSnapshot(100, 0))
        .WillOnce([&] {
          // Record that BrowserProcessSnapshotController triggered a fake
          // snapshot in the child process.
          callbacks.other_process_callback().Run();
        });
  } else {
    EXPECT_CALL(mock_child_snapshot_controller, TakeSnapshot(_, _)).Times(0);
  }

  AddOneSampleAndWait();
  EXPECT_EQ(sample_received_, feature_params().stable.expect_browser_sample);
}

TEST_P(HeapProfilerControllerProcessTest, ChildProcess) {
  // TakeSnapshot() is only called in the child process when the browser process
  // triggers it. Nothing to test if sampling in the browser process is
  // disabled.
  if (!feature_params().stable.expect_browser_sample) {
    return;
  }
  const bool profiling_enabled =
      feature_params().stable.expect_child_sample &&
      child_process_type() != ProfilerProcessType::kUnknown;

  ScopedCallbacks callbacks = CreateScopedCallbacks(
      /*expect_take_snapshot=*/profiling_enabled,
      /*expect_sampled_profile=*/profiling_enabled,
      /*use_other_process_callback=*/true);

  // Simulate the browser side of child process launching.
  constexpr int kTestChildProcessId = 1;

  // Create a snapshot controller to hold the browser end of the mojo pipe.
  auto snapshot_task_runner = base::SequencedTaskRunner::GetCurrentDefault();
  auto fake_browser_snapshot_controller =
      std::make_unique<BrowserProcessSnapshotController>(snapshot_task_runner);

  // This callback should be invoked from AppendCommandLineSwitchForTesting to
  // bind the child end of the mojo pipe.
  fake_browser_snapshot_controller->SetBindRemoteForChildProcessCallback(
      base::BindLambdaForTesting(
          [&](int child_process_id,
              mojo::PendingReceiver<mojom::SnapshotController> receiver) {
            // Should not be called if profiling is unsupported in the child.
            ASSERT_TRUE(feature_params().stable.expect_child_sample);
            ASSERT_NE(child_process_type(), ProfilerProcessType::kUnknown);
            EXPECT_EQ(child_process_id, kTestChildProcessId);
            ChildProcessSnapshotController::CreateSelfOwnedReceiver(
                std::move(receiver));
          }));

  base::test::ScopedCommandLine scoped_command_line;
  HeapProfilerController::AppendCommandLineSwitchForTesting(
      scoped_command_line.GetProcessCommandLine(), child_process_type(),
      kTestChildProcessId, fake_browser_snapshot_controller.get());

  // Simulate the browser process taking a sample after a delay.
  snapshot_task_runner->PostDelayedTask(
      FROM_HERE,
      base::BindOnce(
          &BrowserProcessSnapshotController::TakeSnapshotsOnSnapshotSequence,
          std::move(fake_browser_snapshot_controller))
          .Then(callbacks.other_process_callback()),
      TestTimeouts::action_timeout());

  StartHeapProfiling(version_info::Channel::STABLE, child_process_type(),
                     profiling_enabled, callbacks.first_snapshot_callback(),
                     callbacks.collector_callback());

  // If the child process is unknown, no histograms are logged. If it's known
  // but disabled, 0 is logged.
  if (child_process_type() != ProfilerProcessType::kUnknown) {
    switch (child_process_type()) {
      case ProfilerProcessType::kGpu:
        histogram_tester_.ExpectUniqueSample(
            "HeapProfiling.InProcess.Enabled.GPU", profiling_enabled, 1);
        break;
      case ProfilerProcessType::kNetworkService:
        histogram_tester_.ExpectUniqueSample(
            "HeapProfiling.InProcess.Enabled.NetworkService", profiling_enabled,
            1);
        break;
      case ProfilerProcessType::kRenderer:
        histogram_tester_.ExpectUniqueSample(
            "HeapProfiling.InProcess.Enabled.Renderer", profiling_enabled, 1);
        break;
      case ProfilerProcessType::kUtility:
        histogram_tester_.ExpectUniqueSample(
            "HeapProfiling.InProcess.Enabled.Utility", profiling_enabled, 1);
        break;
      default:
        FAIL() << "Unexpected processs type "
               << static_cast<int>(child_process_type());
    }
    histogram_tester_.ExpectUniqueSample("HeapProfiling.InProcess.Enabled",
                                         profiling_enabled, 1);
  }

  // The child process HeapProfilerController should never have a
  // BrowserProcessSnapshotController. (`fake_browser_snapshot_controller`
  // simulates the browser side of the connection.)
  EXPECT_EQ(controller_->GetBrowserProcessSnapshotController(), nullptr);

  AddOneSampleAndWait();
  EXPECT_EQ(sample_received_, profiling_enabled);
}

#if ENABLE_MULTIPROCESS_TESTS

// Returns a lambda that can be called from a GMock matcher. It will return the
// value of a MetadataItem named `name` in a given CallStackProfile, or nullopt
// if there's no metadata with that name.
auto GetProfileMetadataFunc(std::string_view name) {
  auto get_metadata =
      [name_hash = base::HashMetricName(name)](
          const metrics::CallStackProfile& profile) -> std::optional<int64_t> {
    for (int32_t i = 0; i < profile.metadata_name_hash_size(); ++i) {
      if (profile.metadata_name_hash(i) == name_hash) {
        // Found index of `name_hash`.
        for (const auto& metadata_item : profile.profile_metadata()) {
          if (metadata_item.name_hash_index() == i) {
            return metadata_item.value();
          }
        }
      }
    }
    // No metadata matched `name_hash`.
    return std::nullopt;
  };
  return get_metadata;
}

// End-to-end test with multiple child processes.
constexpr const auto kMultipleChildConfigs = std::to_array<FeatureTestParams>({
    {
        .gpu_snapshot_prob = 100,
        .network_snapshot_prob = 100,
        .renderer_snapshot_prob = 66,
        .utility_snapshot_prob = 50,
    },
});

using HeapProfilerControllerMultipleChildTest = HeapProfilerControllerTest;

INSTANTIATE_TEST_SUITE_P(All,
                         HeapProfilerControllerMultipleChildTest,
                         Combine(ValuesIn(kMultipleChildConfigs),
                                 // Children are manually specified.
                                 Values(ProfilerProcessType::kUnknown)));

MULTIPROCESS_TEST_MAIN(HeapProfilerControllerChildMain) {
  MultiprocessTestChild child(kMultipleChildConfigs[0].GetEnabledFeatures(),
                              kMultipleChildConfigs[0].GetDisabledFeatures());
  child.RunTestInChild();
  return ::testing::Test::HasFailure();
}

TEST_P(HeapProfilerControllerMultipleChildTest, EndToEnd) {
  // Initialize mojo IPC support.
  mojo::core::ScopedIPCSupport enable_mojo(
      base::SingleThreadTaskRunner::GetCurrentDefault(),
      mojo::core::ScopedIPCSupport::ShutdownPolicy::CLEAN);

  // Process types to test. Each will make a different
  // number of memory allocations so their reports are all different.
  const std::vector<std::pair<ProfilerProcessType, size_t>> kProcessesToTest{
      {ProfilerProcessType::kBrowser, 0},
      {ProfilerProcessType::kGpu, 1},
      // 2 utility processes.
      {ProfilerProcessType::kUtility, 2},
      {ProfilerProcessType::kUtility, 3},
      // 5 renderer processes including one with no samples. The first one will
      // be ignored to simulate the embedder refusing to profile it.
      {ProfilerProcessType::kRenderer, 10},
      {ProfilerProcessType::kRenderer, 0},
      {ProfilerProcessType::kRenderer, 4},
      {ProfilerProcessType::kRenderer, 5},
      {ProfilerProcessType::kRenderer, 6},
  };

  // Expect only 1 utility process and 3 renderer processes to be sampled due
  // to the "renderer-prob" and "utility-prob" params.
  constexpr size_t kExpectedSampledProfiles =
      /*browser*/ 1 + /*gpu*/ 1 + /*utility*/ 1 + /*renderer*/ 3;

  // Create callbacks that store profiles from all processes in a vector.
  std::vector<metrics::SampledProfile> received_profiles;
  auto collector_callback = base::BindLambdaForTesting(
      [&](base::TimeTicks, metrics::SampledProfile profile) {
        received_profiles.push_back(std::move(profile));
      });
  ScopedCallbacks callbacks(
      /*expect_take_snapshot=*/true, kExpectedSampledProfiles,
      /*use_other_process_callback=*/false, std::move(collector_callback),
      task_env().QuitClosure());

  // Snapshots from the children take real time to be passed back to the parent.
  // The mock clock will advance to the next snapshot time while waiting, so
  // stop profiling after the first snapshot by deleting the controller.
  auto stop_after_first_snapshot_callback =
      callbacks.first_snapshot_callback().Then(base::BindLambdaForTesting(
          [this, task_runner = base::SequencedTaskRunner::GetCurrentDefault()] {
            task_runner->DeleteSoon(FROM_HERE, controller_.release());
          }));

  CreateHeapProfiler(
      version_info::Channel::STABLE, ProfilerProcessType::kBrowser,
      /*expect_enabled=*/true, std::move(stop_after_first_snapshot_callback),
      callbacks.collector_callback());
  ASSERT_TRUE(controller_);

  // Start all processes in `kProcessesToTest` except the browser.
  MultiprocessTestParent test_parent;
  test_parent.StartHeapProfilingWhenChildrenConnected(
      kProcessesToTest.size() - 1, controller_.get());

  // On every process launch, create a TestCallStackProfileCollector to collect
  // profiles from the child. BrowserProcessSnapshotController will create a
  // SnapshotController to trigger snapshots in the child.
  auto* browser_snapshot_controller =
      controller_->GetBrowserProcessSnapshotController();
  ASSERT_TRUE(browser_snapshot_controller);
  auto binder_callback = base::BindLambdaForTesting(
      [&](int id, mojo::PendingReceiver<mojom::SnapshotController> receiver) {
        mojo::PendingRemote<metrics::mojom::CallStackProfileCollector> remote =
            AddTestProfileCollector(callbacks.collector_callback());
        test_parent.BindTestConnector(id, std::move(receiver),
                                      std::move(remote));
      });
  browser_snapshot_controller->SetBindRemoteForChildProcessCallback(
      std::move(binder_callback));

  bool renderer_was_skipped = false;
  for (const auto [process_type, num_allocations] : kProcessesToTest) {
    if (process_type != ProfilerProcessType::kBrowser) {
      // Skip the first renderer.
      bool should_profile = true;
      if (process_type == ProfilerProcessType::kRenderer &&
          !renderer_was_skipped) {
        should_profile = false;
        renderer_was_skipped = true;
      }
      test_parent.LaunchTestChild(controller_.get(), process_type,
                                  num_allocations, should_profile);
    }
  }

  // Loop until all children are connected and all processes send snapshots.
  task_env().RunUntilQuit();

  // GMock matcher that tests that the given CallStackProfile contains `count`
  // stack samples with metadata containing `process_percent` and
  // "process_index" < `sampled_processes`.
  auto call_stack_profile_matches = [](size_t count, int64_t process_percent,
                                       int64_t sampled_processes) {
    using StackSample = metrics::CallStackProfile::StackSample;
    return AllOf(
        Property(
            "stack_sample", &metrics::CallStackProfile::stack_sample,
            Conditional(
                count > 0,
                // The test makes allocations at addresses without symbols, so
                // they're all counted in the same stack frame.
                ElementsAre(AllOf(Property("count", &StackSample::count, count),
                                  Property("weight", &StackSample::weight,
                                           count * kAllocationSize))),
                // No allocations means no stack frames.
                IsEmpty())),
        ResultOf("process_percent metadata",
                 GetProfileMetadataFunc("process_percent"),
                 Optional(process_percent)),
        ResultOf("process_index metadata",
                 GetProfileMetadataFunc("process_index"),
                 // Processes can be sampled in any order, so just check the
                 // range of "process_index".
                 Optional(AllOf(Ge(0), Lt(sampled_processes)))),
        ResultOf("sampling_interval metadata",
                 GetProfileMetadataFunc("sampling_interval"),
                 Optional(testing::Gt(0))));
  };

  // GMock matcher that tests that the given SampledProfile is a heap snapshot
  // for the given `process_type` containing `count` stack samples with metadata
  // containing `process_percent` and "process_index" < `sampled_processes`.
  auto sampled_profile_matches = [&](metrics::Process process_type,
                                     size_t count, int64_t process_percent,
                                     int64_t sampled_processes) {
    return AllOf(
        Property("process", &metrics::SampledProfile::process, process_type),
        Property("call_stack_profile",
                 &metrics::SampledProfile::call_stack_profile,
                 call_stack_profile_matches(count, process_percent,
                                            sampled_processes)));
  };

  // Only the first 1/2 of utility processes and 2/3 of renderers should be
  // included due to sampling. Renderers are rounded up to 3 of the 4 that can
  // be profiled - the 5th is invisible to the profiler.
  EXPECT_THAT(
      received_profiles,
      UnorderedElementsAre(
          sampled_profile_matches(metrics::Process::BROWSER_PROCESS, 0, 100, 1),
          sampled_profile_matches(metrics::Process::GPU_PROCESS, 1, 100, 1),
          sampled_profile_matches(metrics::Process::UTILITY_PROCESS, 2, 50, 1),
          // The first renderer should be skipped.
          sampled_profile_matches(metrics::Process::RENDERER_PROCESS, 0, 66, 3),
          sampled_profile_matches(metrics::Process::RENDERER_PROCESS, 4, 66, 3),
          sampled_profile_matches(metrics::Process::RENDERER_PROCESS, 5, 66,
                                  3)));

  // Make sure both per-process and aggregate profiler stats are logged.
  // Subprocess metrics aren't hooked up in this test, so `histogram_tester_`
  // only sees the browser process histograms.
  histogram_tester_.ExpectTotalCount(
      "HeapProfiling.InProcess.SamplesPerSnapshot.Browser", 1);
  histogram_tester_.ExpectTotalCount(
      "HeapProfiling.InProcess.SamplesPerSnapshot", 1);
}

#endif  // ENABLE_MULTIPROCESS_TESTS

}  // namespace

}  // namespace heap_profiling