chrome/browser/private_working_set_snapshot_win.cc - chromium/src - Git at Google

 // 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 <stddef.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() {}

 PrivateWorkingSetSnapshot::~PrivateWorkingSetSnapshot() {}

 void PrivateWorkingSetSnapshot::Initialize() {
   DCHECK(!initialized_);

   // Set this to true whether initialization succeeds or not. That is, only try
   // once.
   initialized_ = true;

   // 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) {
     process_names_.clear();
     return;
   }

   // Pdh.dll has a format-string bug that causes crashes on Windows 8 and 8.1.
   // This was patched (around October 2014) but the broken versions are still on
   // a few machines. CreateFileVersionInfoForModule touches the disk so it
   // be used on the main thread. If this call gets moved off of the main thread
   // then pdh.dll version checking can be found in the history of
   // https://codereview.chromium.org/1269223005

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

   query_handle_.Set(query_handle);

   for (const auto& process_name : process_names_) {
     AddToMonitorList(process_name);
   }
   process_names_.clear();
 }

 void PrivateWorkingSetSnapshot::AddToMonitorList(
     const std::string& process_name) {
   if (!query_handle_.IsValid()) {
     // Save the name for later.
     if (!initialized_)
       process_names_.push_back(process_name);
     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() {
   // Make sure this is called once.
   if (!initialized_)
     Initialize();

   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) != static_cast<PDH_STATUS>(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) != static_cast<PDH_STATUS>(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;
 }
	// 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 <stddef.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() {}

	PrivateWorkingSetSnapshot::~PrivateWorkingSetSnapshot() {}

	void PrivateWorkingSetSnapshot::Initialize() {
	DCHECK(!initialized_);

	// Set this to true whether initialization succeeds or not. That is, only try
	// once.
	initialized_ = true;

	// 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) {
	process_names_.clear();
	return;
	}

	// Pdh.dll has a format-string bug that causes crashes on Windows 8 and 8.1.
	// This was patched (around October 2014) but the broken versions are still on
	// a few machines. CreateFileVersionInfoForModule touches the disk so it
	// be used on the main thread. If this call gets moved off of the main thread
	// then pdh.dll version checking can be found in the history of
	// https://codereview.chromium.org/1269223005

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

	query_handle_.Set(query_handle);

	for (const auto& process_name : process_names_) {
	AddToMonitorList(process_name);
	}
	process_names_.clear();
	}

	void PrivateWorkingSetSnapshot::AddToMonitorList(
	const std::string& process_name) {
	if (!query_handle_.IsValid()) {
	// Save the name for later.
	if (!initialized_)
	process_names_.push_back(process_name);
	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() {
	// Make sure this is called once.
	if (!initialized_)
	Initialize();

	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) != static_cast<PDH_STATUS>(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) != static_cast<PDH_STATUS>(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;
	}