Windows-native sampling profiler wrapper class.

base/win/sampling_profiler.cc

Index: base/win/sampling_profiler.cc
diff --git a/base/win/sampling_profiler.cc b/base/win/sampling_profiler.cc
new file mode 100644
index 0000000000000000000000000000000000000000..268542a94fd8b57f6fa600b93879d579feb37d53
--- /dev/null
+++ b/base/win/sampling_profiler.cc
@@ -0,0 +1,232 @@
+// Copyright (c) 2011 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 "base/win/sampling_profiler.h"
+
+#include <winternl.h>  // for NTSTATUS.
+
+#include "base/lazy_instance.h"
+
+// Copied from wdm.h in the WDK as we don't want to take
+// a dependency on the WDK.
+typedef enum _KPROFILE_SOURCE {
+    ProfileTime,
+    ProfileAlignmentFixup,
+    ProfileTotalIssues,
+    ProfilePipelineDry,
+    ProfileLoadInstructions,
+    ProfilePipelineFrozen,
+    ProfileBranchInstructions,
+    ProfileTotalNonissues,
+    ProfileDcacheMisses,
+    ProfileIcacheMisses,
+    ProfileCacheMisses,
+    ProfileBranchMispredictions,
+    ProfileStoreInstructions,
+    ProfileFpInstructions,
+    ProfileIntegerInstructions,
+    Profile2Issue,
+    Profile3Issue,
+    Profile4Issue,
+    ProfileSpecialInstructions,
+    ProfileTotalCycles,
+    ProfileIcacheIssues,
+    ProfileDcacheAccesses,
+    ProfileMemoryBarrierCycles,
+    ProfileLoadLinkedIssues,
+    ProfileMaximum
+} KPROFILE_SOURCE;
+
+
+namespace {
+
+// Signatures for the native functions we need to access the sampling profiler.
+typedef NTSTATUS (NTAPI *ZwSetIntervalProfileFunc)(ULONG, KPROFILE_SOURCE);
+typedef NTSTATUS (NTAPI *ZwQueryIntervalProfileFunc)(KPROFILE_SOURCE, PULONG);
+
+typedef NTSTATUS (NTAPI *ZwCreateProfileFunc)(PHANDLE profile,
+                                              HANDLE process,
+                                              PVOID code_start,
+                                              ULONG code_size,
+                                              ULONG eip_bucket_shift,
+                                              PULONG buckets,
+                                              ULONG buckets_byte_size,
+                                              KPROFILE_SOURCE source,
+                                              DWORD_PTR processor_mask);
+
+typedef NTSTATUS (NTAPI *ZwStartProfileFunc)(HANDLE);
+typedef NTSTATUS (NTAPI *ZwStopProfileFunc)(HANDLE);
+
+// This class is used to lazy-initialize pointers to the native
+// functions we need to access.
+class ProfilerFuncs {
+ public:
+  ProfilerFuncs();
+
+  ZwSetIntervalProfileFunc ZwSetIntervalProfile;
+  ZwQueryIntervalProfileFunc ZwQueryIntervalProfile;
+  ZwCreateProfileFunc ZwCreateProfile;
+  ZwStartProfileFunc ZwStartProfile;
+  ZwStopProfileFunc ZwStopProfile;
+
+  // True iff all of the function pointers above were successfully initialized.
+  bool initialized_;
+};
+
+ProfilerFuncs::ProfilerFuncs() : initialized_(false) {
+  HMODULE ntdll = ::GetModuleHandle(L"ntdll.dll");
+  if (ntdll != NULL) {
+    ZwSetIntervalProfile = reinterpret_cast<ZwSetIntervalProfileFunc>(
+        ::GetProcAddress(ntdll, "ZwSetIntervalProfile"));
+    ZwQueryIntervalProfile = reinterpret_cast<ZwQueryIntervalProfileFunc>(
+        ::GetProcAddress(ntdll, "ZwQueryIntervalProfile"));
+    ZwCreateProfile = reinterpret_cast<ZwCreateProfileFunc>(
+        ::GetProcAddress(ntdll, "ZwCreateProfile"));
+    ZwStartProfile = reinterpret_cast<ZwStartProfileFunc>(
+        ::GetProcAddress(ntdll, "ZwStartProfile"));
+    ZwStopProfile = reinterpret_cast<ZwStopProfileFunc>(
+        ::GetProcAddress(ntdll, "ZwStopProfile"));
+
+    if (ZwSetIntervalProfile &&
+        ZwQueryIntervalProfile &&
+        ZwCreateProfile &&
+        ZwStartProfile &&
+        ZwStopProfile) {
+      initialized_ = true;
+    }
+  }
+}
+
+base::LazyInstance<ProfilerFuncs, base::LeakyLazyInstanceTraits<ProfilerFuncs>>
+    funcs = LAZY_INSTANCE_INITIALIZER;
+
+}  // namespace
+
+
+namespace base {
+namespace win {
+
+SamplingProfiler::SamplingProfiler() : is_started_(false) {
+}
+
+SamplingProfiler::~SamplingProfiler() {
+  if (is_started_) {
+    CHECK(Stop()) <<
+        "Unable to stop sampling profiler, this will cause memory corruption.";
+  }
+}
+
+bool SamplingProfiler::Initialize(HANDLE process,
+                                  void* start,
+                                  size_t size,
+                                  size_t log2_bucket_size) {
+  // You only get to initialize each instance once.
+  DCHECK(!profile_handle_.IsValid());
+  DCHECK(!is_started_);
+  DCHECK(start != NULL);
+  DCHECK_NE(0U, size);
+  DCHECK_LE(2, log2_bucket_size);
+  DCHECK_GE(32, log2_bucket_size);
+
+  // Bail if the native functions weren't found.
+  if (!funcs.Get().initialized_)
+    return false;
+
+  size_t bucket_size = 1 << log2_bucket_size;
+  size_t num_buckets = (size + bucket_size - 1) / bucket_size;
+  DCHECK(num_buckets != 0);
+  buckets_.resize(num_buckets);
+
+  // Get our affinity mask for the call below.
+  DWORD_PTR process_affinity = 0;
+  DWORD_PTR system_affinity = 0;
+  if (!::GetProcessAffinityMask(process, &process_affinity, &system_affinity)) {
+    LOG(ERROR) << "Failed to get process affinity mask.";
+    return false;
+  }
+
+  HANDLE profile = NULL;
+  NTSTATUS status =
+      funcs.Get().ZwCreateProfile(&profile,
+                                  process,
+                                  start,
+                                  size,
+                                  log2_bucket_size,
+                                  &buckets_[0],
+                                  sizeof(buckets_[0]) * num_buckets,
+                                  ProfileTime,
+                                  process_affinity);
+
+  if (!NT_SUCCESS(status)) {
+    // Might as well deallocate the buckets.
+    buckets_.resize(0);
+    LOG(ERROR) << "Failed to create profile, error 0x" << std::hex << status;
+    return false;
+  }
+
+  DCHECK(profile != NULL);
+  profile_handle_.Set(profile);
+
+  return true;
+}
+
+bool SamplingProfiler::Start() {
+  DCHECK(profile_handle_.IsValid());
+  DCHECK(!is_started_);
+  DCHECK(funcs.Get().initialized_);
+
+  NTSTATUS status = funcs.Get().ZwStartProfile(profile_handle_.Get());
+  if (!NT_SUCCESS(status))
+    return false;
+
+  is_started_ = true;
+
+  return true;
+}
+
+bool SamplingProfiler::Stop() {
+  DCHECK(profile_handle_.IsValid());
+  DCHECK(is_started_);
+  DCHECK(funcs.Get().initialized_);
+
+  NTSTATUS status = funcs.Get().ZwStopProfile(profile_handle_.Get());
+  if (!NT_SUCCESS(status))
+    return false;
+  is_started_ = false;
+
+  return true;
+}
+
+bool SamplingProfiler::SetSamplingInterval(base::TimeDelta sampling_interval) {
+  if (!funcs.Get().initialized_)
+    return false;
+
+  // According to Nebbet, the sampling interval is in units of 100ns.
+  ULONG interval = sampling_interval.InMicroseconds() * 10;
+  NTSTATUS status = funcs.Get().ZwSetIntervalProfile(interval, ProfileTime);
+  if (!NT_SUCCESS(status))
+    return false;
+
+  return true;
+}
+
+bool SamplingProfiler::GetSamplingInterval(base::TimeDelta* sampling_interval) {
+  DCHECK(sampling_interval != NULL);
+
+  if (!funcs.Get().initialized_)
+    return false;
+
+  ULONG interval = 0;
+  NTSTATUS status = funcs.Get().ZwQueryIntervalProfile(ProfileTime, &interval);
+  if (!NT_SUCCESS(status))
+    return false;
+
+  // According to Nebbet, the sampling interval is in units of 100ns.
+  *sampling_interval = base::TimeDelta::FromMicroseconds(interval / 10);
+
+  return true;
+}
+
+}  // namespace win
+}  // namespace base