/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "mozilla/ProfilerThreadRegistrationData.h"

#include "mozilla/CycleCollectedJSContext.h"
#include "mozilla/FlowMarkers.h"
#include "mozilla/ProfilerMarkers.h"
#include "js/AllocationRecording.h"
#include "js/ProfilingStack.h"

#if defined(XP_WIN)
#  include <windows.h>
#elif defined(XP_DARWIN)
#  include <pthread.h>
#elif defined(XP_LINUX) && !defined(ANDROID)
#  include "mozilla/ScopeExit.h"
#  include <pthread.h>
#endif

#ifdef NIGHTLY_BUILD
namespace geckoprofiler::markers {

using namespace mozilla;

struct ThreadCpuUseMarker {
  static constexpr Span<const char> MarkerTypeName() {
    return MakeStringSpan("ThreadCpuUse");
  }
  static void StreamJSONMarkerData(baseprofiler::SpliceableJSONWriter& aWriter,
                                   ProfilerThreadId aThreadId,
                                   int64_t aCpuTimeMs, int64_t aWakeUps,
                                   const ProfilerString8View& aThreadName) {
    aWriter.IntProperty("threadId", static_cast<int64_t>(aThreadId.ToNumber()));
    aWriter.IntProperty("time", aCpuTimeMs);
    aWriter.IntProperty("wakeups", aWakeUps);
    aWriter.StringProperty("label", aThreadName);
  }
  static MarkerSchema MarkerTypeDisplay() {
    using MS = MarkerSchema;
    MS schema{MS::Location::MarkerChart, MS::Location::MarkerTable};
    schema.AddKeyLabelFormat("time", "CPU Time", MS::Format::Milliseconds);
    schema.AddKeyLabelFormat("wakeups", "Wake ups", MS::Format::Integer);
    schema.AddKeyFormat("label", MS::Format::String, MS::PayloadFlags::Hidden);
    schema.SetTooltipLabel("{marker.name} - {marker.data.label}");
    schema.SetTableLabel(
        "{marker.data.label}: {marker.data.time} of CPU time, "
        "{marker.data.wakeups} wake ups");
    return schema;
  }
};

}  // namespace geckoprofiler::markers
#endif

namespace mozilla::profiler {

#if defined(XP_LINUX) && !defined(ANDROID)
static const void* pthread_get_stacktop_linux(const void* aStackTop) {
  pthread_attr_t attr;
  if (pthread_getattr_np(pthread_self(), &attr) != 0) {
    return aStackTop;
  }
  auto attrGuard = MakeScopeExit([&]() { pthread_attr_destroy(&attr); });
  void* stackBase = nullptr;
  size_t stackSize = 0;
  if (pthread_attr_getstack(&attr, &stackBase, &stackSize) != 0 ||
      !(stackBase && stackSize > 0)) {
    return aStackTop;
  }
  // > The base (lowest addressable byte) of the storage shall be
  // > stackaddr, and the size of the storage shall be stacksize
  // > bytes.
  // <https://www.man7.org/linux/man-pages/man3/pthread_attr_getstack.3p.html>
  return static_cast<const char*>(stackBase) + stackSize;
}
#endif

ThreadRegistrationData::ThreadRegistrationData(const char* aName,
                                               const void* aStackTop)
    : mInfo(aName),
      mPlatformData(mInfo.ThreadId()),
      mStackTop(
#if defined(XP_WIN)
          // We don't have to guess on Windows.
          reinterpret_cast<const void*>(
              reinterpret_cast<PNT_TIB>(NtCurrentTeb())->StackBase)
#elif defined(XP_DARWIN)
          // We don't have to guess on Mac/Darwin.
          reinterpret_cast<const void*>(
              pthread_get_stackaddr_np(pthread_self()))
#elif defined(XP_LINUX) && !defined(ANDROID)
          // We don't have to guess on non-Android Linux.
          pthread_get_stacktop_linux(aStackTop)
#else
          // Otherwise use the given guess.
          aStackTop
#endif
      ) {
}

// This is a simplified version of profiler_add_marker that can be easily passed
// into the JS engine.
static void profiler_add_js_marker(mozilla::MarkerCategory aCategory,
                                   const char* aMarkerName,
                                   const char* aMarkerText) {
  AUTO_PROFILER_STATS(js_marker);
  profiler_add_marker(
      mozilla::ProfilerString8View::WrapNullTerminatedString(aMarkerName),
      aCategory, {}, ::geckoprofiler::markers::TextMarker{},
      mozilla::ProfilerString8View::WrapNullTerminatedString(aMarkerText));
}

static void profiler_add_js_interval(mozilla::MarkerCategory aCategory,
                                     const char* aMarkerName,
                                     mozilla::TimeStamp aStartTime,
                                     const char* aMarkerText) {
  AUTO_PROFILER_STATS(js_interval);
  profiler_add_marker(
      mozilla::ProfilerString8View::WrapNullTerminatedString(aMarkerName),
      aCategory, mozilla::MarkerTiming::IntervalUntilNowFrom(aStartTime),
      ::geckoprofiler::markers::TextMarker{},
      mozilla::ProfilerString8View::WrapNullTerminatedString(aMarkerText));
}

static void profiler_add_js_flow(mozilla::MarkerCategory aCategory,
                                 const char* aMarkerName, uint64_t aFlowId) {
  if (!profiler_feature_active(ProfilerFeature::Flows)) {
    return;
  }
  AUTO_PROFILER_STATS(js_flow);
  profiler_add_marker(
      mozilla::ProfilerString8View::WrapNullTerminatedString(aMarkerName),
      aCategory, {}, ::geckoprofiler::markers::FlowMarker{},
      Flow::ProcessScoped(aFlowId));
}

static void profiler_add_js_terminating_flow(mozilla::MarkerCategory aCategory,
                                             const char* aMarkerName,
                                             uint64_t aFlowId) {
  if (!profiler_feature_active(ProfilerFeature::Flows)) {
    return;
  }
  AUTO_PROFILER_STATS(js_terminating_flow);
  profiler_add_marker(
      mozilla::ProfilerString8View::WrapNullTerminatedString(aMarkerName),
      aCategory, {}, ::geckoprofiler::markers::TerminatingFlowMarker{},
      Flow::ProcessScoped(aFlowId));
}

static void profiler_add_js_allocation_marker(JS::RecordAllocationInfo&& info) {
  if (!profiler_thread_is_being_profiled_for_markers()) {
    return;
  }

  struct JsAllocationMarker {
    static constexpr mozilla::Span<const char> MarkerTypeName() {
      return mozilla::MakeStringSpan("JS allocation");
    }
    static void StreamJSONMarkerData(
        mozilla::baseprofiler::SpliceableJSONWriter& aWriter,
        const mozilla::ProfilerString16View& aTypeName,
        const mozilla::ProfilerString8View& aClassName,
        const mozilla::ProfilerString16View& aDescriptiveTypeName,
        const mozilla::ProfilerString8View& aCoarseType, uint64_t aSize,
        bool aInNursery) {
      if (aClassName.Length() != 0) {
        aWriter.StringProperty("className", aClassName);
      }
      if (aTypeName.Length() != 0) {
        aWriter.StringProperty("typeName", NS_ConvertUTF16toUTF8(aTypeName));
      }
      if (aDescriptiveTypeName.Length() != 0) {
        aWriter.StringProperty("descriptiveTypeName",
                               NS_ConvertUTF16toUTF8(aDescriptiveTypeName));
      }
      aWriter.StringProperty("coarseType", aCoarseType);
      aWriter.IntProperty("size", aSize);
      aWriter.BoolProperty("inNursery", aInNursery);
    }
    static mozilla::MarkerSchema MarkerTypeDisplay() {
      return mozilla::MarkerSchema::SpecialFrontendLocation{};
    }
  };

  profiler_add_marker(
      "JS allocation", geckoprofiler::category::JS,
      mozilla::MarkerStack::Capture(), JsAllocationMarker{},
      mozilla::ProfilerString16View::WrapNullTerminatedString(info.typeName),
      mozilla::ProfilerString8View::WrapNullTerminatedString(info.className),
      mozilla::ProfilerString16View::WrapNullTerminatedString(
          info.descriptiveTypeName),
      mozilla::ProfilerString8View::WrapNullTerminatedString(info.coarseType),
      info.size, info.inNursery);
}

JSContext* ThreadRegistrationUnlockedReaderAndAtomicRWOnThread::GetJSContext()
    const {
  if (!mCCJSContext) {
    return nullptr;
  }
  return mCCJSContext->Context();
}

void ThreadRegistrationLockedRWFromAnyThread::SetProfilingFeaturesAndData(
    ThreadProfilingFeatures aProfilingFeatures,
    ProfiledThreadData* aProfiledThreadData, const PSAutoLock&) {
  MOZ_ASSERT(mProfilingFeatures == ThreadProfilingFeatures::NotProfiled);
  mProfilingFeatures = aProfilingFeatures;

  MOZ_ASSERT(!mProfiledThreadData);
  MOZ_ASSERT(aProfiledThreadData);
  mProfiledThreadData = aProfiledThreadData;

  if (mCCJSContext) {
    // The thread is now being profiled, and we already have a JSContext,
    // allocate a JsFramesBuffer to allow profiler-unlocked on-thread sampling.
    MOZ_ASSERT(!mJsFrameBuffer);
    mJsFrameBuffer = new JsFrame[MAX_JS_FRAMES];
  }

  // Check invariants.
  MOZ_ASSERT((mProfilingFeatures != ThreadProfilingFeatures::NotProfiled) ==
             !!mProfiledThreadData);
  MOZ_ASSERT((mCCJSContext &&
              (mProfilingFeatures != ThreadProfilingFeatures::NotProfiled)) ==
             !!mJsFrameBuffer);
}

void ThreadRegistrationLockedRWFromAnyThread::ClearProfilingFeaturesAndData(
    const PSAutoLock&) {
  mProfilingFeatures = ThreadProfilingFeatures::NotProfiled;
  mProfiledThreadData = nullptr;

  if (mJsFrameBuffer) {
    delete[] mJsFrameBuffer;
    mJsFrameBuffer = nullptr;
  }

  // Check invariants.
  MOZ_ASSERT((mProfilingFeatures != ThreadProfilingFeatures::NotProfiled) ==
             !!mProfiledThreadData);
  MOZ_ASSERT((mCCJSContext &&
              (mProfilingFeatures != ThreadProfilingFeatures::NotProfiled)) ==
             !!mJsFrameBuffer);
}

void ThreadRegistrationLockedRWOnThread::SetCycleCollectedJSContext(
    CycleCollectedJSContext* aCx) {
  MOZ_ASSERT(aCx && !mCCJSContext);
  MOZ_ASSERT(aCx->Context());

  mCCJSContext = aCx;

  if (mProfiledThreadData) {
    MOZ_ASSERT((mProfilingFeatures != ThreadProfilingFeatures::NotProfiled) ==
               !!mProfiledThreadData);
    // We now have a JSContext, and the thread is already being profiled,
    // allocate a JsFramesBuffer to allow profiler-unlocked on-thread sampling.
    MOZ_ASSERT(!mJsFrameBuffer);
    mJsFrameBuffer = new JsFrame[MAX_JS_FRAMES];
  }

  // We give the JS engine a non-owning reference to the ProfilingStack. It's
  // important that the JS engine doesn't touch this once the thread dies.
  js::SetContextProfilingStack(aCx->Context(), &ProfilingStackRef());

  // Check invariants.
  MOZ_ASSERT((mCCJSContext &&
              (mProfilingFeatures != ThreadProfilingFeatures::NotProfiled)) ==
             !!mJsFrameBuffer);
}

void ThreadRegistrationLockedRWOnThread::ClearCycleCollectedJSContext() {
  mCCJSContext = nullptr;

  if (mJsFrameBuffer) {
    delete[] mJsFrameBuffer;
    mJsFrameBuffer = nullptr;
  }

  // Check invariants.
  MOZ_ASSERT((mCCJSContext &&
              (mProfilingFeatures != ThreadProfilingFeatures::NotProfiled)) ==
             !!mJsFrameBuffer);
}

void ThreadRegistrationLockedRWOnThread::PollJSSampling() {
  // We can't start/stop profiling until we have the thread's JSContext.
  if (mCCJSContext) {
    // It is possible for mJSSampling to go through the following sequences.
    //
    // - INACTIVE, ACTIVE_REQUESTED, INACTIVE_REQUESTED, INACTIVE
    //
    // - ACTIVE, INACTIVE_REQUESTED, ACTIVE_REQUESTED, ACTIVE
    //
    // Therefore, the if and else branches here aren't always interleaved.
    // This is ok because the JS engine can handle that.
    //
    JSContext* cx = mCCJSContext->Context();
    if (mJSSampling == ACTIVE_REQUESTED) {
      mJSSampling = ACTIVE;
      js::EnableContextProfilingStack(cx, true);

      if (JSAllocationsEnabled()) {
        // TODO - This probability should not be hardcoded. See Bug 1547284.
        JS::EnableRecordingAllocations(cx, profiler_add_js_allocation_marker,
                                       0.01);
      }
      js::RegisterContextProfilerMarkers(
          cx, profiler_add_js_marker, profiler_add_js_interval,
          profiler_add_js_flow, profiler_add_js_terminating_flow);

    } else if (mJSSampling == INACTIVE_REQUESTED) {
      mJSSampling = INACTIVE;
      js::EnableContextProfilingStack(cx, false);

      if (JSAllocationsEnabled()) {
        JS::DisableRecordingAllocations(cx);
      }
    }
  }
}

#ifdef NIGHTLY_BUILD
bool ThreadRegistrationUnlockedConstReaderAndAtomicRW::RecordWakeCount(
    nsACString& aThreadName, uint64_t& aCpuTimeMs, uint64_t& aWakeCount) const {
  baseprofiler::detail::BaseProfilerAutoLock lock(mRecordWakeCountMutex);

  aWakeCount = mWakeCount - mAlreadyRecordedWakeCount;
  if (aWakeCount == 0 && mSleep != AWAKE) {
    // If no new wake-up was counted, and the thread is not marked awake,
    // we can be pretty sure there is no CPU activity to record.
    // Threads that are never annotated as asleep/awake (typically rust threads)
    // start as awake.
    return false;
  }

  uint64_t cpuTimeNs;
  if (!GetCpuTimeSinceThreadStartInNs(&cpuTimeNs, PlatformDataCRef())) {
    cpuTimeNs = 0;
  }

  constexpr uint64_t NS_PER_MS = 1'000'000;
  uint64_t cpuTimeMs = cpuTimeNs / NS_PER_MS;

  aCpuTimeMs = MOZ_LIKELY(cpuTimeMs > mAlreadyRecordedCpuTimeInMs)
                   ? cpuTimeMs - mAlreadyRecordedCpuTimeInMs
                   : 0;

  if (!aWakeCount && !aCpuTimeMs) {
    // Nothing to report, avoid computing the Glean friendly thread name.
    return false;
  }

  aThreadName.Assign(mInfo.Name());
  // Trim the trailing number of threads that are part of a thread pool.
  for (size_t length = aThreadName.Length(); length > 0; --length) {
    const char c = aThreadName.CharAt(length - 1);
    if ((c < '0' || c > '9') && c != '#' && c != ' ') {
      if (length != aThreadName.Length()) {
        aThreadName.SetLength(length);
      }
      break;
    }
  }

  // The thread id is provided as part of the payload because this call is
  // inside a ThreadRegistration data function, which could be invoked with
  // the ThreadRegistry locked. We cannot call any function/option that could
  // attempt to lock the ThreadRegistry again, like MarkerThreadId.
  PROFILER_MARKER("Thread CPU use", OTHER, {}, ThreadCpuUseMarker,
                  mInfo.ThreadId(), aCpuTimeMs, aWakeCount, aThreadName);
  mAlreadyRecordedCpuTimeInMs = cpuTimeMs;
  mAlreadyRecordedWakeCount += aWakeCount;

  return true;
}
#endif

}  // namespace mozilla::profiler