StackSamplingProfiler clean up

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
index ba46cf076dbd01a26dcaf2d65df3e9c5830dd2dd..4d0f1d6992f21a8e2072cd32689edd5e181467a1 100644
--- a/base/profiler/stack_sampling_profiler_win.cc
+++ b/base/profiler/stack_sampling_profiler_win.cc
@@ -2,14 +2,13 @@
 // 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 <windows.h>
 
-#include <dbghelp.h>
 #include <map>
 #include <utility>
-#include <windows.h>
 
 #include "base/logging.h"
+#include "base/profiler/native_stack_sampler.h"
 #include "base/time/time.h"
 #include "base/win/pe_image.h"
 #include "base/win/scoped_handle.h"
@@ -18,38 +17,7 @@ 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,
+// Walks 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,
@@ -58,66 +26,65 @@ int RecordStack(CONTEXT* context,
 #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) {
+  int i = 0;
+  for (; (i < max_stack_size) && context->Rip; ++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);
+    instruction_pointers[i] = reinterpret_cast<const 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,
-          &handler_data, &establisher_frame, &nvcontext);
+                       &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.
+      // If we don't have a RUNTIME_FUNCTION, then we've encountered a leaf
+      // function.  Adjust the stack appropriately prior to the next function
+      // lookup.
       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;
+  return i;
 #else
   return 0;
 #endif
 }
 
-// Fills in |modules| corresponding to the pointers to code in |addresses|. The
-// modules are returned with reference counts incremented should be freed with
-// FreeModules.
-void FindModulesForAddresses(const void* const addresses[], HMODULE modules[],
-                             int stack_depth,
+// Fills in |module_handles| corresponding to the pointers to code in
+// |addresses|. The module handles are returned with reference counts
+// incremented and should be freed with FreeModuleHandles. See note in
+// SuspendThreadAndRecordStack for why |addresses| and |module_handles| are
+// arrays.
+void FindModuleHandlesForAddresses(const void* const addresses[],
+                             HMODULE module_handles[], int stack_depth,
                              bool last_frame_is_unknown_function) {
-  const int module_frames = last_frame_is_unknown_function ? stack_depth - 1 :
-      stack_depth;
+  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;
+    HMODULE module_handle = NULL;
     if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
                           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;
+                          &module_handle)) {
+      // HMODULE actually represents the base address of the module, so we can
+      // use it directly as an address.
+      DCHECK_LE(reinterpret_cast<const void*>(module_handle), addresses[i]);
+      module_handles[i] = module_handle;
     }
   }
 }
 
-// Free the modules returned by FindModulesForAddresses.
-void FreeModules(int stack_depth, HMODULE modules[]) {
+// Frees the modules handles returned by FindModuleHandlesForAddresses. See note
+// in SuspendThreadAndRecordStack for why |module_handles| is an array.
+void FreeModuleHandles(int stack_depth, HMODULE module_handles[]) {
   for (int i = 0; i < stack_depth; ++i) {
-    if (modules[i])
-      ::FreeLibrary(modules[i]);
+    if (module_handles[i])
+      ::FreeLibrary(module_handles[i]);
   }
 }
 
@@ -141,185 +108,212 @@ ScopedDisablePriorityBoost::ScopedDisablePriorityBoost(HANDLE thread_handle)
       boost_state_was_disabled_(false) {
   got_previous_boost_state_ =
       ::GetThreadPriorityBoost(thread_handle_, &boost_state_was_disabled_);
-  if (got_previous_boost_state_ && !boost_state_was_disabled_) {
-    // Confusingly, TRUE disables priority boost ...
+  if (got_previous_boost_state_) {
+    // Confusingly, TRUE disables priority boost.
     ::SetThreadPriorityBoost(thread_handle_, TRUE);
   }
 }
 
 ScopedDisablePriorityBoost::~ScopedDisablePriorityBoost() {
-  if (got_previous_boost_state_ && !boost_state_was_disabled_) {
-    // ... and FALSE enables priority boost.
-    ::SetThreadPriorityBoost(thread_handle_, FALSE);
-  }
+  if (got_previous_boost_state_)
+    ::SetThreadPriorityBoost(thread_handle_, boost_state_was_disabled_);
 }
 
 // Suspends the thread with |thread_handle|, records the stack into
 // |instruction_pointers|, then resumes the thread. Returns the size of the
 // stack.
+//
+// IMPORTANT NOTE: No heap allocations may occur between SuspendThread and
+// ResumeThread. Otherwise this code can deadlock on heap locks acquired by the
+// target thread before it was suspended. This is why we pass instruction
+// pointers and module handles as preallocated arrays rather than vectors, since
+// vectors make it too easy to subtly allocate memory.
 int SuspendThreadAndRecordStack(HANDLE thread_handle, int max_stack_size,
                                 const void* instruction_pointers[],
                                 bool* last_frame_is_unknown_function) {
-#if defined(_WIN64)
-  if (RtlVirtualUnwind == nullptr || RtlLookupFunctionEntry == nullptr)
+  if (::SuspendThread(thread_handle) == -1)
     return 0;
-#endif
-
-  if (::SuspendThread(thread_handle) == -1) {
-    LOG(ERROR) << "SuspendThread failed: " << GetLastError();
-    return 0;
-  }
 
+  int stack_depth = 0;
   CONTEXT thread_context = {0};
   thread_context.ContextFlags = CONTEXT_FULL;
-  if (!::GetThreadContext(thread_handle, &thread_context)) {
-    LOG(ERROR) << "GetThreadContext failed: " << GetLastError();
+  if (::GetThreadContext(thread_handle, &thread_context)) {
+    stack_depth = RecordStack(&thread_context, max_stack_size,
+                              instruction_pointers,
+                              last_frame_is_unknown_function);
   }
 
-  int stack_depth = RecordStack(&thread_context, max_stack_size,
-                                instruction_pointers,
-                                last_frame_is_unknown_function);
-
-  {
-    ScopedDisablePriorityBoost disable_priority_boost(thread_handle);
-    if (::ResumeThread(thread_handle) == -1)
-      LOG(ERROR) << "ResumeThread failed: " << GetLastError();
-  }
+  // Disable the priority boost that the thread would otherwise receive on
+  // resume. We do this to avoid artificially altering the dynamics of the
+  // executing application any more than we already are by suspending and
+  // resuming the thread.
+  //
+  // Note that this can racily disable a priority boost that otherwise would
+  // have been given to the thread, if the thread is waiting on other wait
+  // conditions at the time of SuspendThread and those conditions are satisfied
+  // before priority boost is reenabled. The measured length of this window is
+  // ~100us, so this should occur fairly rarely.
+  ScopedDisablePriorityBoost disable_priority_boost(thread_handle);
+  bool resume_thread_succeeded = ::ResumeThread(thread_handle) != -1;
+  CHECK(resume_thread_succeeded) << "ResumeThread failed: " << GetLastError();
 
   return stack_depth;
 }
 
-}  // namespace
+class NativeStackSamplerWin : public NativeStackSampler {
+ public:
+  explicit NativeStackSamplerWin(win::ScopedHandle thread_handle);
+  ~NativeStackSamplerWin() override;
 
-scoped_ptr<StackSamplingProfiler::NativeStackSampler>
-StackSamplingProfiler::NativeStackSampler::Create(PlatformThreadId thread_id) {
-#if _WIN64
-  // Get the thread's handle.
-  HANDLE thread_handle = ::OpenThread(
-      THREAD_GET_CONTEXT | THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
-      FALSE,
-      thread_id);
-  DCHECK(thread_handle) << "OpenThread failed";
+  // StackSamplingProfiler::NativeStackSampler:
+  void ProfileRecordingStarting(
+      std::vector<StackSamplingProfiler::Module>* modules) override;
+  void RecordStackSample(StackSamplingProfiler::Sample* sample) override;
+  void ProfileRecordingStopped() override;
 
-  return scoped_ptr<NativeStackSampler>(new NativeStackSamplerWin(
-      win::ScopedHandle(thread_handle)));
-#else
-  return scoped_ptr<NativeStackSampler>();
-#endif
-}
+ private:
+  // Attempts to query the module filename, base address, and id for
+  // |module_handle|, and store them in |module|. Returns true if it succeeded.
+  static bool GetModuleForHandle(HMODULE module_handle,
+                                 StackSamplingProfiler::Module* module);
+
+  // Gets the index for the Module corresponding to |module_handle| in
+  // |modules|, adding it if it's not already present. Returns
+  // StackSamplingProfiler::Frame::kUnknownModuleIndex if no Module can be
+  // determined for |module|.
+  size_t GetModuleIndex(HMODULE module_handle,
+                        std::vector<StackSamplingProfiler::Module>* modules);
+
+  // Copies the stack information represented by |instruction_pointers| into
+  // |sample| and |modules|.
+  void CopyToSample(const void* const instruction_pointers[],
+                    const HMODULE module_handles[],
+                    int stack_depth,
+                    StackSamplingProfiler::Sample* sample,
+                    std::vector<StackSamplingProfiler::Module>* modules);
+
+  win::ScopedHandle thread_handle_;
+  // Weak. Points to the modules associated with the profile being recorded
+  // between ProfileRecordingStarting() and ProfileRecordingStopped().
+  std::vector<StackSamplingProfiler::Module>* current_modules_;
+  // Maps a module handle to the corresponding Module's index within
+  // current_modules_.
+  std::map<HMODULE, size_t> profile_module_index_;
+
+  DISALLOW_COPY_AND_ASSIGN(NativeStackSamplerWin);
+};
 
 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");
-    // This should always be non-null, but handle just in case.
-    if (nt_dll_handle) {
-      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;
+    std::vector<StackSamplingProfiler::Module>* modules) {
+  current_modules_ = modules;
   profile_module_index_.clear();
 }
 
 void NativeStackSamplerWin::RecordStackSample(
     StackSamplingProfiler::Sample* sample) {
-  DCHECK(current_profile_);
+  DCHECK(current_modules_);
 
   const int max_stack_size = 64;
   const void* instruction_pointers[max_stack_size] = {0};
-  HMODULE modules[max_stack_size] = {0};
+  HMODULE module_handles[max_stack_size] = {0};
 
   bool last_frame_is_unknown_function = false;
   int stack_depth = SuspendThreadAndRecordStack(
       thread_handle_.Get(), max_stack_size, instruction_pointers,
       &last_frame_is_unknown_function);
-  FindModulesForAddresses(instruction_pointers, modules, stack_depth,
-                          last_frame_is_unknown_function);
-  CopyToSample(instruction_pointers, modules, stack_depth, sample,
-               &current_profile_->modules);
-  FreeModules(stack_depth, modules);
+  FindModuleHandlesForAddresses(instruction_pointers, module_handles,
+                                stack_depth, last_frame_is_unknown_function);
+  CopyToSample(instruction_pointers, module_handles, stack_depth, sample,
+               current_modules_);
+  FreeModuleHandles(stack_depth, module_handles);
 }
 
 void NativeStackSamplerWin::ProfileRecordingStopped() {
-  current_profile_ = nullptr;
+  current_modules_ = nullptr;
 }
 
 // static
-bool NativeStackSamplerWin::GetModuleInfo(
-    HMODULE module,
-    StackSamplingProfiler::Module* module_info) {
+bool NativeStackSamplerWin::GetModuleForHandle(
+    HMODULE module_handle,
+    StackSamplingProfiler::Module* module) {
   wchar_t module_name[MAX_PATH];
   DWORD result_length =
-      GetModuleFileName(module, module_name, arraysize(module_name));
+      GetModuleFileName(module_handle, module_name, arraysize(module_name));
   if (result_length == 0)
     return false;
 
-  module_info->filename = base::FilePath(module_name);
+  module->filename = base::FilePath(module_name);
 
-  module_info->base_address = reinterpret_cast<const void*>(module);
+  module->base_address = reinterpret_cast<const void*>(module_handle);
 
   GUID guid;
   DWORD age;
-  win::PEImage(module).GetDebugId(&guid, &age);
-  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));
+  win::PEImage(module_handle).GetDebugId(&guid, &age);
+  module->id.assign(reinterpret_cast<char*>(&guid), sizeof(guid));
+  module->id.append(reinterpret_cast<char*>(&age), sizeof(age));
 
   return true;
 }
 
+size_t NativeStackSamplerWin::GetModuleIndex(
+    HMODULE module_handle,
+    std::vector<StackSamplingProfiler::Module>* modules) {
+  if (!module_handle)
+    return StackSamplingProfiler::Frame::kUnknownModuleIndex;
+
+  auto loc = profile_module_index_.find(module_handle);
+  if (loc == profile_module_index_.end()) {
+    StackSamplingProfiler::Module module;
+    if (!GetModuleForHandle(module_handle, &module))
+      return StackSamplingProfiler::Frame::kUnknownModuleIndex;
+    modules->push_back(module);
+    loc = profile_module_index_.insert(std::make_pair(
+        module_handle, modules->size() - 1)).first;
+  }
+
+  return loc->second;
+}
+
 void NativeStackSamplerWin::CopyToSample(
     const void* const instruction_pointers[],
-    const HMODULE modules[],
+    const HMODULE module_handles[],
     int stack_depth,
     StackSamplingProfiler::Sample* sample,
-    std::vector<StackSamplingProfiler::Module>* module_infos) {
+    std::vector<StackSamplingProfiler::Module>* module) {
   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];
+    sample->push_back(StackSamplingProfiler::Frame(
+        instruction_pointers[i],
+        GetModuleIndex(module_handles[i], module)));
+  }
+}
 
-    // Record an invalid module index if we don't have a valid module.
-    if (!modules[i]) {
-      frame.module_index = -1;
-      continue;
-    }
+}  // namespace
 
-    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;
-    }
+scoped_ptr<NativeStackSampler> NativeStackSampler::Create(
+    PlatformThreadId thread_id) {
+#if _WIN64
+  // Get the thread's handle.
+  HANDLE thread_handle = ::OpenThread(
+      THREAD_GET_CONTEXT | THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
+      FALSE,
+      thread_id);
 
-    frame.module_index = loc->second;
+  if (thread_handle) {
+    return scoped_ptr<NativeStackSampler>(new NativeStackSamplerWin(
+        win::ScopedHandle(thread_handle)));
   }
+#endif
+  return scoped_ptr<NativeStackSampler>();
 }
 
 }  // namespace base