Task manager: Get physical memory efficiently for all processes from SharedSampler

chrome/browser/task_manager/sampling/shared_sampler_win.cc

 // 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 "chrome/browser/task_manager/sampling/shared_sampler.h"
 
 #include <windows.h>
 #include <winternl.h>
 
 #include <algorithm>
@@ skipping 201 matching lines @@
   base::PlatformThreadId thread_id;
   ULONG context_switches;
 };
 
 // Per-process data extracted from SYSTEM_PROCESS_INFORMATION and stored in a
 // snapshot. This structure is accessed only on the worker thread.
 struct ProcessData {
   ProcessData() = default;
   ProcessData(ProcessData&&) = default;
 
+  int64_t physical_bytes;
   std::vector<ThreadData> threads;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(ProcessData);
 };
 
 typedef std::map<base::ProcessId, ProcessData> ProcessDataMap;
 
 ULONG CountContextSwitchesDelta(const ProcessData& prev_process_data,
   const ProcessData& new_process_data) {
@@ skipping 73 matching lines @@
   // This object will be created on the UI thread, however the sequenced checker
   // will be used to assert we're running the expensive operations on one of the
   // blocking pool threads.
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
   worker_pool_sequenced_checker_.DetachFromSequence();
 }
 
 SharedSampler::~SharedSampler() {}
 
 int64_t SharedSampler::GetSupportedFlags() const {
-  return REFRESH_TYPE_IDLE_WAKEUPS;
+  return REFRESH_TYPE_IDLE_WAKEUPS | REFRESH_TYPE_PHYSICAL_MEMORY;
 }
 
 SharedSampler::Callbacks::Callbacks() {}
 
 SharedSampler::Callbacks::~Callbacks() {}
 
 SharedSampler::Callbacks::Callbacks(Callbacks&& other) {
   on_idle_wakeups = std::move(other.on_idle_wakeups);
+  on_physical_memory = std::move(other.on_physical_memory);
 }
 
 void SharedSampler::RegisterCallbacks(
     base::ProcessId process_id,
-    const OnIdleWakeupsCallback& on_idle_wakeups) {
+    const OnIdleWakeupsCallback& on_idle_wakeups,
+    const OnPhysicalMemoryCallback& on_physical_memory) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
   if (process_id == 0)
     return;
 
   Callbacks callbacks;
   callbacks.on_idle_wakeups = on_idle_wakeups;
+  callbacks.on_physical_memory = on_physical_memory;
   bool result = callbacks_map_.insert(
       std::make_pair(process_id, std::move(callbacks))).second;
   DCHECK(result);
 }
 
 void SharedSampler::UnregisterCallbacks(base::ProcessId process_id) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
   if (process_id == 0)
     return;
@@ skipping 124 matching lines @@
         // 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.physical_bytes =
+            static_cast<int64_t>(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 = static_cast<base::PlatformThreadId>(
@@ skipping 56 matching lines @@
     } else {
       // Process is missing in the previous snapshot.
       // Use entire number of context switches of the current process.
       idle_wakeups_delta = CountContextSwitchesDelta(ProcessData(), process);
     }
 
     RefreshResult result;
     result.process_id = process_id;
     result.idle_wakeups_per_second =
         static_cast<int>(round(idle_wakeups_delta / time_delta));
+    result.physical_bytes = process.physical_bytes;
     results->push_back(result);
   }
 }
 
 void SharedSampler::MakeResultsFromSnapshot(const ProcessDataSnapshot& snapshot,
                                             RefreshResults* results) {
   for (const auto& pair : snapshot.processes) {
     RefreshResult result;
     result.process_id = pair.first;
     // Use 0 for Idle Wakeups / sec in this case. This is consistent with
     // ProcessMetrics::CalculateIdleWakeupsPerSecond implementation.
     result.idle_wakeups_per_second = 0;
+    result.physical_bytes = pair.second.physical_bytes;
     results->push_back(result);
   }
 }
 
 void SharedSampler::OnRefreshDone(
     std::unique_ptr<RefreshResults> refresh_results) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
   DCHECK_NE(0, refresh_flags_);
 
   size_t result_index = 0;
 
   for (const auto& callback_entry : callbacks_map_) {
     base::ProcessId process_id = callback_entry.first;
     // A sentinel value of -1 is used when the result isn't available.
     // Task manager will use this to display 'N/A'.
     int idle_wakeups_per_second = -1;
+    int64_t physical_bytes = -1;
 
     // Match refresh result by |process_id|.
     // This relies on refresh results being ordered by Process ID.
     // Please note that |refresh_results| might contain some extra entries that
     // don't exist in |callbacks_map_| if there is more than one instance of
     // Chrome. It might be missing some entries too if there is a race condition
     // between getting process information on the worker thread and adding a
     // corresponding TaskGroup to the task manager.
     for (; result_index < refresh_results->size(); ++result_index) {
       const auto& result = (*refresh_results)[result_index];
       if (result.process_id == process_id) {
         // Data matched in |refresh_results|.
         idle_wakeups_per_second = result.idle_wakeups_per_second;
+        physical_bytes = result.physical_bytes;
         ++result_index;
         break;
       }
 
       if (result.process_id > process_id) {
         // An entry corresponding to |process_id| is missing. See above.
         break;
       }
     }
 
     if (TaskManagerObserver::IsResourceRefreshEnabled(REFRESH_TYPE_IDLE_WAKEUPS,
                                                       refresh_flags_)) {
+      DCHECK(callback_entry.second.on_idle_wakeups);
       callback_entry.second.on_idle_wakeups.Run(idle_wakeups_per_second);
     }
+
+    if (TaskManagerObserver::IsResourceRefreshEnabled(
+        REFRESH_TYPE_PHYSICAL_MEMORY, refresh_flags_)) {
+      DCHECK(callback_entry.second.on_physical_memory);
+      callback_entry.second.on_physical_memory.Run(physical_bytes);
+    }
   }
 
   // Reset refresh_results_ to trigger RefreshOnWorkerThread next time Refresh
   // is called.
   refresh_flags_ = 0;
 }
 
 }  // namespace task_manager