// Copyright 2019 The Chromium Authors. All rights reserved.
// 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/perfetto/posix_system_producer.h"

#include <utility>

#include "base/bind.h"
#include "base/strings/strcat.h"
#include "base/strings/string_number_conversions.h"
#include "base/trace_event/trace_log.h"
#include "build/build_config.h"
#include "services/tracing/public/cpp/perfetto/shared_memory.h"
#include "services/tracing/public/cpp/trace_startup.h"
#include "services/tracing/public/cpp/traced_process_impl.h"
#include "services/tracing/public/cpp/tracing_features.h"
#include "third_party/perfetto/include/perfetto/ext/tracing/core/commit_data_request.h"
#include "third_party/perfetto/include/perfetto/ext/tracing/core/shared_memory_arbiter.h"
#include "third_party/perfetto/include/perfetto/ext/tracing/core/trace_writer.h"
#include "third_party/perfetto/include/perfetto/ext/tracing/ipc/producer_ipc_client.h"
#include "third_party/perfetto/include/perfetto/protozero/scattered_heap_buffer.h"
#include "third_party/perfetto/include/perfetto/protozero/scattered_stream_writer.h"
#include "third_party/perfetto/protos/perfetto/common/track_event_descriptor.pbzero.h"

#if defined(OS_ANDROID)
#include "base/android/build_info.h"
#endif  // defined(OS_ANDROID)

namespace tracing {
namespace {
constexpr uint32_t kInitialConnectionBackoffMs = 100;
constexpr uint32_t kMaxConnectionBackoffMs = 30 * 1000;

perfetto::DataSourceConfig EnsureGuardRailsAreFollowed(
    const perfetto::DataSourceConfig& data_source_config) {
  if (!data_source_config.enable_extra_guardrails() ||
      data_source_config.chrome_config().privacy_filtering_enabled()) {
    return data_source_config;
  }
  // If extra_guardrails is enabled then we have to ensure we have privacy
  // filtering enabled.
  perfetto::DataSourceConfig config = data_source_config;
  config.mutable_chrome_config()->set_privacy_filtering_enabled(true);
  return config;
}

uint32_t IncreaseBackoff(uint32_t current, uint32_t max) {
  return std::min(current * 2, max);
}
}  // namespace

PosixSystemProducer::PosixSystemProducer(const char* socket,
                                         PerfettoTaskRunner* task_runner)
    : SystemProducer(task_runner),
      socket_name_(socket),
      connection_backoff_ms_(kInitialConnectionBackoffMs) {
  DETACH_FROM_SEQUENCE(sequence_checker_);
}

PosixSystemProducer::~PosixSystemProducer() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
}

void PosixSystemProducer::SetDisallowPreAndroidPieForTesting(bool disallow) {
  bool was_disallowed = SkipIfOnAndroidAndPreAndroidPie();
  disallow_pre_android_pie_ = disallow;
  if (!disallow && was_disallowed && state_ == State::kDisconnected) {
    // If previously we would not have connected, we now attempt to connect
    // since we are now skipping a check.
    Connect();
  }
}

void PosixSystemProducer::SetNewSocketForTesting(const char* socket) {
  socket_name_ = socket;

  if (state_ == State::kDisconnected) {
    // Not connected yet, wait for ConnectToSystemService().
    return;
  }

  if (state_ == State::kConnected) {
    // If we are fully connected we need to reset the service before we
    // reconnect.
    DisconnectWithReply(base::BindOnce(&PosixSystemProducer::OnDisconnect,
                                       base::Unretained(this)));
    return;
  }

  // In any other case, we just need to do a normal disconnect and
  // DisconnectWithReply will ensure we set up the retries on the new |socket|.
  DisconnectWithReply(base::OnceClosure());
}

perfetto::SharedMemoryArbiter* PosixSystemProducer::MaybeSharedMemoryArbiter() {
  base::AutoLock lock(lock_);
  DCHECK(GetService());
  return GetService()->MaybeSharedMemoryArbiter();
}

void PosixSystemProducer::NewDataSourceAdded(
    const PerfettoTracedProcess::DataSourceBase* const data_source) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (state_ != State::kConnected) {
    return;
  }
  perfetto::DataSourceDescriptor new_registration;
  new_registration.set_name(data_source->name());
  new_registration.set_will_notify_on_start(true);
  new_registration.set_will_notify_on_stop(true);
  new_registration.set_handles_incremental_state_clear(true);

  // Add categories to the DataSourceDescriptor.
  std::set<std::string> category_set;
  tracing::TracedProcessImpl::GetInstance()->GetCategories(&category_set);
  protozero::HeapBuffered<perfetto::protos::pbzero::TrackEventDescriptor> proto;
  for (const std::string& s : category_set) {
    auto* cat = proto->add_available_categories();
    cat->set_name(s);
    if (s.find(TRACE_DISABLED_BY_DEFAULT("")) == 0) {
      cat->add_tags("slow");
    }
  }
  new_registration.set_track_event_descriptor_raw(proto.SerializeAsString());

  GetService()->RegisterDataSource(new_registration);
}

bool PosixSystemProducer::IsTracingActive() {
  base::AutoLock lock(lock_);
  return data_sources_tracing_ > 0;
}

void PosixSystemProducer::ConnectToSystemService() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  CHECK(IsTracingInitialized());
  DCHECK(state_ == State::kDisconnected);

  // Some Telemetry tests use sideloaded Perfetto library on pre-Pie devices.
  // We allow those tests to use system tracing by setting the
  // EnablePerfettoSystemTracing feature.
  disallow_pre_android_pie_ =
      !base::FeatureList::IsEnabled(features::kEnablePerfettoSystemTracing);

  Connect();
}

void PosixSystemProducer::ActivateTriggers(
    const std::vector<std::string>& triggers) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (state_ == State::kConnected) {
    GetService()->ActivateTriggers(triggers);
  }
}

void PosixSystemProducer::DisconnectWithReply(
    base::OnceClosure on_disconnect_complete) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (state_ == State::kConnected) {
    // We are connected and need to unregister the DataSources to
    // inform the service these data sources are going away. If we
    // are currently tracing the service will ask for them to shut
    // down asynchronously.
    //
    // Note that the system service may have concurrently posted a
    // task to request one of these data sources to start. However we
    // will ignore such requests by verifying that we're allowed to
    // trace in StartDataSource().
    for (const auto* const data_source :
         PerfettoTracedProcess::Get()->data_sources()) {
      DCHECK(GetService());
      GetService()->UnregisterDataSource(data_source->name());
    }
    state_ = State::kUnregistered;
  }
  // If we are tracing we need to wait until we're fully disconnected
  // to run the callback, otherwise we run it immediately (we will
  // still unregister the data sources but that can happen async in
  // the background).
  if (!on_disconnect_complete.is_null()) {
    if (IsTracingActive() || !on_disconnect_callbacks_.empty()) {
      on_disconnect_callbacks_.push_back(std::move(on_disconnect_complete));
    } else {
      std::move(on_disconnect_complete).Run();
    }
  }
  DelayedReconnect();
}

void PosixSystemProducer::OnConnect() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (!PerfettoTracedProcess::Get()->CanStartTracing(this,
                                                     base::OnceClosure())) {
    // We are succesfully connected, but we can't register the data sources
    // right now, so move into "kUnregistered".
    state_ = State::kUnregistered;
    DisconnectWithReply();
    return;
  }
  state_ = State::kConnected;
  connection_backoff_ms_ = kInitialConnectionBackoffMs;
  for (const auto* const data_source :
       PerfettoTracedProcess::Get()->data_sources()) {
    NewDataSourceAdded(data_source);
  }
}

void PosixSystemProducer::OnDisconnect() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  DCHECK(GetService());
  // Currently our data sources don't support the concept of the service
  // disappearing and thus can't shut down cleanly (they would attempt to flush
  // data across the broken socket). Add a CHECK to catch this if its a problem.
  //
  // TODO(nuskos): Fix this, make it so we cleanly shut down on IPC errors.
  CHECK(!IsTracingActive());
  // This PostTask is needed because we want to clean up the state AFTER the
  // |ProducerEndpoint| has finished cleaning up.
  task_runner()->GetOrCreateTaskRunner()->PostTask(
      FROM_HERE,
      base::BindOnce(
          [](base::WeakPtr<PosixSystemProducer> weak_ptr) {
            if (!weak_ptr) {
              return;
            }
            if (weak_ptr->state_ == State::kConnecting) {
              base::AutoLock lock(weak_ptr->lock_);
              // We never connected, which means this disconnect is
              // an error from connecting, which means we don't need
              // to keep this endpoint (and associated memory around
              // forever) this prevents the memory leak from getting
              // excessive.
              weak_ptr->services_.erase(weak_ptr->services_.end() - 1);
            }
            weak_ptr->state_ = State::kDisconnected;
            weak_ptr->DelayedReconnect();
          },
          weak_ptr_factory_.GetWeakPtr()));
}

void PosixSystemProducer::OnTracingSetup() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  // Called by the IPC layer when tracing is first started and after shared
  // memory is set up.
}

void PosixSystemProducer::SetupDataSource(perfetto::DataSourceInstanceID,
                                          const perfetto::DataSourceConfig&) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  // Always called before StartDataSource but not used for any setup currently.
}

void PosixSystemProducer::StartDataSource(
    perfetto::DataSourceInstanceID id,
    const perfetto::DataSourceConfig& config) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (state_ != State::kConnected) {
    // Because StartDataSource is async, its possible a previous StartDataSource
    // is still in the PostTask queue, so we just ignore it here (We'll get
    // a new one if/when we re-register the DataSource once we've moved into
    // kConnected).
    return;
  }

  for (auto* const data_source : PerfettoTracedProcess::Get()->data_sources()) {
    if (data_source->name() == config.name()) {
      auto can_trace = PerfettoTracedProcess::Get()->CanStartTracing(
          this,
          base::BindOnce(
              [](base::WeakPtr<PosixSystemProducer> weak_ptr,
                 PerfettoTracedProcess::DataSourceBase* data_source,
                 perfetto::DataSourceInstanceID id,
                 const perfetto::DataSourceConfig& data_source_config) {
                if (!weak_ptr) {
                  return;
                }
                DCHECK_CALLED_ON_VALID_SEQUENCE(weak_ptr->sequence_checker_);
                {
                  base::AutoLock lock(weak_ptr->lock_);
                  ++weak_ptr->data_sources_tracing_;
                }
                data_source->StartTracingWithID(
                    id, weak_ptr.get(),
                    EnsureGuardRailsAreFollowed(data_source_config));
                weak_ptr->GetService()->NotifyDataSourceStarted(id);
              },
              weak_ptr_factory_.GetWeakPtr(), data_source, id, config));
      if (!can_trace) {
        DisconnectWithReply(base::OnceClosure());
      }
      return;
    }
  }
}

void PosixSystemProducer::StopDataSource(perfetto::DataSourceInstanceID id) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  for (auto* const data_source : PerfettoTracedProcess::Get()->data_sources()) {
    if (data_source->data_source_id() == id &&
        data_source->producer() == this) {
      data_source->StopTracing(base::BindOnce(
          [](base::WeakPtr<PosixSystemProducer> weak_ptr,
             perfetto::DataSourceInstanceID id) {
            if (!weak_ptr) {
              return;
            }
            DCHECK_CALLED_ON_VALID_SEQUENCE(weak_ptr->sequence_checker_);
            weak_ptr->GetService()->NotifyDataSourceStopped(id);
            {
              base::AutoLock lock(weak_ptr->lock_);
              --weak_ptr->data_sources_tracing_;
            }
            if (!weak_ptr->IsTracingActive()) {
              // If this is the last data source to be shut down then
              // perhaps we need to invoke any callbacks that were stored
              // (there might be none).
              weak_ptr->InvokeStoredOnDisconnectCallbacks();
            }
          },
          weak_ptr_factory_.GetWeakPtr(), id));
      return;
    }
  }
}

void PosixSystemProducer::Flush(
    perfetto::FlushRequestID id,
    const perfetto::DataSourceInstanceID* data_source_ids,
    size_t num_data_sources) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  pending_replies_for_latest_flush_ = {id, num_data_sources};
  for (auto* const data_source : PerfettoTracedProcess::Get()->data_sources()) {
    if (std::find(data_source_ids, data_source_ids + num_data_sources,
                  data_source->data_source_id()) !=
        data_source_ids + num_data_sources) {
      data_source->Flush(base::BindRepeating(
          [](base::WeakPtr<PosixSystemProducer> weak_ptr,
             perfetto::FlushRequestID flush_id) {
            if (weak_ptr) {
              weak_ptr->NotifyDataSourceFlushComplete(flush_id);
            }
          },
          weak_ptr_factory_.GetWeakPtr(), id));
    }
  }
}

void PosixSystemProducer::ClearIncrementalState(
    const perfetto::DataSourceInstanceID* data_source_ids,
    size_t num_data_sources) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  DCHECK(data_source_ids);
  DCHECK_GT(num_data_sources, 0u);
  std::unordered_set<perfetto::DataSourceInstanceID> to_clear{
      data_source_ids, data_source_ids + num_data_sources};
  for (auto* data_source : PerfettoTracedProcess::Get()->data_sources()) {
    if (to_clear.find(data_source->data_source_id()) != to_clear.end()) {
      data_source->ClearIncrementalState();
    }
  }
}

bool PosixSystemProducer::SetupSharedMemoryForStartupTracing() {
  // TODO(eseckler): Support startup tracing using an unbound SMA.
  NOTIMPLEMENTED();
  return false;
}

void PosixSystemProducer::ConnectSocket() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  state_ = State::kConnecting;
  const char* host_package_name = nullptr;
#if defined(OS_ANDROID)
  host_package_name =
      base::android::BuildInfo::GetInstance()->host_package_name();
#endif  // defined(OS_ANDROID)

  // On android we want to include if this is webview inside of an app or
  // Android Chrome. To aid this we add the host_package_name to differentiate
  // the various apps and sources.
  std::string producer_name;
  if (host_package_name) {
    producer_name = base::StrCat(
        {mojom::kPerfettoProducerNamePrefix, host_package_name, "-",
         base::NumberToString(
             base::trace_event::TraceLog::GetInstance()->process_id())});
  } else {
    producer_name = base::StrCat(
        {mojom::kPerfettoProducerNamePrefix,
         base::NumberToString(
             base::trace_event::TraceLog::GetInstance()->process_id())});
  }

  auto service = perfetto::ProducerIPCClient::Connect(
      socket_name_.c_str(), this, std::move(producer_name), task_runner(),
      perfetto::TracingService::ProducerSMBScrapingMode::kEnabled);

  base::AutoLock lock(lock_);
  services_.push_back(std::move(service));
}

bool PosixSystemProducer::SkipIfOnAndroidAndPreAndroidPie() const {
#if defined(OS_ANDROID)
  return disallow_pre_android_pie_ &&
         base::android::BuildInfo::GetInstance()->sdk_int() <
             base::android::SDK_VERSION_P;
#endif  // defined(OS_ANDROID)
  return false;
}

void PosixSystemProducer::InvokeStoredOnDisconnectCallbacks() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  for (auto& callback : on_disconnect_callbacks_) {
    DCHECK(!callback.is_null());
    std::move(callback).Run();
  }
  on_disconnect_callbacks_.clear();
}

void PosixSystemProducer::Connect() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (SkipIfOnAndroidAndPreAndroidPie()) {
    return;
  }
  switch (state_) {
    case State::kDisconnected:
      ConnectSocket();
      break;
    case State::kConnecting:
    case State::kConnected:
      // We are already connected (in which case do nothing). Or we're
      // currently connecting the socket and waiting for the OnConnect call
      // from the service.
      return;
    case State::kUnregistered:
      DCHECK(GetService());
      // We unregistered all our data sources due to a concurrent tracing
      // session but still have an open connection so just reregister
      // everything.
      OnConnect();
      break;
  }
}

void PosixSystemProducer::DelayedReconnect() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (SkipIfOnAndroidAndPreAndroidPie()) {
    return;
  }
  if (retrying_) {
    return;
  }
  retrying_ = true;

  task_runner()->GetOrCreateTaskRunner()->PostDelayedTask(
      FROM_HERE,
      base::BindOnce(
          [](base::WeakPtr<PosixSystemProducer> weak_ptr) {
            if (!weak_ptr) {
              return;
            }
            weak_ptr->retrying_ = false;
            if (PerfettoTracedProcess::Get()->CanStartTracing(
                    weak_ptr.get(), base::OnceClosure())) {
              weak_ptr->Connect();
            } else {
              weak_ptr->DelayedReconnect();
            }
          },
          weak_ptr_factory_.GetWeakPtr()),
      base::TimeDelta::FromMilliseconds(connection_backoff_ms_));

  connection_backoff_ms_ =
      IncreaseBackoff(connection_backoff_ms_, kMaxConnectionBackoffMs);
}

void PosixSystemProducer::NotifyDataSourceFlushComplete(
    perfetto::FlushRequestID id) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (pending_replies_for_latest_flush_.first != id) {
    // Ignore; completed flush was for an earlier request.
    return;
  }

  DCHECK_NE(pending_replies_for_latest_flush_.second, 0u);
  if (--pending_replies_for_latest_flush_.second == 0) {
    DCHECK(MaybeSharedMemoryArbiter());
    MaybeSharedMemoryArbiter()->NotifyFlushComplete(id);
  }
}

perfetto::TracingService::ProducerEndpoint* PosixSystemProducer::GetService() {
#if DCHECK_IS_ON()
  // Requires lock to be held when called on a non-producer sequence/thread.
  if (!sequence_checker_.CalledOnValidSequence()) {
    lock_.AssertAcquired();
  }
#endif  // DCHECK_IS_ON()

  switch (state_) {
    case State::kConnecting:
    case State::kConnected:
    case State::kUnregistered:
      return services_.back().get();
    default:
      return nullptr;
  }
}

}  // namespace tracing
