[Telemetry] Reduce memory usage in presence of TraceValues

Previously, creating a TraceValue would store the associated trace_data in
memory and defer serialization until needed. In this CL, the trace_data is
serialized right away into a temporary file to prevent it from being stored in
memory.

BUG=457385

Review URL: https://codereview.chromium.org/920523002

Cr-Commit-Position: refs/heads/master@{#317106}

tools/telemetry/telemetry/benchmark.py

@@ -189,22 +189,22 @@ class Benchmark(command_line.Command):
     self._DownloadGeneratedProfileArchive(finder_options)
 
     benchmark_metadata = self.GetMetadata()
-    results = results_options.CreateResults(
-        benchmark_metadata, finder_options, self.ValueCanBeAddedPredicate)
-    try:
-      user_story_runner.Run(pt, us, expectations, finder_options, results,
-                            max_failures=self._max_failures)
-      return_code = min(254, len(results.failures))
-    except Exception:
-      exception_formatter.PrintFormattedException()
-      return_code = 255
-
-    bucket = cloud_storage.BUCKET_ALIASES[finder_options.upload_bucket]
-    if finder_options.upload_results:
-      results.UploadTraceFilesToCloud(bucket)
-      results.UploadProfilingFilesToCloud(bucket)
-
-    results.PrintSummary()
+    with results_options.CreateResults(benchmark_metadata,
+                                       finder_options) as results:
+      try:
+        user_story_runner.Run(pt, us, expectations, finder_options, results,
+                              max_failures=self._max_failures)
+        return_code = min(254, len(results.failures))
+      except Exception:
+        exception_formatter.PrintFormattedException()
+        return_code = 255
+
+      bucket = cloud_storage.BUCKET_ALIASES[finder_options.upload_bucket]
+      if finder_options.upload_results:
+        results.UploadTraceFilesToCloud(bucket)
+        results.UploadProfilingFilesToCloud(bucket)
+
+      results.PrintSummary()
     return return_code
 
   def _DownloadGeneratedProfileArchive(self, options):

tools/telemetry/telemetry/results/page_test_results.py

@@ -134,6 +134,20 @@ class PageTestResults(object):
   def _GetStringFromExcInfo(self, err):
     return ''.join(traceback.format_exception(*err))
 
+  def CleanUp(self):
+    """Clean up any TraceValues contained within this results object."""
+    for run in self._all_page_runs:
+      for v in run.values:
+        if isinstance(v, trace.TraceValue):
+          v.CleanUp()
+          run.values.remove(v)
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, _, __, ___):
+    self.CleanUp()
+
   def WillRunPage(self, page):
     assert not self._current_page_run, 'Did not call DidRunPage.'
     self._current_page_run = page_run.PageRun(page)

tools/telemetry/telemetry/results/page_test_results_unittest.py

@@ -15,7 +15,6 @@ from telemetry.value import scalar
 from telemetry.value import skip
 from telemetry.value import trace
 
-
 class PageTestResultsTest(base_test_results_unittest.BaseTestResultsUnittest):
   def setUp(self):
     ps = page_set.PageSet(file_path=os.path.dirname(__file__))
@@ -234,3 +233,36 @@ class PageTestResultsTest(base_test_results_unittest.BaseTestResultsUnittest):
 
     values = results.FindAllTraceValues()
     self.assertEquals(2, len(values))
+
+  def testCleanUpCleansUpTraceValues(self):
+    results = page_test_results.PageTestResults()
+    v0 = trace.TraceValue(None, trace_data.TraceData({'test': 1}))
+    v1 = trace.TraceValue(None, trace_data.TraceData({'test': 2}))
+
+    results.WillRunPage(self.pages[0])
+    results.AddValue(v0)
+    results.DidRunPage(self.pages[0])
+
+    results.WillRunPage(self.pages[1])
+    results.AddValue(v1)
+    results.DidRunPage(self.pages[1])
+
+    results.CleanUp()
+    self.assertTrue(v0.cleaned_up)
+    self.assertTrue(v1.cleaned_up)
+
+  def testNoTracesLeftAfterCleanUp(self):
+    results = page_test_results.PageTestResults()
+    v0 = trace.TraceValue(None, trace_data.TraceData({'test': 1}))
+    v1 = trace.TraceValue(None, trace_data.TraceData({'test': 2}))
+
+    results.WillRunPage(self.pages[0])
+    results.AddValue(v0)
+    results.DidRunPage(self.pages[0])
+
+    results.WillRunPage(self.pages[1])
+    results.AddValue(v1)
+    results.DidRunPage(self.pages[1])
+
+    results.CleanUp()
+    self.assertFalse(results.FindAllTraceValues())

tools/telemetry/telemetry/value/trace.py

@@ -18,8 +18,7 @@ import telemetry.web_components # pylint: disable=W0611
 from trace_viewer.build import trace2html
 
 class TraceValue(value_module.Value):
-  def __init__(self, page, trace_data, important=False,
-               description=None):
+  def __init__(self, page, trace_data, important=False, description=None):
     """A value that contains a TraceData object and knows how to
     output it.
 
@@ -30,18 +29,18 @@ class TraceValue(value_module.Value):
     super(TraceValue, self).__init__(
         page, name='trace', units='', important=important,
         description=description)
-    self._trace_data = trace_data
+    self._temp_file = self._GetTempFileHandle(trace_data)
     self._cloud_url = None
     self._serialized_file_handle = None
 
-  def _GetTempFileHandle(self):
+  def _GetTempFileHandle(self, trace_data):
     tf = tempfile.NamedTemporaryFile(delete=False, suffix='.html')
     if self.page:
       title = self.page.display_name
     else:
       title = ''
     trace2html.WriteHTMLForTraceDataToFile(
-        [self._trace_data.GetEventsFor(trace_data_module.CHROME_TRACE_PART)],
+        [trace_data.GetEventsFor(trace_data_module.CHROME_TRACE_PART)],
         title,
         tf)
     tf.close()
@@ -54,6 +53,27 @@ class TraceValue(value_module.Value):
       page_name = None
     return 'TraceValue(%s, %s)' % (page_name, self.name)
 
+  def CleanUp(self):
+    """Cleans up tempfile after it is no longer needed.
+
+    A cleaned up TraceValue cannot be used for further operations. CleanUp()
+    may be called more than once without error.
+    """
+    if self._temp_file is None:
+      return
+    os.remove(self._temp_file.GetAbsPath())
+    self._temp_file = None
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, _, __, ___):
+    self.CleanUp()
+
+  @property
+  def cleaned_up(self):
+    return self._temp_file is None
+
   def GetBuildbotDataType(self, output_context):
     return None
 
@@ -84,6 +104,8 @@ class TraceValue(value_module.Value):
     return None
 
   def AsDict(self):
+    if self._temp_file is None:
+      raise ValueError('Tried to serialize TraceValue without tempfile.')
     d = super(TraceValue, self).AsDict()
     if self._serialized_file_handle:
       d['file_id'] = self._serialized_file_handle.id
@@ -92,21 +114,22 @@ class TraceValue(value_module.Value):
     return d
 
   def Serialize(self, dir_path):
-    fh = self._GetTempFileHandle()
-    file_name = str(fh.id) + fh.extension
+    if self._temp_file is None:
+      raise ValueError('Tried to serialize nonexistent trace.')
+    file_name = str(self._temp_file.id) + self._temp_file.extension
     file_path = os.path.abspath(os.path.join(dir_path, file_name))
-    shutil.move(fh.GetAbsPath(), file_path)
+    shutil.copy(self._temp_file.GetAbsPath(), file_path)
     self._serialized_file_handle = file_handle.FromFilePath(file_path)
     return self._serialized_file_handle
 
   def UploadToCloud(self, bucket):
-    temp_fh = None
+    if self._temp_file is None:
+      raise ValueError('Tried to upload nonexistent trace to Cloud Storage.')
     try:
       if self._serialized_file_handle:
         fh = self._serialized_file_handle
       else:
-        temp_fh = self._GetTempFileHandle()
-        fh = temp_fh
+        fh = self._temp_file
       remote_path = ('trace-file-id_%s-%s-%d%s' % (
           fh.id,
           datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'),
@@ -121,6 +144,3 @@ class TraceValue(value_module.Value):
     except cloud_storage.PermissionError as e:
       logging.error('Cannot upload trace files to cloud storage due to '
                     ' permission error: %s' % e.message)
-    finally:
-      if temp_fh:
-        os.remove(temp_fh.GetAbsPath())

tools/telemetry/telemetry/value/trace_unittest.py

@@ -106,23 +106,14 @@ class NoLeakedTempfilesTests(TestBase):
     self.actual_tempdir = trace.tempfile.tempdir
     trace.tempfile.tempdir = self.temp_test_dir
 
-  def testNoLeakedTempFileWhenTraceSerialize(self):
-    tempdir = tempfile.mkdtemp()
-    v = trace.TraceValue(None, trace_data.TraceData({'test': 1}))
-    fh = v.Serialize(tempdir)
-    try:
-      shutil.rmtree(fh.GetAbsPath(), ignore_errors=True)
-      self.assertTrue(tempdir)
-    finally:
-      shutil.rmtree(tempdir)
-      self.assertTrue(_IsEmptyDir(self.temp_test_dir))
+  def testNoLeakedTempFileOnImplicitCleanUp(self):
+    with trace.TraceValue(None, trace_data.TraceData({'test': 1})):
+      pass
+    self.assertTrue(_IsEmptyDir(self.temp_test_dir))
 
   def testNoLeakedTempFileWhenUploadingTrace(self):
     v = trace.TraceValue(None, trace_data.TraceData({'test': 1}))
-    trace.cloud_storage.SetCalculatedHashesForTesting(
-        TestDefaultDict(123))
-    bucket = trace.cloud_storage.PUBLIC_BUCKET
-    v.UploadToCloud(bucket)
+    v.CleanUp()
     self.assertTrue(_IsEmptyDir(self.temp_test_dir))
 
   def tearDown(self):