Add JS heap snapshotting infrastructure to Telemetry (part 1)

tools/telemetry/telemetry/core/backends/chrome/inspector_backend.py

 # Copyright 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.
 import json
 import logging
 import socket
 import sys
 
 from telemetry.core import util
 from telemetry.core import exceptions
 from telemetry.core.backends import png_bitmap
 from telemetry.core.backends.chrome import inspector_console
 from telemetry.core.backends.chrome import inspector_memory
 from telemetry.core.backends.chrome import inspector_network
 from telemetry.core.backends.chrome import inspector_page
 from telemetry.core.backends.chrome import inspector_runtime
 from telemetry.core.backends.chrome import inspector_timeline
 from telemetry.core.backends.chrome import websocket
+from telemetry.core.jsheap import model
 
 class InspectorException(Exception):
   pass
 
 class InspectorBackend(object):
   def __init__(self, browser, browser_backend, debugger_url):
     assert debugger_url
     self._browser = browser
     self._browser_backend = browser_backend
     self._debugger_url = debugger_url
@@ skipping 258 matching lines @@
           raise util.TimeoutException(
             'Timed out waiting for reply. This is unusual.')
         raise exceptions.TabCrashException(sys.exc_info()[1])
 
       res = json.loads(data)
       logging.debug('got [%s]', data)
       if 'method' in res:
         self._HandleNotification(res)
         continue
 
-      if res['id'] != req['id']:
+      if 'id' not in res or res['id'] != req['id']:
         logging.debug('Dropped reply: %s', json.dumps(res))
         continue
       return res
 
   def RegisterDomain(self,
       domain_name, notification_handler, will_close_handler):
     """Registers a given domain for handling notification methods.
 
     For example, given inspector_backend:
        def OnConsoleNotification(msg):
           if msg['method'] == 'Console.messageAdded':
              print msg['params']['message']
           return
        def OnConsoleClose(self):
           pass
        inspector_backend.RegisterDomain('Console',
                                         OnConsoleNotification, OnConsoleClose)
        """
     assert domain_name not in self._domain_handlers
     self._domain_handlers[domain_name] = (notification_handler,
                                           will_close_handler)
 
   def UnregisterDomain(self, domain_name):
     """Unregisters a previously registered domain."""
     assert domain_name in self._domain_handlers
     self._domain_handlers.pop(domain_name)
 
   def CollectGarbage(self):
     self._page.CollectGarbage()
+
+  def TakeJSHeapSnapshot(self, timeout=120):
+    # This is a hack to make the nested function be able to modify the
+    # variables.
+    snapshot_uid = [0]
+    snapshot = [[]]
+
+    def OnNotification(res):
+      if res['method'] == 'HeapProfiler.addProfileHeader':
+        snapshot_uid[0] = res['params']['header']['uid']
+      elif res['method'] == 'HeapProfiler.addHeapSnapshotChunk':
+        snapshot[0].append(res['params']['chunk'])
+      elif res['method'] == 'HeapProfiler.finishHeapSnapshot':
+        snapshot[0] = ''.join(snapshot[0])
+
+    def OnClose():
+      pass
+
+    self.RegisterDomain('HeapProfiler', OnNotification, OnClose)
+
+    self.SyncRequest({'method': 'Page.getResourceTree'}, timeout)
+    self.SyncRequest({'method': 'Debugger.enable'}, timeout)
+    self.SyncRequest({'method': 'HeapProfiler.clearProfiles'}, timeout)
+    self.SyncRequest({'method': 'HeapProfiler.takeHeapSnapshot',
+                      'params': {'detailed': True}}, timeout)
+    self.SyncRequest({'method': 'HeapProfiler.getHeapSnapshot',
+                      'params': {'uid': snapshot_uid[0]}}, timeout)
+
+    self.UnregisterDomain('HeapProfiler')
+
+    return model.JsHeapSnapshotModel(snapshot[0])