Statistical stack profiler for Windows x64

base/profiler/stack_sampling_profiler_win.cc

Index: base/profiler/stack_sampling_profiler_win.cc
diff --git a/base/profiler/stack_sampling_profiler_win.cc b/base/profiler/stack_sampling_profiler_win.cc
new file mode 100644
index 0000000000000000000000000000000000000000..066798dc6faa481438584c3c197dc7a34e40fd1a
--- /dev/null
+++ b/base/profiler/stack_sampling_profiler_win.cc
@@ -0,0 +1,270 @@
+// Copyright 2015 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/profiler/stack_sampling_profiler.h"
+
+#include <dbghelp.h>
+#include <map>
+#include <utility>
+#include <windows.h>
+
+#include "base/logging.h"
+#include "base/time/time.h"
+#include "base/win/pe_image.h"
+#include "base/win/scoped_handle.h"
+
+namespace base {
+
+namespace {
+
+class NativeStackSamplerWin : public StackSamplingProfiler::NativeStackSampler {
+ public:
+  explicit NativeStackSamplerWin(win::ScopedHandle thread_handle);
+  ~NativeStackSamplerWin() override;
+
+  // StackSamplingProfiler::NativeStackSampler:
+  void ProfileRecordingStarting(
+      StackSamplingProfiler::Profile* profile) override;
+  void RecordStackSample(StackSamplingProfiler::Sample* sample) override;
+  void ProfileRecordingStopped() override;
+
+ private:
+  static bool GetModuleInfo(HMODULE module,
+                            StackSamplingProfiler::Module* module_info);
+
+  void CopyToSample(const void* const instruction_pointers[],
+                    const HMODULE modules[],
+                    int stack_depth,
+                    StackSamplingProfiler::Sample* sample,
+                    std::vector<StackSamplingProfiler::Module>* module_infos);
+
+  win::ScopedHandle thread_handle_;
+  // Weak. Points to the profile being recorded between
+  // ProfileRecordingStarting() and ProfileRecordingStopped().
+  StackSamplingProfiler::Profile* current_profile_;
+  // Maps a module to the module's index within current_profile_->modules.
+  std::map<HMODULE, int> profile_module_index_;
+
+  DISALLOW_COPY_AND_ASSIGN(NativeStackSamplerWin);
+};
+
+// Walk the stack represented by |context| from the current frame downwards,
+// recording the instruction pointers for each frame in |instruction_pointers|.
+int RecordStack(CONTEXT* context,
+                int max_stack_size,
+                const void* instruction_pointers[],
+                bool* last_frame_is_unknown_function) {
+#ifdef _WIN64
+  *last_frame_is_unknown_function = false;
+
+  IMAGEHLP_SYMBOL64 sym;
+  sym.SizeOfStruct = sizeof(sym);
+  sym.MaxNameLength = 0;
+
+  for (int i = 0; i < max_stack_size; ++i) {
+    // Try to look up unwind metadata for the current function.
+    ULONG64 image_base;
+    PRUNTIME_FUNCTION runtime_function =
+        RtlLookupFunctionEntry(context->Rip, &image_base, nullptr);
+
+    instruction_pointers[i] = reinterpret_cast<void*>(context->Rip);
+
+    if (runtime_function) {
+      KNONVOLATILE_CONTEXT_POINTERS nvcontext = {0};
+      void* handler_data;
+      ULONG64 establisher_frame;
+      RtlVirtualUnwind(0, image_base, context->Rip, runtime_function, context,

[zturner, 2015/03/18 23:04:37]

Did you consider using RtlCaptureStackBacktrace?  It's a function which is
implemented on top of RtlVirtualUnwind, with its only limitation being that it
does not return the non-volatile register contexts for each frame.  You don't
appear to use any of this data, so you could simplify this code quite a bit by
just using RtlCaptureStackBackTrace.

[danduong, 2015/03/18 23:20:48]

As far as I know, RtlCaptureStackBackTrace cannot grab a stacktrace of another
thread.

+          &handler_data, &establisher_frame, &nvcontext);
+    } else {
+      // If we don't have a RUNTIME_FUNCTION, then we've encountered
+      // a leaf function.  Adjust the stack appropriately.
+      context->Rip = *reinterpret_cast<PDWORD64>(context->Rsp);
+      context->Rsp += 8;
+      *last_frame_is_unknown_function = true;
+    }
+
+    if (!context->Rip)
+      return i;
+  }
+  return max_stack_size;
+#else
+  return 0;

[zturner, 2015/03/18 23:04:37]

Is there any desire to support x86?  RtlVirtualUnwind isn't supported on x86,
but StackWalk64 is, and you can use StackWalk64 to walk yourself.  You need to
use RtlCaptureContext() to get a CONTEXT for yourself, then fill out the fields
of a STACKFRAME64 with the values from the CONTEXT and call it with
GetCurrentProcess() and GetCurrentThread().  I'm not sure how well it works
without symbols, but it might be worth an experiment?  Unless you don't care
about x86 by design, in which case ignore this.

[danduong, 2015/03/18 23:20:47]

We'll focus on x64 for now. StackWalk64, I've heard, can deadlock if you're
using it after SuspendThread.  It also tries to resolve symbols which is
unnecessary overhead.

+#endif
+}
+
+// Fills in |modules| corresponding to the pointers to code in |addresses|.
+void FindModulesForAddresses(const void* const addresses[], HMODULE modules[],
+                             int stack_depth,
+                             bool last_frame_is_unknown_function) {
+  const int module_frames = last_frame_is_unknown_function ? stack_depth - 1 :
+      stack_depth;
+  for (int i = 0; i < module_frames; ++i) {
+    HMODULE module = NULL;
+    if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+                          GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,

[zturner, 2015/03/18 23:04:37]

Why not increment the module's refcount?  You store the handle, so it seems
risky to not increment the refcount.

[Mike Wittman, 2015/03/19 00:15:44]

Incrementing the refcount led to deadlock previously when this was called
between SuspendThread/ResumeThread. Now that it's outside of those calls it
appears to work properly.

+                          reinterpret_cast<LPCTSTR>(addresses[i]),
+                          &module)) {
+      // HMODULE is the base address of the module.
+      DCHECK_LT(reinterpret_cast<const void*>(module), addresses[i]);
+      modules[i] = module;
+    }
+  }
+  for (int i = module_frames; i < stack_depth; ++i)

[zturner, 2015/03/18 23:04:37]

Alternatively, you could memset() the entire array to 0 upon startup.  Seems a
little clearer what the intent is in that case.

[Mike Wittman, 2015/03/19 00:15:44]

Used a zero initializer.

+    modules[stack_depth - 1] = NULL;

[zturner, 2015/03/18 23:04:37]

What about addresses?  You're going to end up with junk in the remainder of the
addresses[] array.

[Mike Wittman, 2015/03/19 00:15:44]

Done.

+}
+
+// Suspends the thread with |thread_handle|, records the stack into
+// |instruction_pointers| and the corresponding modules into |modules|, then
+// resumes the thread. Returns the size of the stack.
+int SuspendThreadAndRecordStack(HANDLE thread_handle, int max_stack_size,
+                                const void* instruction_pointers[],
+                                HMODULE modules[]) {
+  if (::SuspendThread(thread_handle) == -1) {

[zturner, 2015/03/18 23:04:37]

Just to be certain, this will never be equal to the current thread right?

[Mike Wittman, 2015/03/19 00:15:43]

Yes, this is only ever called on the profiler thread.

+    LOG(ERROR) << "SuspendThread failed: " << GetLastError();
+    return 0;
+  }
+
+  CONTEXT thread_context = {0};
+  thread_context.ContextFlags = CONTEXT_ALL;

[zturner, 2015/03/18 23:04:37]

Do you actually need CONTEXT_ALL?  CONTEXT_FULL is (oddly) more lightweight. 
CONTEXT_ALL will get debug registers, floating point registers, and extended
registers like the MMX registers.

[Mike Wittman, 2015/03/19 00:15:44]

CONTEXT_FULL appears to be sufficient.

+  if (!::GetThreadContext(thread_handle, &thread_context)) {
+    LOG(ERROR) << "GetThreadContext failed: " << GetLastError();
+  }
+
+  bool last_frame_is_unknown_function = false;
+  int stack_depth = RecordStack(&thread_context, max_stack_size,
+                                instruction_pointers,
+                                &last_frame_is_unknown_function);
+
+  if (::ResumeThread(thread_handle) == -1)
+    LOG(ERROR) << "ResumeThread failed: " << GetLastError();
+
+  FindModulesForAddresses(instruction_pointers, modules, stack_depth,
+                          last_frame_is_unknown_function);
+
+  return stack_depth;
+}
+
+}  // namespace
+
+scoped_ptr<StackSamplingProfiler::NativeStackSampler>
+StackSamplingProfiler::NativeStackSampler::Create(PlatformThreadId thread_id) {
+#if _WIN64
+  // Get the thread's handle.
+  HANDLE thread_handle = ::OpenThread(THREAD_ALL_ACCESS, FALSE, thread_id);

[zturner, 2015/03/18 23:04:37]

THREAD_GET_CONTEXT | THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION should be
all you need.  I'm not sure which is less likely to fail in practice when you're
talking about a billion users.  But I feel like asking for only as much
privilege as you need is better, even if we do control the thread.

[Mike Wittman, 2015/03/19 00:15:44]

Done.

+  DCHECK(thread_handle) << "OpenThread failed";
+
+  return scoped_ptr<NativeStackSampler>(new NativeStackSamplerWin(
+      win::ScopedHandle(thread_handle)));
+#else
+  return scoped_ptr<NativeStackSampler>();
+#endif
+}
+
+NativeStackSamplerWin::NativeStackSamplerWin(win::ScopedHandle thread_handle)
+    : thread_handle_(thread_handle.Take()) {
+#ifdef _WIN64
+  if (RtlVirtualUnwind == nullptr && RtlLookupFunctionEntry == nullptr) {
+    const HMODULE nt_dll_handle = ::GetModuleHandle(L"ntdll.dll");

[zturner, 2015/03/18 23:04:37]

You should probably handle the case where nt_dll_handle == nullptr.  It
shouldn't happen in practice, but still should handle it.

[Mike Wittman, 2015/03/19 00:15:44]

Done.

+    reinterpret_cast<void*&>(RtlVirtualUnwind) =
+        ::GetProcAddress(nt_dll_handle, "RtlVirtualUnwind");
+    reinterpret_cast<void*&>(RtlLookupFunctionEntry) =
+        ::GetProcAddress(nt_dll_handle, "RtlLookupFunctionEntry");
+  }
+#endif
+}
+
+NativeStackSamplerWin::~NativeStackSamplerWin() {
+}
+
+void NativeStackSamplerWin::ProfileRecordingStarting(
+    StackSamplingProfiler::Profile* profile) {
+  current_profile_ = profile;
+  profile_module_index_.clear();
+}
+
+void NativeStackSamplerWin::RecordStackSample(
+    StackSamplingProfiler::Sample* sample) {
+  DCHECK(current_profile_);
+
+  const int max_stack_size = 64;
+  const void* instruction_pointers[max_stack_size];
+  HMODULE modules[max_stack_size];

[zturner, 2015/03/18 23:04:37]

If you initialize these two arrays with "= {0}" then you can skip the loop in
the SuspendThreadAndRecordStack function and also fix the issue I pointed out
with that function earlier where instruction_pointers[] ends up with junk.

[Mike Wittman, 2015/03/19 00:15:44]

Done.

+
+  int stack_depth = SuspendThreadAndRecordStack(
+      thread_handle_.Get(), max_stack_size, instruction_pointers, modules);
+  CopyToSample(instruction_pointers, modules, stack_depth, sample,
+               &current_profile_->modules);
+}
+
+void NativeStackSamplerWin::ProfileRecordingStopped() {
+  current_profile_ = nullptr;
+}
+
+// static
+bool NativeStackSamplerWin::GetModuleInfo(
+    HMODULE module,
+    StackSamplingProfiler::Module* module_info) {
+  wchar_t module_name[MAX_PATH];
+  DWORD result_length =
+      GetModuleFileName(module, module_name, arraysize(module_name));
+  if (result_length == 0)
+    return false;
+
+  module_info->filename = base::FilePath(module_name);
+
+  module_info->base_address = reinterpret_cast<const void*>(module);
+
+  GUID guid;
+  DWORD age;
+  win::PEImage(module).GetDebugId(&guid, &age);

[zturner, 2015/03/18 23:04:37]

Is this present for release / optimized builds?

[danduong, 2015/03/18 23:20:47]

Should be. Minidumps include the debugIds of the modules.

+  module_info->id.insert(module_info->id.end(),
+                         reinterpret_cast<char*>(&guid),
+                         reinterpret_cast<char*>(&guid + 1));
+  module_info->id.insert(module_info->id.end(),
+                         reinterpret_cast<char*>(&age),
+                         reinterpret_cast<char*>(&age + 1));
+
+  return true;
+}
+
+void NativeStackSamplerWin::CopyToSample(
+    const void* const instruction_pointers[],
+    const HMODULE modules[],
+    int stack_depth,
+    StackSamplingProfiler::Sample* sample,
+    std::vector<StackSamplingProfiler::Module>* module_infos) {
+  sample->clear();
+  sample->reserve(stack_depth);
+
+  for (int i = 0; i < stack_depth; ++i) {
+    sample->push_back(StackSamplingProfiler::Frame());
+    StackSamplingProfiler::Frame& frame = sample->back();
+
+    frame.instruction_pointer = instruction_pointers[i];
+
+    // Record an invalid module index if we don't have a valid module.
+    if (!modules[i]) {
+      frame.module_index = -1;
+      continue;
+    }
+
+    auto loc = profile_module_index_.find(modules[i]);
+    if (loc == profile_module_index_.end()) {
+      StackSamplingProfiler::Module module_info;
+      // Record an invalid module index if we have a module but can't find
+      // information on it.
+      if (!GetModuleInfo(modules[i], &module_info)) {
+        frame.module_index = -1;
+        continue;
+      }
+      module_infos->push_back(module_info);
+      loc = profile_module_index_.insert(std::make_pair(
+          modules[i], static_cast<int>(module_infos->size() - 1))).first;
+    }
+
+    frame.module_index = loc->second;
+  }
+}
+
+}  // namespace base