tools/win/IdleWakeups/system_information_sampler.cpp - chromium/src - Git at Google

 // Copyright 2016 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 "stdafx.h"

 #include <algorithm>

 #include "system_information_sampler.h"

 // From ntdef.h
 typedef struct _UNICODE_STRING {
   USHORT Length;
   USHORT MaximumLength;
   PWCH Buffer;
 } UNICODE_STRING;

 // From <wdm.h>
 typedef LONG KPRIORITY;
 typedef LONG KWAIT_REASON;  // Full definition is in wdm.h

 // From ntddk.h
 typedef struct _VM_COUNTERS {
   SIZE_T PeakVirtualSize;
   SIZE_T VirtualSize;
   ULONG PageFaultCount;
   // Padding here in 64-bit
   SIZE_T PeakWorkingSetSize;
   SIZE_T WorkingSetSize;
   SIZE_T QuotaPeakPagedPoolUsage;
   SIZE_T QuotaPagedPoolUsage;
   SIZE_T QuotaPeakNonPagedPoolUsage;
   SIZE_T QuotaNonPagedPoolUsage;
   SIZE_T PagefileUsage;
   SIZE_T PeakPagefileUsage;
 } VM_COUNTERS;

 // Two possibilities available from here:
 // http://stackoverflow.com/questions/28858849/where-is-system-information-class-defined

 typedef enum _SYSTEM_INFORMATION_CLASS {
   SystemBasicInformation = 0,
   SystemPerformanceInformation = 2,
   SystemTimeOfDayInformation = 3,
   SystemProcessInformation = 5,  // This is the number that we need
   SystemProcessorPerformanceInformation = 8,
   SystemInterruptInformation = 23,
   SystemExceptionInformation = 33,
   SystemRegistryQuotaInformation = 37,
   SystemLookasideInformation = 45
 } SYSTEM_INFORMATION_CLASS;

 // https://msdn.microsoft.com/en-us/library/gg750647.aspx?f=255&MSPPError=-2147217396
 typedef struct {
   HANDLE UniqueProcess;  // Actually process ID
   HANDLE UniqueThread;   // Actually thread ID
 } CLIENT_ID;

 // From http://alax.info/blog/1182, with corrections and modifications
 // Originally from
 // http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
 struct SYSTEM_THREAD_INFORMATION {
   ULONGLONG KernelTime;
   ULONGLONG UserTime;
   ULONGLONG CreateTime;
   ULONG WaitTime;
   // Padding here in 64-bit
   PVOID StartAddress;
   CLIENT_ID ClientId;
   KPRIORITY Priority;
   LONG BasePriority;
   ULONG ContextSwitchCount;
   ULONG State;
   KWAIT_REASON WaitReason;
 };
 #if _M_X64
 static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 80,
               "Structure size mismatch");
 #else
 static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 64,
               "Structure size mismatch");
 #endif

 // From http://alax.info/blog/1182, with corrections and modifications
 // Originally from
 // http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
 struct SYSTEM_PROCESS_INFORMATION {
   ULONG NextEntryOffset;
   ULONG NumberOfThreads;
   // http://processhacker.sourceforge.net/doc/struct___s_y_s_t_e_m___p_r_o_c_e_s_s___i_n_f_o_r_m_a_t_i_o_n.html
   ULONGLONG WorkingSetPrivateSize;
   ULONG HardFaultCount;
   ULONG Reserved1;
   ULONGLONG CycleTime;
   ULONGLONG CreateTime;
   ULONGLONG UserTime;
   ULONGLONG KernelTime;
   UNICODE_STRING ImageName;
   KPRIORITY BasePriority;
   HANDLE ProcessId;
   HANDLE ParentProcessId;
   ULONG HandleCount;
   ULONG Reserved2[2];
   // Padding here in 64-bit
   VM_COUNTERS VirtualMemoryCounters;
   size_t Reserved3;
   IO_COUNTERS IoCounters;
   SYSTEM_THREAD_INFORMATION Threads[1];
 };
 #if _M_X64
 static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 336,
               "Structure size mismatch");
 #else
 static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 248,
               "Structure size mismatch");
 #endif

 // ntstatus.h conflicts with windows.h so define this locally.
 #define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
 #define STATUS_BUFFER_TOO_SMALL ((NTSTATUS)0xC0000023L)
 #define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004L)

 typedef NTSTATUS(WINAPI* NTQUERYSYSTEMINFORMATION)(
     SYSTEM_INFORMATION_CLASS SystemInformationClass,
     PVOID SystemInformation,
     ULONG SystemInformationLength,
     PULONG ReturnLength);

 __declspec(noreturn) void oops(const char* pMessage) {
   printf("%s\n", pMessage);
   exit(0);
 }

 // Simple memory buffer wrapper for passing the data out of
 // QuerySystemProcessInformation.
 class ByteBuffer {
  public:
   explicit ByteBuffer(size_t capacity) : size_(0), capacity_(0) {
     if (capacity > 0)
       grow(capacity);
   }

   ~ByteBuffer() {}

   BYTE* data() { return data_.get(); }

   size_t size() { return size_; }

   void set_size(size_t new_size) { size_ = new_size; }

   size_t capacity() { return capacity_; }

   void grow(size_t new_capacity) {
     capacity_ = new_capacity;
     data_.reset(new BYTE[new_capacity]);
   }

  private:
   std::unique_ptr<BYTE[]> data_;
   size_t size_;
   size_t capacity_;

   ByteBuffer& operator=(const ByteBuffer&) = delete;
   ByteBuffer(const ByteBuffer&) = delete;
 };

 // Wrapper for NtQuerySystemProcessInformation with buffer reallocation logic.
 bool QuerySystemProcessInformation(ByteBuffer* buffer) {
   typedef NTSTATUS(WINAPI * NTQUERYSYSTEMINFORMATION)(
       SYSTEM_INFORMATION_CLASS SystemInformationClass, PVOID SystemInformation,
       ULONG SystemInformationLength, PULONG ReturnLength);

   HMODULE ntdll = GetModuleHandle(L"ntdll.dll");
   if (!ntdll) {
     oops("Couldn't load ntdll.dll");
   }

   NTQUERYSYSTEMINFORMATION nt_query_system_information_ptr =
       reinterpret_cast<NTQUERYSYSTEMINFORMATION>(
           GetProcAddress(ntdll, "NtQuerySystemInformation"));
   if (!nt_query_system_information_ptr)
     oops("Couldn't find NtQuerySystemInformation");

   NTSTATUS result;

   // There is a potential race condition between growing the buffer and new
   // processes being created. Try a few times before giving up.
   for (int i = 0; i < 10; i++) {
     ULONG data_size = 0;
     ULONG buffer_size = static_cast<ULONG>(buffer->capacity());
     result = nt_query_system_information_ptr(
         SystemProcessInformation, buffer->data(), buffer_size, &data_size);

     if (result == STATUS_SUCCESS) {
       buffer->set_size(data_size);
       break;
     }

     if (result == STATUS_INFO_LENGTH_MISMATCH ||
         result == STATUS_BUFFER_TOO_SMALL) {
       // Insufficient buffer. Grow to the returned |data_size| plus 10% extra
       // to avoid frequent reallocations and try again.
       buffer->grow(static_cast<ULONG>(data_size * 1.1));
     } else {
       // An error other than the two above.
       break;
     }
   }

   return result == STATUS_SUCCESS;
 }

 SystemInformationSampler::SystemInformationSampler(
     const wchar_t* process_name) {
   lstrcpyn(target_process_name_, process_name,
            sizeof(target_process_name_) / sizeof(wchar_t));

   QueryPerformanceFrequency(&perf_frequency_);
   QueryPerformanceCounter(&initial_counter_);
 }

 SystemInformationSampler::~SystemInformationSampler() {}

 std::unique_ptr<ProcessDataSnapshot> SystemInformationSampler::TakeSnapshot() {
   // Preallocate the buffer with the size determined on the previous call to
   // QuerySystemProcessInformation. This should be sufficient most of the time.
   // QuerySystemProcessInformation will grow the buffer if necessary.
   ByteBuffer data_buffer(previous_buffer_size_);

   if (!QuerySystemProcessInformation(&data_buffer))
     return std::unique_ptr<ProcessDataSnapshot>();

   previous_buffer_size_ = data_buffer.capacity();

   std::unique_ptr<ProcessDataSnapshot> snapshot(new ProcessDataSnapshot);

   LARGE_INTEGER perf_counter_value;
   QueryPerformanceCounter(&perf_counter_value);
   snapshot->timestamp = static_cast<double>(
       (perf_counter_value.QuadPart - initial_counter_.QuadPart) /
       perf_frequency_.QuadPart);

   for (size_t offset = 0; offset < data_buffer.size();) {
     auto pi = reinterpret_cast<const SYSTEM_PROCESS_INFORMATION*>(
         data_buffer.data() + offset);

     // Validate that the offset is valid and all needed data is within
     // the buffer boundary.
     if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) > data_buffer.size())
       break;
     if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) +
             (pi->NumberOfThreads - 1) * sizeof(SYSTEM_THREAD_INFORMATION) >
         data_buffer.size())
       break;

     if (pi->ImageName.Buffer) {
       // Validate that the image name is within the buffer boundary.
       // ImageName.Length seems to be in bytes rather than characters.
       size_t image_name_offset =
           reinterpret_cast<BYTE*>(pi->ImageName.Buffer) - data_buffer.data();
       if (image_name_offset + pi->ImageName.Length > data_buffer.size())
         break;

       // Check if this is a chrome process. Ignore all other processes.
       if (wcsncmp(target_process_name_filter(), pi->ImageName.Buffer,
                   lstrlen(target_process_name_filter())) == 0) {
         // Collect enough data to be able to do a diff between two snapshots.
         // Some threads might stop or new threads might be created between two
         // snapshots. If a thread with a large number of context switches gets
         // terminated the total number of context switches for the process might
         // go down and the delta would be negative.
         // To avoid that we need to compare thread IDs between two snapshots and
         // not count context switches for threads that are missing in the most
         // recent snapshot.
         ProcessData process_data;

         process_data.cpu_time = pi->KernelTime + pi->UserTime;
         process_data.working_set = pi->WorkingSetPrivateSize;

         // Iterate over threads and store each thread's ID and number of context
         // switches.
         for (ULONG thread_index = 0; thread_index < pi->NumberOfThreads;
              ++thread_index) {
           const SYSTEM_THREAD_INFORMATION* ti = &pi->Threads[thread_index];
           if (ti->ClientId.UniqueProcess != pi->ProcessId)
             continue;

           ThreadData thread_data;
           thread_data.thread_id = ti->ClientId.UniqueThread;
           thread_data.context_switches = ti->ContextSwitchCount;
           process_data.threads.push_back(thread_data);
         }

         // Order thread data by thread ID to help diff two snapshots.
         std::sort(process_data.threads.begin(), process_data.threads.end(),
                   [](const ThreadData& l, const ThreadData r) {
                     return l.thread_id < r.thread_id;
                   });

         snapshot->processes.insert(
             std::make_pair(pi->ProcessId, std::move(process_data)));
       }
     }

     // Check for end of the list.
     if (!pi->NextEntryOffset)
       break;

     // Jump to the next entry.
     offset += pi->NextEntryOffset;
   }

   return snapshot;
 }
	// Copyright 2016 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 "stdafx.h"

	#include <algorithm>

	#include "system_information_sampler.h"

	// From ntdef.h
	typedef struct _UNICODE_STRING {
	USHORT Length;
	USHORT MaximumLength;
	PWCH Buffer;
	} UNICODE_STRING;

	// From <wdm.h>
	typedef LONG KPRIORITY;
	typedef LONG KWAIT_REASON; // Full definition is in wdm.h

	// From ntddk.h
	typedef struct _VM_COUNTERS {
	SIZE_T PeakVirtualSize;
	SIZE_T VirtualSize;
	ULONG PageFaultCount;
	// Padding here in 64-bit
	SIZE_T PeakWorkingSetSize;
	SIZE_T WorkingSetSize;
	SIZE_T QuotaPeakPagedPoolUsage;
	SIZE_T QuotaPagedPoolUsage;
	SIZE_T QuotaPeakNonPagedPoolUsage;
	SIZE_T QuotaNonPagedPoolUsage;
	SIZE_T PagefileUsage;
	SIZE_T PeakPagefileUsage;
	} VM_COUNTERS;

	// Two possibilities available from here:
	// http://stackoverflow.com/questions/28858849/where-is-system-information-class-defined

	typedef enum _SYSTEM_INFORMATION_CLASS {
	SystemBasicInformation = 0,
	SystemPerformanceInformation = 2,
	SystemTimeOfDayInformation = 3,
	SystemProcessInformation = 5, // This is the number that we need
	SystemProcessorPerformanceInformation = 8,
	SystemInterruptInformation = 23,
	SystemExceptionInformation = 33,
	SystemRegistryQuotaInformation = 37,
	SystemLookasideInformation = 45
	} SYSTEM_INFORMATION_CLASS;

	// https://msdn.microsoft.com/en-us/library/gg750647.aspx?f=255&MSPPError=-2147217396
	typedef struct {
	HANDLE UniqueProcess; // Actually process ID
	HANDLE UniqueThread; // Actually thread ID
	} CLIENT_ID;

	// From http://alax.info/blog/1182, with corrections and modifications
	// Originally from
	// http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
	struct SYSTEM_THREAD_INFORMATION {
	ULONGLONG KernelTime;
	ULONGLONG UserTime;
	ULONGLONG CreateTime;
	ULONG WaitTime;
	// Padding here in 64-bit
	PVOID StartAddress;
	CLIENT_ID ClientId;
	KPRIORITY Priority;
	LONG BasePriority;
	ULONG ContextSwitchCount;
	ULONG State;
	KWAIT_REASON WaitReason;
	};
	#if _M_X64
	static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 80,
	"Structure size mismatch");
	#else
	static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 64,
	"Structure size mismatch");
	#endif

	// From http://alax.info/blog/1182, with corrections and modifications
	// Originally from
	// http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
	struct SYSTEM_PROCESS_INFORMATION {
	ULONG NextEntryOffset;
	ULONG NumberOfThreads;
	// http://processhacker.sourceforge.net/doc/struct___s_y_s_t_e_m___p_r_o_c_e_s_s___i_n_f_o_r_m_a_t_i_o_n.html
	ULONGLONG WorkingSetPrivateSize;
	ULONG HardFaultCount;
	ULONG Reserved1;
	ULONGLONG CycleTime;
	ULONGLONG CreateTime;
	ULONGLONG UserTime;
	ULONGLONG KernelTime;
	UNICODE_STRING ImageName;
	KPRIORITY BasePriority;
	HANDLE ProcessId;
	HANDLE ParentProcessId;
	ULONG HandleCount;
	ULONG Reserved2[2];
	// Padding here in 64-bit
	VM_COUNTERS VirtualMemoryCounters;
	size_t Reserved3;
	IO_COUNTERS IoCounters;
	SYSTEM_THREAD_INFORMATION Threads[1];
	};
	#if _M_X64
	static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 336,
	"Structure size mismatch");
	#else
	static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 248,
	"Structure size mismatch");
	#endif

	// ntstatus.h conflicts with windows.h so define this locally.
	#define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
	#define STATUS_BUFFER_TOO_SMALL ((NTSTATUS)0xC0000023L)
	#define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004L)

	typedef NTSTATUS(WINAPI* NTQUERYSYSTEMINFORMATION)(
	SYSTEM_INFORMATION_CLASS SystemInformationClass,
	PVOID SystemInformation,
	ULONG SystemInformationLength,
	PULONG ReturnLength);

	__declspec(noreturn) void oops(const char* pMessage) {
	printf("%s\n", pMessage);
	exit(0);
	}

	// Simple memory buffer wrapper for passing the data out of
	// QuerySystemProcessInformation.
	class ByteBuffer {
	public:
	explicit ByteBuffer(size_t capacity) : size_(0), capacity_(0) {
	if (capacity > 0)
	grow(capacity);
	}

	~ByteBuffer() {}

	BYTE* data() { return data_.get(); }

	size_t size() { return size_; }

	void set_size(size_t new_size) { size_ = new_size; }

	size_t capacity() { return capacity_; }

	void grow(size_t new_capacity) {
	capacity_ = new_capacity;
	data_.reset(new BYTE[new_capacity]);
	}

	private:
	std::unique_ptr<BYTE[]> data_;
	size_t size_;
	size_t capacity_;

	ByteBuffer& operator=(const ByteBuffer&) = delete;
	ByteBuffer(const ByteBuffer&) = delete;
	};

	// Wrapper for NtQuerySystemProcessInformation with buffer reallocation logic.
	bool QuerySystemProcessInformation(ByteBuffer* buffer) {
	typedef NTSTATUS(WINAPI * NTQUERYSYSTEMINFORMATION)(
	SYSTEM_INFORMATION_CLASS SystemInformationClass, PVOID SystemInformation,
	ULONG SystemInformationLength, PULONG ReturnLength);

	HMODULE ntdll = GetModuleHandle(L"ntdll.dll");
	if (!ntdll) {
	oops("Couldn't load ntdll.dll");
	}

	NTQUERYSYSTEMINFORMATION nt_query_system_information_ptr =
	reinterpret_cast<NTQUERYSYSTEMINFORMATION>(
	GetProcAddress(ntdll, "NtQuerySystemInformation"));
	if (!nt_query_system_information_ptr)
	oops("Couldn't find NtQuerySystemInformation");

	NTSTATUS result;

	// There is a potential race condition between growing the buffer and new
	// processes being created. Try a few times before giving up.
	for (int i = 0; i < 10; i++) {
	ULONG data_size = 0;
	ULONG buffer_size = static_cast<ULONG>(buffer->capacity());
	result = nt_query_system_information_ptr(
	SystemProcessInformation, buffer->data(), buffer_size, &data_size);

	if (result == STATUS_SUCCESS) {
	buffer->set_size(data_size);
	break;
	}

	if (result == STATUS_INFO_LENGTH_MISMATCH \|\|
	result == STATUS_BUFFER_TOO_SMALL) {
	// Insufficient buffer. Grow to the returned \|data_size\| plus 10% extra
	// to avoid frequent reallocations and try again.
	buffer->grow(static_cast<ULONG>(data_size * 1.1));
	} else {
	// An error other than the two above.
	break;
	}
	}

	return result == STATUS_SUCCESS;
	}

	SystemInformationSampler::SystemInformationSampler(
	const wchar_t* process_name) {
	lstrcpyn(target_process_name_, process_name,
	sizeof(target_process_name_) / sizeof(wchar_t));

	QueryPerformanceFrequency(&perf_frequency_);
	QueryPerformanceCounter(&initial_counter_);
	}

	SystemInformationSampler::~SystemInformationSampler() {}

	std::unique_ptr<ProcessDataSnapshot> SystemInformationSampler::TakeSnapshot() {
	// Preallocate the buffer with the size determined on the previous call to
	// QuerySystemProcessInformation. This should be sufficient most of the time.
	// QuerySystemProcessInformation will grow the buffer if necessary.
	ByteBuffer data_buffer(previous_buffer_size_);

	if (!QuerySystemProcessInformation(&data_buffer))
	return std::unique_ptr<ProcessDataSnapshot>();

	previous_buffer_size_ = data_buffer.capacity();

	std::unique_ptr<ProcessDataSnapshot> snapshot(new ProcessDataSnapshot);

	LARGE_INTEGER perf_counter_value;
	QueryPerformanceCounter(&perf_counter_value);
	snapshot->timestamp = static_cast<double>(
	(perf_counter_value.QuadPart - initial_counter_.QuadPart) /
	perf_frequency_.QuadPart);

	for (size_t offset = 0; offset < data_buffer.size();) {
	auto pi = reinterpret_cast<const SYSTEM_PROCESS_INFORMATION*>(
	data_buffer.data() + offset);

	// Validate that the offset is valid and all needed data is within
	// the buffer boundary.
	if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) > data_buffer.size())
	break;
	if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) +
	(pi->NumberOfThreads - 1) * sizeof(SYSTEM_THREAD_INFORMATION) >
	data_buffer.size())
	break;

	if (pi->ImageName.Buffer) {
	// Validate that the image name is within the buffer boundary.
	// ImageName.Length seems to be in bytes rather than characters.
	size_t image_name_offset =
	reinterpret_cast<BYTE*>(pi->ImageName.Buffer) - data_buffer.data();
	if (image_name_offset + pi->ImageName.Length > data_buffer.size())
	break;

	// Check if this is a chrome process. Ignore all other processes.
	if (wcsncmp(target_process_name_filter(), pi->ImageName.Buffer,
	lstrlen(target_process_name_filter())) == 0) {
	// Collect enough data to be able to do a diff between two snapshots.
	// Some threads might stop or new threads might be created between two
	// snapshots. If a thread with a large number of context switches gets
	// terminated the total number of context switches for the process might
	// go down and the delta would be negative.
	// To avoid that we need to compare thread IDs between two snapshots and
	// not count context switches for threads that are missing in the most
	// recent snapshot.
	ProcessData process_data;

	process_data.cpu_time = pi->KernelTime + pi->UserTime;
	process_data.working_set = pi->WorkingSetPrivateSize;

	// Iterate over threads and store each thread's ID and number of context
	// switches.
	for (ULONG thread_index = 0; thread_index < pi->NumberOfThreads;
	++thread_index) {
	const SYSTEM_THREAD_INFORMATION* ti = &pi->Threads[thread_index];
	if (ti->ClientId.UniqueProcess != pi->ProcessId)
	continue;

	ThreadData thread_data;
	thread_data.thread_id = ti->ClientId.UniqueThread;
	thread_data.context_switches = ti->ContextSwitchCount;
	process_data.threads.push_back(thread_data);
	}

	// Order thread data by thread ID to help diff two snapshots.
	std::sort(process_data.threads.begin(), process_data.threads.end(),
	[](const ThreadData& l, const ThreadData r) {
	return l.thread_id < r.thread_id;
	});

	snapshot->processes.insert(
	std::make_pair(pi->ProcessId, std::move(process_data)));
	}
	}

	// Check for end of the list.
	if (!pi->NextEntryOffset)
	break;

	// Jump to the next entry.
	offset += pi->NextEntryOffset;
	}

	return snapshot;
	}