summaryrefslogtreecommitdiffstats
path: root/chrome/browser/private_working_set_snapshot_win.cc
blob: a3955a43da8dd99b06b717348fa3e4169be8a3f6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
// 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 "chrome/browser/private_working_set_snapshot.h"

#include <pdh.h>
#include <pdhmsg.h>

#include <algorithm>

#include "base/numerics/safe_conversions.h"
#include "base/win/windows_version.h"

// Link pdh.lib (import library for pdh.dll) whenever this file is linked in.
#pragma comment(lib, "pdh.lib")

PrivateWorkingSetSnapshot::PrivateWorkingSetSnapshot() {
  // The Pdh APIs are supported on Windows XP and above, but the "Working Set -
  // Private" counter that PrivateWorkingSetSnapshot depends on is not defined
  // until Windows Vista and is not reliable until Windows 7. Early-out to avoid
  // wasted effort. All queries will return zero and will have to use the
  // fallback calculations.
  if (base::win::GetVersion() <= base::win::VERSION_VISTA)
    return;

  // Create a Pdh query
  PDH_HQUERY query_handle;
  if (PdhOpenQuery(NULL, NULL, &query_handle) != ERROR_SUCCESS) {
    return;
  }

  query_handle_.Set(query_handle);
}

PrivateWorkingSetSnapshot::~PrivateWorkingSetSnapshot() {
}

void PrivateWorkingSetSnapshot::AddToMonitorList(
    const std::string& process_name) {
  if (!query_handle_.IsValid())
    return;

  // Create the magic strings that will return a list of process IDs and a list
  // of private working sets. The 'process_name' variable should be something
  // like "chrome". The '*' character indicates that we want records for all
  // processes whose names start with process_name - all chrome processes, but
  // also all 'chrome_editor.exe' processes or other matching names. The excess
  // information is unavoidable but harmless.
  std::string process_id_query = "\\Process(" + process_name + "*)\\ID Process";
  std::string private_ws_query =
      "\\Process(" + process_name + "*)\\Working Set - Private";

  // Add the two counters to the query.
  PdhCounterPair new_counters;
  if (PdhAddCounterA(query_handle_.Get(), process_id_query.c_str(), NULL,
                     &new_counters.process_id_handle) != ERROR_SUCCESS) {
    return;
  }

  // If adding the second counter fails then we should remove the first one.
  if (PdhAddCounterA(query_handle_.Get(), private_ws_query.c_str(), NULL,
                     &new_counters.private_ws_handle) != ERROR_SUCCESS) {
    PdhRemoveCounter(new_counters.process_id_handle);
  }

  // Record the pair of counter query handles so that we can query them later.
  counter_pairs_.push_back(new_counters);
}

void PrivateWorkingSetSnapshot::Sample() {
  if (counter_pairs_.empty())
    return;

  // Destroy all previous data.
  records_.resize(0);
  // Record the requested data into PDH's internal buffers.
  if (PdhCollectQueryData(query_handle_.Get()) != ERROR_SUCCESS)
    return;

  for (auto& counter_pair : counter_pairs_) {
    // Find out how much space is required for the two counter arrays.
    // A return code of PDH_MORE_DATA indicates that we should call again with
    // the buffer size returned.
    DWORD buffer_size1 = 0;
    DWORD item_count1 = 0;
    // Process IDs should be retrieved as PDH_FMT_LONG
    if (PdhGetFormattedCounterArray(counter_pair.process_id_handle,
                                    PDH_FMT_LONG, &buffer_size1, &item_count1,
                                    nullptr) != PDH_MORE_DATA)
      continue;
    if (buffer_size1 == 0 || item_count1 == 0)
      continue;

    DWORD buffer_size2 = 0;
    DWORD item_count2 = 0;
    // Working sets should be retrieved as PDH_FMT_LARGE (LONGLONG)
    // Note that if this second call to PdhGetFormattedCounterArray with the
    // buffer size and count variables being zero is omitted then the PID and
    // working-set results are not reliably correlated.
    if (PdhGetFormattedCounterArray(counter_pair.private_ws_handle,
                                    PDH_FMT_LARGE, &buffer_size2, &item_count2,
                                    nullptr) != PDH_MORE_DATA)
      continue;

    // It is not clear whether Pdh guarantees that the two counters in the same
    // query will execute atomically - if they will see the same set of
    // processes. If they do not then the correspondence between "ID Process"
    // and "Working Set - Private" is lost and we have to discard these results.
    // In testing these values have always matched. If this check fails then
    // the old per-process memory calculations will be used instead.
    if (buffer_size1 != buffer_size2 || item_count1 != item_count2)
      continue;

    // Allocate enough space for the results of both queries.
    std::vector<char> buffer(buffer_size1 * 2);
    // Retrieve the process ID data.
    auto process_id_data =
        reinterpret_cast<PDH_FMT_COUNTERVALUE_ITEM*>(&buffer[0]);
    if (PdhGetFormattedCounterArray(counter_pair.process_id_handle,
                                    PDH_FMT_LONG, &buffer_size1, &item_count1,
                                    process_id_data) != ERROR_SUCCESS)
      continue;
    // Retrieve the private working set data.
    auto private_ws_data =
        reinterpret_cast<PDH_FMT_COUNTERVALUE_ITEM*>(&buffer[buffer_size1]);
    if (PdhGetFormattedCounterArray(counter_pair.private_ws_handle,
                                    PDH_FMT_LARGE, &buffer_size1, &item_count1,
                                    private_ws_data) != ERROR_SUCCESS)
      continue;

    // Make room for the new set of records.
    size_t start_offset = records_.size();
    records_.resize(start_offset + item_count1);

    for (DWORD i = 0; i < item_count1; ++i) {
      records_[start_offset + i].process_id =
          process_id_data[i].FmtValue.longValue;
      // Integer overflow can happen here if a 32-bit process is monitoring a
      // 64-bit process so we do a saturated_cast.
      records_[start_offset + i].private_ws =
          base::saturated_cast<size_t>(private_ws_data[i].FmtValue.largeValue);
    }
  }

  // The results will include all processes that match the passed in name,
  // regardless of whether they are spawned by the calling process.
  // The results must be sorted by process ID for efficient lookup.
  std::sort(records_.begin(), records_.end());
}

size_t PrivateWorkingSetSnapshot::GetPrivateWorkingSet(
    base::ProcessId process_id) const {
  // Do a binary search for the requested process ID and return the working set
  // if found.
  auto p = std::lower_bound(records_.begin(), records_.end(), process_id);
  if (p != records_.end() && p->process_id == process_id)
    return p->private_ws;

  return 0;
}