Add a PartialPreReadImage function to file_util (on Windows).

base/file_util_win.cc

 // Copyright (c) 2012 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/file_util.h"
 
 #include <windows.h>
 #include <propvarutil.h>
 #include <psapi.h>
 #include <shellapi.h>
 #include <shlobj.h>
 #include <time.h>
 
+#include <algorithm>
 #include <limits>
 #include <string>
 
 #include "base/file_path.h"
 #include "base/logging.h"
+#include "base/memory/scoped_ptr.h"
 #include "base/metrics/histogram.h"
 #include "base/string_number_conversions.h"
 #include "base/string_util.h"
 #include "base/threading/thread_restrictions.h"
 #include "base/time.h"
 #include "base/utf_string_conversions.h"
 #include "base/win/pe_image.h"
 #include "base/win/scoped_comptr.h"
 #include "base/win/scoped_handle.h"
 #include "base/win/win_util.h"
@@ skipping 1110 matching lines @@
     size_t total_read = 0;
     while (::ReadFile(file, buffer, actual_step_size, &len, NULL) &&
            len > 0 &&
            (size_to_read ? total_read < size_to_read : true)) {
       total_read += static_cast<size_t>(len);
     }
     ::VirtualFree(buffer, 0, MEM_RELEASE);
   } else {
     // WinXP branch. Here, reading the DLL from disk doesn't do
     // what we want so instead we pull the pages into memory by loading
     // the DLL and touching pages at a stride.

[robertshield, 2012/01/26 22:43:44]

Consider adding a note indicating that we intentionally disregard step_size in
this branch.

[Roger McFarlane (Google), 2012/01/27 05:24:09]

Done.

     HMODULE dll_module = ::LoadLibraryExW(
         file_path,
         NULL,
         LOAD_WITH_ALTERED_SEARCH_PATH | DONT_RESOLVE_DLL_REFERENCES);
 
     if (!dll_module)
       return false;
 
     base::win::PEImage pe_image(dll_module);
     PIMAGE_NT_HEADERS nt_headers = pe_image.GetNTHeaders();
+    SYSTEM_INFO system_info;
+    GetSystemInfo(&system_info);
+
     size_t actual_size_to_read = size_to_read ? size_to_read :
                                  nt_headers->OptionalHeader.SizeOfImage;
     volatile uint8* touch = reinterpret_cast<uint8*>(dll_module);
     size_t offset = 0;
     while (offset < actual_size_to_read) {
       uint8 unused = *(touch + offset);
-      offset += step_size;
+      offset += system_info.dwPageSize;
     }
+    uint8 unused2 = *(touch + size_to_read - 1);
     FreeLibrary(dll_module);
   }
 
   return true;
 }
 
+namespace {
+
+struct ScopedPtrVirtualFree {
+  void operator() (void* ptr) {
+    ::VirtualFree(ptr, 0, MEM_RELEASE);
+  }
+};
+
+bool SetFilePointer(HANDLE file_handle, size_t position) {
+  return position <= std::numeric_limits<LONG>::max() &&
+      ::SetFilePointer(file_handle,
+                       static_cast<LONG>(position),
+                       NULL,
+                       FILE_BEGIN) != INVALID_SET_FILE_POINTER;
+}
+
+bool ReadNextBytes(HANDLE file_handle, void* buffer, size_t bytes_to_read) {
+  DWORD bytes_read = 0;
+  return bytes_to_read <= std::numeric_limits<DWORD>::max() &&
+      ::ReadFile(file_handle,
+                 buffer,
+                 static_cast<DWORD>(bytes_to_read),
+                 &bytes_read,
+                 NULL) &&
+      bytes_read == bytes_to_read;
+}
+
+}  // namespace
+
+namespace internal {
+
+BASE_EXPORT bool PartialPreReadImageOnDisk(const wchar_t* file_path,
+                                           size_t percentage,
+                                           size_t max_chunk_size) {
+  // TODO(rogerm): change this to be driven by a table within the PE file
+  //     (defaulting to full read if it's not there?) that's initialized
+  //     by the optimization toolchain.
+  DCHECK(file_path != NULL);
+
+  if (percentage == 0)
+    return true;
+
+  const size_t kOneHundredPercent = 100;
+  if (percentage > kOneHundredPercent)
+    percentage = kOneHundredPercent;
+
+  // Validate/setup the percentage and max_chunk_size.
+  const size_t kMinChunkSize = 1024 * 1024;
+  percentage = std::min(percentage, kOneHundredPercent);

[robertshield, 2012/01/26 22:43:44]

This line is redundant with lines 1223-1224.

[Roger McFarlane (Google), 2012/01/27 05:24:09]

Oops, was supposed to have been deleted.

+  max_chunk_size = std::max(max_chunk_size, kMinChunkSize);
+
+  // Open the file.
+  base::win::ScopedHandle file(
+      CreateFile(file_path,
+                 GENERIC_READ,
+                 FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
+                 NULL,
+                 OPEN_EXISTING,
+                 FILE_FLAG_SEQUENTIAL_SCAN,
+                 NULL));
+
+  if (!file.IsValid())
+    return false;
+
+  // Allocate a resizable buffer for the headers. We initially reserve as much
+  // space as we typically see as the header size for chrome.dll and other
+  // PE images.
+  const size_t kInitialBufferSize = 0x400;
+  std::vector<uint8> headers(kInitialBufferSize);
+
+  // Read, hopefully, all of the headers.
+  if (!ReadNextBytes(file, &headers[0], headers.size()))
+    return false;
+
+  // The DOS headers start at offset 0, which allow us to get the offset of the
+  // NT headers. Let's ensure we've read enough to capture the NT headers.
+  size_t nt_headers_start =
+      reinterpret_cast<IMAGE_DOS_HEADER*>(&headers[0])->e_lfanew;
+  size_t nt_headers_end = nt_headers_start + sizeof(IMAGE_NT_HEADERS);
+  if (nt_headers_end > headers.size()) {
+    size_t current_length = headers.size();
+    size_t bytes_to_read = nt_headers_end - current_length;
+    headers.resize(nt_headers_end);
+    if (!ReadNextBytes(file, &headers[current_length], bytes_to_read))
+      return false;
+  }
+
+  // Now that we've got the NT headers we can get the total header size,
+  // including all of the section headers. Let's ensure we've read enough
+  // to capture all of the section headers.
+  size_t size_of_headers = reinterpret_cast<IMAGE_NT_HEADERS*>(
+      &headers[nt_headers_start])->OptionalHeader.SizeOfHeaders;
+  if (size_of_headers > headers.size()) {
+    size_t current_length = headers.size();
+    size_t bytes_to_read = size_of_headers - current_length;
+    headers.resize(size_of_headers);
+    if (!ReadNextBytes(file, &headers[current_length], bytes_to_read))
+      return false;
+  }
+
+  // Now we have all of the headers. This is enough to let us use the PEImage
+  // wrapper to query the structure of the image.
+  base::win::PEImage pe_image(reinterpret_cast<HMODULE>(&headers[0]));
+  CHECK(pe_image.VerifyMagic());
+
+  // Allocate a buffer to hold the pre-read bytes.
+  scoped_ptr_malloc<uint8, ScopedPtrVirtualFree> buffer(
+      reinterpret_cast<uint8*>(
+          ::VirtualAlloc(NULL, max_chunk_size, MEM_COMMIT, PAGE_READWRITE)));
+  if (buffer.get() == NULL)
+    return false;
+
+  // Iterate over each section, reading in a percentage of each.
+  for (UINT i = 0; i < pe_image.GetNumSections(); ++i) {
+    // Get the current section and figure out how much of it to touch.
+    const IMAGE_SECTION_HEADER* section = pe_image.GetSectionHeader(i);
+    CHECK(section != NULL);
+    CHECK_LE(reinterpret_cast<const uint8*>(section + 1),
+             &headers[0] + headers.size());
+
+    if (!SetFilePointer(file, section->PointerToRawData))
+      return false;
+
+    size_t max_bytes =
+        (section->SizeOfRawData * percentage) / kOneHundredPercent;
+    while (max_bytes > 0) {
+      size_t chunk_size = std::min(max_bytes, max_chunk_size);
+      if (!ReadNextBytes(file, buffer.get(), chunk_size))
+        return false;
+      max_bytes -= chunk_size;
+    }
+  }
+
+  // We're done.
+  return true;
+}
+
+BASE_EXPORT bool PartialPreReadImageInMemory(const wchar_t* file_path,
+                                             size_t percentage) {
+  // TODO(rogerm): change this to be driven by a table within the PE image
+  //     (defaulting to full read if it's not there?) that's initialized
+  //     by the optimization toolchain.
+  DCHECK(file_path != NULL);
+
+  if (percentage == 0)
+    return true;
+
+  const size_t kOneHundredPercent = 100;

[robertshield, 2012/01/26 22:43:44]

consider pulling this up out of the function since it is used in a couple of
places

[Roger McFarlane (Google), 2012/01/27 05:24:09]

Done.

+  if (percentage > kOneHundredPercent)
+    percentage = kOneHundredPercent;
+
+  HMODULE dll_module = ::LoadLibraryExW(
+      file_path,
+      NULL,
+      LOAD_WITH_ALTERED_SEARCH_PATH | DONT_RESOLVE_DLL_REFERENCES);
+
+  if (!dll_module)
+    return false;
+
+  base::win::PEImage pe_image(dll_module);
+  CHECK(pe_image.VerifyMagic());
+
+  // Iterate over each section, stepping through a percentage of each to page
+  // it in off the disk.
+  for (UINT i = 0; i < pe_image.GetNumSections(); ++i) {
+    // Get the current section and figure out how much of it to touch. We
+    // start from the beginning of the section, and step through all of the
+    // initialized (on disk) data in the section; this is the lesser of the
+    // size of the raw data in the section and the virtual size of the section
+    // because sections can have the tails implicitly initialized to zero (i.e,
+    // their virtual size is larger than the raw size) and the raw data is
+    // padded to the PE page size if the entire section is initialized (i.e.,
+    // raw data size will be larger than the virtual size).
+    const IMAGE_SECTION_HEADER* section = pe_image.GetSectionHeader(i);
+    DCHECK(section != NULL);
+    volatile uint8* touch =
+        reinterpret_cast<uint8*>(pe_image.RVAToAddr(section->VirtualAddress));
+    size_t section_length =
+        std::min(section->SizeOfRawData, section->Misc.VirtualSize);
+    volatile uint8* max_touch =
+        touch + ((section_length * percentage) / kOneHundredPercent) - 1;
+
+    // Verify that the extent we're going to touch falls inside the section
+    // we expect it to (and by implication, inside the pe_image).
+    DCHECK_EQ(section,
+              pe_image.GetImageSectionFromAddr(const_cast<uint8*>(touch)));
+    DCHECK_EQ(section,
+              pe_image.GetImageSectionFromAddr(const_cast<uint8*>(max_touch)));
+
+    // Get the system info so we know the page size.
+    SYSTEM_INFO system_info;
+    GetSystemInfo(&system_info);
+
+    // Read the memory in the range [touch, max_touch] with a stride of
+    // the system page size. This ensures that it's been paged in.
+    while (touch < max_touch) {
+      uint8 unused1 = *touch;
+      touch += system_info.dwPageSize;
+    }
+    uint8 unused2 = *max_touch;
+  }
+
+  FreeLibrary(dll_module);
+
+  return true;
+}
+
+}  // namespace file_util::internal
+
+bool PartialPreReadImage(const wchar_t* file_path,
+                         size_t percentage,
+                         size_t max_chunk_size) {
+  base::ThreadRestrictions::AssertIOAllowed();
+
+  if (percentage >= 100) {
+    // If we're reading the whole image, we don't need to parse headers and
+    // navigate sections, the basic PreReadImage() can be used to just step
+    // blindly through the entire file / address-space.
+    return PreReadImage(file_path, 0, max_chunk_size);
+  }
+
+  using internal::PartialPreReadImageInMemory;
+  using internal::PartialPreReadImageOnDisk;
+
+  if (base::win::GetVersion() > base::win::VERSION_XP) {
+    // Vista+ branch. On these OSes, we warm up the Image by reading its
+    // file off the disk.
+    return PartialPreReadImageOnDisk(file_path, percentage, max_chunk_size);
+  }
+
+  // WinXP branch. For XP, reading the image from disk doesn't do what we want
+  // so instead we pull the pages into memory by loading the DLL and touching
+  // initialized pages at a stride.
+  return PartialPreReadImageInMemory(file_path, percentage);
+}
+
 }  // namespace file_util