Fixing OOM crash in base::ProcessMetrics::GetWorkingSetKBytes

base/process/process_metrics_win.cc

Index: base/process/process_metrics_win.cc
diff --git a/base/process/process_metrics_win.cc b/base/process/process_metrics_win.cc
index fe4b84988f1a2736da6fef1b9a3c5ca0552f5ba0..a4f408f525f66ad8e8d2d9c6e47f2d50746f37bc 100644
--- a/base/process/process_metrics_win.cc
+++ b/base/process/process_metrics_win.cc
@@ -14,6 +14,7 @@
 
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
+#include "base/process/memory.h"
 #include "base/sys_info.h"
 
 namespace base {
@@ -140,6 +141,39 @@ void ProcessMetrics::GetCommittedKBytes(CommittedKBytes* usage) const {
   usage->priv = committed_private / 1024;
 }
 
+namespace {
+
+class WorkingSetInformationBuffer {

[brucedawson, 2016/09/08 20:25:27]

Needs to be tagged as no copy/assign.

[stanisc, 2016/09/08 23:10:14]

Done.

+ public:
+  WorkingSetInformationBuffer() {}
+  ~WorkingSetInformationBuffer() { Clear(); }
+
+  bool Reserve(size_t size) {
+    Clear();
+    // Use UncheckedMalloc here because this can be called from the code
+    // that handles low memory condition.
+    return base::UncheckedMalloc(size, &buffer_);

[Lei Zhang, 2016/09/08 21:19:56]

You can omit base:: in namespace base.

[stanisc, 2016/09/08 23:10:14]

Done.

+  }
+
+  PSAPI_WORKING_SET_INFORMATION* get() {
+    return reinterpret_cast<PSAPI_WORKING_SET_INFORMATION*>(buffer_);
+  }
+
+  const PSAPI_WORKING_SET_INFORMATION* operator ->() const {
+    return reinterpret_cast<const PSAPI_WORKING_SET_INFORMATION*>(buffer_);
+  }
+
+ private:
+  void Clear() {
+    free(buffer_);
+    buffer_ = nullptr;
+  }
+
+  void* buffer_ = nullptr;

[Lei Zhang, 2016/09/08 21:19:56]

If you'd like, you can also use a std::unique_ptr with base::FreeDeleter.

[Lei Zhang, 2016/09/08 21:19:57]

Can this be a PSAPI_WORKING_SET_INFORMATION* ?

[stanisc, 2016/09/08 23:10:14]

Yeah, I thought about that and even initially implemented this with unique_ptr
but it turned out to be slightly awkward because:
a) I'd need an extra intermediate pointer to pass to UncheckedMalloc before
assigning to unique_ptr
b) I want the old buffer to be deleted before allocating the new one to reduce
chances of OOM. That would require an extra .reset() call on unique_ptr which
might seem redundant.

[Lei Zhang, 2016/09/08 23:12:48]

Ack. Thanks for explaining.

+};
+
+}  // namespace
+
 bool ProcessMetrics::GetWorkingSetKBytes(WorkingSetKBytes* ws_usage) const {
   size_t ws_private = 0;
   size_t ws_shareable = 0;
@@ -149,40 +183,31 @@ bool ProcessMetrics::GetWorkingSetKBytes(WorkingSetKBytes* ws_usage) const {
   memset(ws_usage, 0, sizeof(*ws_usage));
 
   DWORD number_of_entries = 4096;  // Just a guess.
-  PSAPI_WORKING_SET_INFORMATION* buffer = NULL;
+  WorkingSetInformationBuffer buffer;
   int retries = 5;
   for (;;) {
+

[stanisc, 2016/09/07 21:41:05]

Will remove this empty line.

[stanisc, 2016/09/08 23:10:14]

Done.

     DWORD buffer_size = sizeof(PSAPI_WORKING_SET_INFORMATION) +
                         (number_of_entries * sizeof(PSAPI_WORKING_SET_BLOCK));
 
-    // if we can't expand the buffer, don't leak the previous
-    // contents or pass a NULL pointer to QueryWorkingSet
-    PSAPI_WORKING_SET_INFORMATION* new_buffer =
-        reinterpret_cast<PSAPI_WORKING_SET_INFORMATION*>(
-            realloc(buffer, buffer_size));
-    if (!new_buffer) {
-      free(buffer);
+    if (!buffer.Reserve(buffer_size))
       return false;
-    }
-    buffer = new_buffer;
 
     // Call the function once to get number of items
-    if (QueryWorkingSet(process_, buffer, buffer_size))
+    if (QueryWorkingSet(process_, buffer.get(), buffer_size))
       break;  // Success
 
-    if (GetLastError() != ERROR_BAD_LENGTH) {
-      free(buffer);
+    if (GetLastError() != ERROR_BAD_LENGTH)
       return false;
-    }
 
     number_of_entries = static_cast<DWORD>(buffer->NumberOfEntries);
 
     // Maybe some entries are being added right now. Increase the buffer to
     // take that into account.

[Lei Zhang, 2016/09/08 21:19:56]

Is it helpful to expand the comment to explain how you chose 1.1?

[stanisc, 2016/09/08 23:10:15]

Done.

-    number_of_entries = static_cast<DWORD>(number_of_entries * 1.25);
+    number_of_entries = static_cast<DWORD>(number_of_entries * 1.1);
 
     if (--retries == 0) {
-      free(buffer);  // If we're looping, eventually fail.
+      // If we're looping, eventually fail.
       return false;
     }
   }
@@ -205,7 +230,6 @@ bool ProcessMetrics::GetWorkingSetKBytes(WorkingSetKBytes* ws_usage) const {
   ws_usage->priv = ws_private * PAGESIZE_KB;
   ws_usage->shareable = ws_shareable * PAGESIZE_KB;
   ws_usage->shared = ws_shared * PAGESIZE_KB;
-  free(buffer);
   return true;
 }