chrome/browser/private_working_set_snapshot_win.cc

   1 // Copyright 2015 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "chrome/browser/private_working_set_snapshot.h"
   6
   7 #include <pdh.h>
   8 #include <pdhmsg.h>
   9
  10 #include <algorithm>
  11
  12 #include "base/numerics/safe_conversions.h"
  13 #include "base/win/windows_version.h"
  14
  15 // Link pdh.lib (import library for pdh.dll) whenever this file is linked in.
  16 #pragma comment(lib, "pdh.lib")
  17
  18 PrivateWorkingSetSnapshot::PrivateWorkingSetSnapshot() {
  19   // The Pdh APIs are supported on Windows XP and above, but the "Working Set -
  20   // Private" counter that PrivateWorkingSetSnapshot depends on is not defined
  21   // until Windows Vista and is not reliable until Windows 7. Early-out to avoid
  22   // wasted effort. All queries will return zero and will have to use the
  23   // fallback calculations.
  24   if (base::win::GetVersion() <= base::win::VERSION_VISTA)
  25     return;
  26
  27   // Create a Pdh query
  28   PDH_HQUERY query_handle;
  29   if (PdhOpenQuery(NULL, NULL, &query_handle) != ERROR_SUCCESS) {
  30     return;
  31   }
  32
  33   query_handle_.Set(query_handle);
  34 }
  35
  36 PrivateWorkingSetSnapshot::~PrivateWorkingSetSnapshot() {
  37 }
  38
  39 void PrivateWorkingSetSnapshot::AddToMonitorList(
  40     const std::string& process_name) {
  41   if (!query_handle_.IsValid())
  42     return;
  43
  44   // Create the magic strings that will return a list of process IDs and a list
  45   // of private working sets. The 'process_name' variable should be something
  46   // like "chrome". The '*' character indicates that we want records for all
  47   // processes whose names start with process_name - all chrome processes, but
  48   // also all 'chrome_editor.exe' processes or other matching names. The excess
  49   // information is unavoidable but harmless.
  50   std::string process_id_query = "\\Process(" + process_name + "*)\\ID Process";
  51   std::string private_ws_query =
  52       "\\Process(" + process_name + "*)\\Working Set - Private";
  53
  54   // Add the two counters to the query.
  55   PdhCounterPair new_counters;
  56   if (PdhAddCounterA(query_handle_.Get(), process_id_query.c_str(), NULL,
  57                      &new_counters.process_id_handle) != ERROR_SUCCESS) {
  58     return;
  59   }
  60
  61   // If adding the second counter fails then we should remove the first one.
  62   if (PdhAddCounterA(query_handle_.Get(), private_ws_query.c_str(), NULL,
  63                      &new_counters.private_ws_handle) != ERROR_SUCCESS) {
  64     PdhRemoveCounter(new_counters.process_id_handle);
  65   }
  66
  67   // Record the pair of counter query handles so that we can query them later.
  68   counter_pairs_.push_back(new_counters);
  69 }
  70
  71 void PrivateWorkingSetSnapshot::Sample() {
  72   if (counter_pairs_.empty())
  73     return;
  74
  75   // Destroy all previous data.
  76   records_.resize(0);
  77   // Record the requested data into PDH's internal buffers.
  78   if (PdhCollectQueryData(query_handle_.Get()) != ERROR_SUCCESS)
  79     return;
  80
  81   for (auto& counter_pair : counter_pairs_) {
  82     // Find out how much space is required for the two counter arrays.
  83     // A return code of PDH_MORE_DATA indicates that we should call again with
  84     // the buffer size returned.
  85     DWORD buffer_size1 = 0;
  86     DWORD item_count1 = 0;
  87     // Process IDs should be retrieved as PDH_FMT_LONG
  88     if (PdhGetFormattedCounterArray(counter_pair.process_id_handle,
  89                                     PDH_FMT_LONG, &buffer_size1, &item_count1,
  90                                     nullptr) != PDH_MORE_DATA)
  91       continue;
  92     if (buffer_size1 == 0 || item_count1 == 0)
  93       continue;
  94
  95     DWORD buffer_size2 = 0;
  96     DWORD item_count2 = 0;
  97     // Working sets should be retrieved as PDH_FMT_LARGE (LONGLONG)
  98     // Note that if this second call to PdhGetFormattedCounterArray with the
  99     // buffer size and count variables being zero is omitted then the PID and
 100     // working-set results are not reliably correlated.
 101     if (PdhGetFormattedCounterArray(counter_pair.private_ws_handle,
 102                                     PDH_FMT_LARGE, &buffer_size2, &item_count2,
 103                                     nullptr) != PDH_MORE_DATA)
 104       continue;
 105
 106     // It is not clear whether Pdh guarantees that the two counters in the same
 107     // query will execute atomically - if they will see the same set of
 108     // processes. If they do not then the correspondence between "ID Process"
 109     // and "Working Set - Private" is lost and we have to discard these results.
 110     // In testing these values have always matched. If this check fails then
 111     // the old per-process memory calculations will be used instead.
 112     if (buffer_size1 != buffer_size2 || item_count1 != item_count2)
 113       continue;
 114
 115     // Allocate enough space for the results of both queries.
 116     std::vector<char> buffer(buffer_size1 * 2);
 117     // Retrieve the process ID data.
 118     auto process_id_data =
 119         reinterpret_cast<PDH_FMT_COUNTERVALUE_ITEM*>(&buffer[0]);
 120     if (PdhGetFormattedCounterArray(counter_pair.process_id_handle,
 121                                     PDH_FMT_LONG, &buffer_size1, &item_count1,
 122                                     process_id_data) != ERROR_SUCCESS)
 123       continue;
 124     // Retrieve the private working set data.
 125     auto private_ws_data =
 126         reinterpret_cast<PDH_FMT_COUNTERVALUE_ITEM*>(&buffer[buffer_size1]);
 127     if (PdhGetFormattedCounterArray(counter_pair.private_ws_handle,
 128                                     PDH_FMT_LARGE, &buffer_size1, &item_count1,
 129                                     private_ws_data) != ERROR_SUCCESS)
 130       continue;
 131
 132     // Make room for the new set of records.
 133     size_t start_offset = records_.size();
 134     records_.resize(start_offset + item_count1);
 135
 136     for (DWORD i = 0; i < item_count1; ++i) {
 137       records_[start_offset + i].process_id =
 138           process_id_data[i].FmtValue.longValue;
 139       // Integer overflow can happen here if a 32-bit process is monitoring a
 140       // 64-bit process so we do a saturated_cast.
 141       records_[start_offset + i].private_ws =
 142           base::saturated_cast<size_t>(private_ws_data[i].FmtValue.largeValue);
 143     }
 144   }
 145
 146   // The results will include all processes that match the passed in name,
 147   // regardless of whether they are spawned by the calling process.
 148   // The results must be sorted by process ID for efficient lookup.
 149   std::sort(records_.begin(), records_.end());
 150 }
 151
 152 size_t PrivateWorkingSetSnapshot::GetPrivateWorkingSet(
 153     base::ProcessId process_id) const {
 154   // Do a binary search for the requested process ID and return the working set
 155   // if found.
 156   auto p = std::lower_bound(records_.begin(), records_.end(), process_id);
 157   if (p != records_.end() && p->process_id == process_id)
 158     return p->private_ws;
 159
 160   return 0;
 161 }