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

#include "services/tracing/public/cpp/system_metrics_sampler.h"

#include "base/check.h"
#include "base/no_destructor.h"
#include "base/power_monitor/cpu_frequency_utils.h"
#include "base/task/sequenced_task_runner.h"
#include "base/task/thread_pool.h"
#include "base/time/time.h"
#include "base/trace_event/common/trace_event_common.h"
#include "base/trace_event/trace_event.h"
#include "services/tracing/public/mojom/perfetto_service.mojom.h"
#include "third_party/perfetto/include/perfetto/tracing/core/data_source_descriptor.h"
#include "third_party/perfetto/protos/perfetto/config/chrome/system_metrics.gen.h"
#include "third_party/perfetto/protos/perfetto/config/data_source_config.gen.h"

#if BUILDFLAG(IS_WIN)
#include <windows.h>

// Psapi.h must come after Windows.h.
#include <psapi.h>
#endif  // BUILDFLAG(IS_WIN)

namespace tracing {

namespace {

constexpr base::TimeDelta kDefaultSamplingInterval = base::Seconds(5);

}  // namespace

void SystemMetricsSampler::Register(bool system_wide) {
  perfetto::DataSourceDescriptor desc;
  desc.set_name(tracing::mojom::kSystemMetricsSourceName);
  perfetto::DataSource<SystemMetricsSampler>::Register(desc, system_wide);
}

SystemMetricsSampler::SystemMetricsSampler(bool system_wide)
    : system_wide_(system_wide), sampling_interval_(kDefaultSamplingInterval) {}
SystemMetricsSampler::~SystemMetricsSampler() = default;

void SystemMetricsSampler::OnSetup(const SetupArgs& args) {
  if (args.config->chromium_system_metrics_raw().empty()) {
    return;
  }
  perfetto::protos::gen::ChromiumSystemMetricsConfig config;
  if (!config.ParseFromString(args.config->chromium_system_metrics_raw())) {
    DLOG(ERROR) << "Failed to parse chromium_system_metrics";
    return;
  }
  if (config.has_sampling_interval_ms()) {
    sampling_interval_ = base::Milliseconds(config.sampling_interval_ms());
  }
}

void SystemMetricsSampler::OnStart(const StartArgs&) {
  if (system_wide_) {
    system_sampler_ = base::SequenceBound<SystemSampler>(
        base::ThreadPool::CreateSequencedTaskRunner({}), sampling_interval_);
  }
  process_sampler_ = base::SequenceBound<ProcessSampler>(
      base::ThreadPool::CreateSequencedTaskRunner({}), sampling_interval_);
}

void SystemMetricsSampler::OnStop(const StopArgs&) {
  system_sampler_.Reset();
  process_sampler_.Reset();
}

SystemMetricsSampler::SystemSampler::SystemSampler(
    base::TimeDelta sampling_interval)
    : cpu_probe_{system_cpu::CpuProbe::Create()} {
  if (cpu_probe_) {
    cpu_probe_->StartSampling();
  }
  sample_timer_.Start(FROM_HERE, sampling_interval, this,
                      &SystemSampler::SampleSystemMetrics);
}

SystemMetricsSampler::SystemSampler::~SystemSampler() = default;

void SystemMetricsSampler::SystemSampler::SampleSystemMetrics() {
  if (cpu_probe_) {
    cpu_probe_->RequestSample(base::BindOnce(&SystemSampler::OnCpuProbeResult,
                                             base::Unretained(this)));
  }
  std::optional<base::CpuThroughputEstimationResult> cpu_throughput =
      base::EstimateCpuThroughput();
  if (cpu_throughput) {
    TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"),
                  perfetto::CounterTrack("EstimatedCpuThroughput",
                                         perfetto::Track::Global(0)),
                  cpu_throughput->estimated_frequency);
  }

#if BUILDFLAG(IS_WIN)
  base::CpuFrequencyInfo cpu_info = base::GetCpuFrequencyInfo();
  TRACE_COUNTER(
      TRACE_DISABLED_BY_DEFAULT("system_metrics"),
      perfetto::CounterTrack("NumActiveCpus", perfetto::Track::Global(0)),
      cpu_info.num_active_cpus);

  SampleMemoryMetrics();
#endif
}

#if BUILDFLAG(IS_WIN)
void SystemMetricsSampler::SystemSampler::SampleMemoryMetrics() {
  MEMORYSTATUSEX mem_status = {};
  mem_status.dwLength = sizeof(mem_status);
  if (!::GlobalMemoryStatusEx(&mem_status)) {
    return;
  }

  TRACE_COUNTER(
      TRACE_DISABLED_BY_DEFAULT("system_metrics"),
      perfetto::CounterTrack("CommitMemoryLimit", perfetto::Track::Global(0)),
      mem_status.ullTotalPageFile);

  TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"),
                perfetto::CounterTrack("CommitMemoryAvailable",
                                       perfetto::Track::Global(0)),
                mem_status.ullAvailPageFile);
  TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"),
                perfetto::CounterTrack("AvailablePhysicalMemory",
                                       perfetto::Track::Global(0)),
                mem_status.ullAvailPhys);
}
#endif  // BUILDFLAG(IS_WIN)

void SystemMetricsSampler::SystemSampler::OnCpuProbeResult(
    std::optional<system_cpu::CpuSample> cpu_sample) {
  if (!cpu_sample) {
    return;
  }
  TRACE_COUNTER(
      TRACE_DISABLED_BY_DEFAULT("system_metrics"),
      perfetto::CounterTrack("SystemCpuUsage", perfetto::Track::Global(0)),
      cpu_sample->cpu_utilization);
}

SystemMetricsSampler::ProcessSampler::ProcessSampler(
    base::TimeDelta sampling_interval) {
  process_metrics_ = base::ProcessMetrics::CreateCurrentProcessMetrics();
  SampleProcessMetrics();
  sample_timer_.Start(FROM_HERE, sampling_interval, this,
                      &ProcessSampler::SampleProcessMetrics);
}

SystemMetricsSampler::ProcessSampler::~ProcessSampler() = default;

void SystemMetricsSampler::ProcessSampler::SampleProcessMetrics() {
  auto cpu_usage = process_metrics_->GetPlatformIndependentCPUUsage();
  if (cpu_usage.has_value()) {
    TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"), "CpuUsage",
                  *cpu_usage);
  }
  auto memory_info = process_metrics_->GetMemoryInfo();
  if (memory_info.has_value()) {
    TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"), "ResidentSet",
                  memory_info->resident_set_bytes);
#if BUILDFLAG(IS_MAC)
    TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"),
                  "PhysicalMemoryFootprint",
                  memory_info->physical_footprint_bytes);
#elif BUILDFLAG(IS_LINUX) || BUILDFLAG(IS_CHROMEOS) || \
    BUILDFLAG(IS_ANDROID) || BUILDFLAG(IS_FUCHSIA)
    TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"), "VmSwapMemory",
                  memory_info->vm_swap_bytes);
    TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"), "RssAnonMemory",
                  memory_info->rss_anon_bytes);
#elif BUILDFLAG(IS_WIN)
    TRACE_COUNTER(TRACE_DISABLED_BY_DEFAULT("system_metrics"), "PrivateMemory",
                  memory_info->private_bytes);
#endif  // BUILDFLAG(IS_WIN)
  }
}

}  // namespace tracing