sandwich: Fixes two sources of KeyError task failures

tools/android/loading/sandwich_prefetch.py

 # 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.
 
 """
 Implements a task builder for benchmarking effects of NoState Prefetch.
 Noticeable steps of the task pipeline:
   * Save a WPR archive
   * Process the WPR archive to make all resources cacheable
   * Process cache archive to patch response headers back to their original
       values.
   * Find out which resources are discoverable by NoState Prefetch
       (HTMLPreloadScanner)
   * Load pages with empty/full/prefetched cache
   * Extract most important metrics to a CSV
 """
 
 import csv
 import logging
 import json
 import os
 import re
 import shutil
-from urlparse import urlparse
+import urlparse
 
 import chrome_cache
 import common_util
 import loading_trace
 from prefetch_view import PrefetchSimulationView
 from request_dependencies_lens import RequestDependencyLens
 import sandwich_metrics
 import sandwich_runner
+import sandwich_utils
 import task_manager
 import wpr_backend
 
 
 class Discoverer(object):
   # Do not prefetch anything.
   EmptyCache = 'empty-cache'
 
   # Prefetches everything to load fully from cache (impossible in practice).
   FullCache = 'full-cache'
@@ skipping 19 matching lines @@
   Discoverer.MainDocument,
   Discoverer.Parser,
   Discoverer.HTMLPreloadScanner,
   Discoverer.HTMLPreloadScannerStore,
 ])
 
 
 _UPLOAD_DATA_STREAM_REQUESTS_REGEX = re.compile(r'^\d+/(?P<url>.*)$')
 
 
+def _NormalizeUrl(url):
+  """Returns normalized URL such as removing trailing slashes."""
+  parsed_url = list(urlparse.urlparse(url))
+  parsed_url[2] = re.sub(r'/{2,}', r'/', parsed_url[2])
+  return urlparse.urlunparse(parsed_url)
+
+
 def _PatchWpr(wpr_archive):
   """Patches a WPR archive to get all resources into the HTTP cache and avoid
   invalidation and revalidations.
 
   Args:
     wpr_archive: wpr_backend.WprArchiveBackend WPR archive to patch.
   """
   # Sets the resources cache max-age to 10 years.
   MAX_AGE = 10 * 365 * 24 * 60 * 60
   CACHE_CONTROL = 'public, max-age={}'.format(MAX_AGE)
@@ skipping 112 matching lines @@
   logging.info('number of requests discovered by %s: %d',
       subresource_discoverer, len(requests))
   return requests
 
 
 def _PruneOutOriginalNoStoreRequests(original_headers_path, requests):
   with open(original_headers_path) as file_input:
     original_headers = json.load(file_input)
   pruned_requests = set()
   for request in requests:
-    request_original_headers = original_headers[request.url]
+    url = _NormalizeUrl(request.url)
+    if url not in original_headers:
+      # TODO(gabadie): Investigate why these requests were not in WPR.
+      assert request.failed
+      logging.warning(
+          'could not find original headers for: %s (failure: %s)',
+          url, request.error_text)
+      continue
+    request_original_headers = original_headers[url]
     if ('cache-control' in request_original_headers and
         'no-store' in request_original_headers['cache-control'].lower()):
       pruned_requests.add(request)
   return [r for r in requests if r not in pruned_requests]
 
 
 def _ExtractDiscoverableUrls(
     original_headers_path, loading_trace_path, subresource_discoverer):
   """Extracts discoverable resource urls from a loading trace according to a
   sub-resource discoverer.
@@ skipping 147 matching lines @@
   def VerifyWprLog(self, wpr_log_path):
     """Verifies WPR log with previously verified traces."""
     all_wpr_requests = wpr_backend.ExtractRequestsFromLog(wpr_log_path)
     all_wpr_urls = set()
     unserved_wpr_urls = set()
     wpr_command_colliding_urls = set()
 
     for request in all_wpr_requests:
       if request.is_wpr_host:
         continue
-      if urlparse(request.url).path.startswith('/web-page-replay'):
+      if urlparse.urlparse(request.url).path.startswith('/web-page-replay'):
         wpr_command_colliding_urls.add(request.url)
       elif request.is_served is False:
         unserved_wpr_urls.add(request.url)
       all_wpr_urls.add(request.url)
 
     _PrintUrlSetComparison(set(), unserved_wpr_urls,
                            'Distinct unserved resources from WPR')
     _PrintUrlSetComparison(set(), wpr_command_colliding_urls,
                            'Distinct resources colliding to WPR commands')
     _PrintUrlSetComparison(all_wpr_urls, self._all_sent_url_requests,
@@ skipping 91 matching lines @@
     logging.info('loading trace: %s', trace_path)
     trace = loading_trace.LoadingTrace.FromJsonFile(trace_path)
 
     logging.info('verifying trace: %s', trace_path)
     run_output_verifier.VerifyTrace(trace)
 
     logging.info('extracting metrics from trace: %s', trace_path)
     served_from_network_bytes = 0
     served_from_cache_bytes = 0
     urls_hitting_network = set()
+    response_sizes = {}
     for request in _FilterOutDataAndIncompleteRequests(
         trace.request_track.GetEvents()):
       # Ignore requests served from the blink's cache.
       if request.served_from_cache:
         continue
       urls_hitting_network.add(request.url)
       if request.from_disk_cache:
-        served_from_cache_bytes += cached_encoded_data_lengths[request.url]
+        if request.url in cached_encoded_data_lengths:
+          response_size = cached_encoded_data_lengths[request.url]
+        else:
+          # Some fat webpages may overflow the Memory cache, and so some
+          # requests might be served from disk cache couple of times per page
+          # load.
+          logging.warning('Looks like could be served from memory cache: %s',
+              request.url)
+          response_size = response_sizes[request.url]
+        served_from_cache_bytes += response_size
       else:
-        served_from_network_bytes += request.GetEncodedDataLength()
+        response_size = request.GetEncodedDataLength()
+        served_from_network_bytes += response_size
+      response_sizes[request.url] = response_size
 
     # Make sure the served from blink's cache requests have at least one
     # corresponding request that was not served from the blink's cache.
     for request in _FilterOutDataAndIncompleteRequests(
         trace.request_track.GetEvents()):
       assert (request.url in urls_hitting_network or
               not request.served_from_cache)
 
     run_metrics = {
         'url': trace.url,
@@ skipping 63 matching lines @@
                        dependencies=[self._common_builder.original_wpr_task])
     def BuildPatchedWpr():
       common_util.EnsureParentDirectoryExists(BuildPatchedWpr.path)
       shutil.copyfile(
           self._common_builder.original_wpr_task.path, BuildPatchedWpr.path)
       wpr_archive = wpr_backend.WprArchiveBackend(BuildPatchedWpr.path)
 
       # Save up original response headers.
       original_response_headers = {e.url: e.GetResponseHeadersDict() \
           for e in wpr_archive.ListUrlEntries()}
+      logging.info('save up response headers for %d resources',
+                   len(original_response_headers))
+      if not original_response_headers:
+        # TODO(gabadie): How is it possible to not even have the main resource
+        # in the WPR archive? Example URL can be found in:
+        # http://crbug.com/623966#c5
+        raise Exception(
+            'Looks like no resources were recorded in WPR during: {}'.format(
+                self._common_builder.original_wpr_task.name))
       with open(self._original_headers_path, 'w') as file_output:
         json.dump(original_response_headers, file_output)
 
       # Patch WPR.
       _PatchWpr(wpr_archive)
       wpr_archive.Persist()
 
     @self.RegisterTask('common/original-cache.zip', [BuildPatchedWpr])
     def BuildOriginalCache():
       runner = self._common_builder.CreateSandwichRunner()
@@ skipping 111 matching lines @@
       run_metrics_list = _ProcessRunOutputDir(
           cache_validation_result, benchmark_setup, RunBenchmark.path)
       with open(ProcessRunOutputDir.path, 'w') as csv_file:
         writer = csv.DictWriter(csv_file, fieldnames=(additional_column_names +
                                     sandwich_metrics.COMMON_CSV_COLUMN_NAMES))
         writer.writeheader()
         for trace_metrics in run_metrics_list:
           writer.writerow(trace_metrics)
 
     self._common_builder.default_final_tasks.append(ProcessRunOutputDir)