Fixing OOM crash in base::ProcessMetrics::GetWorkingSetKBytes

base/process/process_metrics_win.cc

 // Copyright (c) 2013 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 "base/process/process_metrics.h"
 
 #include <windows.h>
 #include <psapi.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <winternl.h>
 
 #include <algorithm>
 
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
+#include "base/process/memory.h"
 #include "base/sys_info.h"
 
 namespace base {
 namespace {
 
 // System pagesize. This value remains constant on x86/64 architectures.
 const int PAGESIZE_KB = 4;
 
 typedef NTSTATUS(WINAPI* NTQUERYSYSTEMINFORMATION)(
     SYSTEM_INFORMATION_CLASS SystemInformationClass,
@@ skipping 106 matching lines @@
       usage->priv = 0;
       return;
     }
     base_address = new_base;
   }
   usage->image = committed_image / 1024;
   usage->mapped = committed_mapped / 1024;
   usage->priv = committed_private / 1024;
 }
 
+namespace {
+
+class WorkingSetInformationBuffer {
+ 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 UncheckedMalloc(size, reinterpret_cast<void**>(&buffer_));
+  }
+
+  PSAPI_WORKING_SET_INFORMATION* get() { return buffer_; }
+  const PSAPI_WORKING_SET_INFORMATION* operator ->() const { return buffer_; }
+
+ private:
+  void Clear() {
+    free(buffer_);
+    buffer_ = nullptr;
+  }
+
+  PSAPI_WORKING_SET_INFORMATION* buffer_ = nullptr;
+
+  DISALLOW_COPY_AND_ASSIGN(WorkingSetInformationBuffer);
+};
+
+}  // namespace
+
 bool ProcessMetrics::GetWorkingSetKBytes(WorkingSetKBytes* ws_usage) const {
   size_t ws_private = 0;
   size_t ws_shareable = 0;
   size_t ws_shared = 0;
 
   DCHECK(ws_usage);
   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 (;;) {
     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.
-    number_of_entries = static_cast<DWORD>(number_of_entries * 1.25);
+    // take that into account. Increasing by 10% should generally be enough,
+    // especially considering the potentially low memory condition during the
+    // call (when called from OomMemoryDetails) and the potentially high
+    // number of entries (300K was observed in crash dumps).
+    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;
     }
   }
 
   // On windows 2000 the function returns 1 even when the buffer is too small.
   // The number of entries that we are going to parse is the minimum between the
   // size we allocated and the real number of entries.
   number_of_entries =
       std::min(number_of_entries, static_cast<DWORD>(buffer->NumberOfEntries));
   for (unsigned int i = 0; i < number_of_entries; i++) {
     if (buffer->WorkingSetInfo[i].Shared) {
       ws_shareable++;
       if (buffer->WorkingSetInfo[i].ShareCount > 1)
         ws_shared++;
     } else {
       ws_private++;
     }
   }
 
   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;
 }
 
 static uint64_t FileTimeToUTC(const FILETIME& ftime) {
   LARGE_INTEGER li;
   li.LowPart = ftime.dwLowDateTime;
   li.HighPart = ftime.dwHighDateTime;
   return li.QuadPart;
 }
 
@@ skipping 76 matching lines @@
 
   meminfo->total = mem_status.ullTotalPhys / 1024;
   meminfo->free = mem_status.ullAvailPhys / 1024;
   meminfo->swap_total = mem_status.ullTotalPageFile / 1024;
   meminfo->swap_free = mem_status.ullAvailPageFile / 1024;
 
   return true;
 }
 
 }  // namespace base