Creates a memory-infra dump using the Chrome Remote Debugging protocol

memory-infra-dump-devtools.py

#!/usr/bin/env python
import argparse
import json
import os
import sys
import tempfile
import time

try:
  import websocket
except ImportError:
  print 'Please run: pip install --user websocket-client'
  sys.exit(1)


class TracingDevtoolsClient(object):
  def __init__(self, host, port):
    url = 'ws://%s:%s/devtools/browser' % (host, port)
    print 'Connecting to ' + url
    self.ws = websocket.create_connection(url);
    self.cmd = 0;

  def send(self, method, params={}):
    self.cmd += 1;
    self.ws.send(
      json.dumps({'id': self.cmd, 'method': method, 'params': params}))
    resp = self.recv()
    assert(resp['id'] == self.cmd)
    return resp.get('result', {})

  def recv(self):
    return json.loads(self.ws.recv())

  def dump(self, trace_fd):
    trace_config = {
        'excludedCategories': ['*'],
        'includedCategories': ['disabled-by-default-memory-infra'],
        'memoryDumpConfig': {'triggers': []}
    }
    print 'Starting trace with trace_config', trace_config
    params = {'traceConfig': trace_config, 'transferMode':'ReturnAsStream'}
    self.send('Tracing.start', params)
    print 'Requesting memory dump'
    resp = self.send('Tracing.requestMemoryDump')
    assert(resp['success'] == True)
    self.send('Tracing.end')

    # Wait for trace completion
    print 'Flushing trace'
    resp = self.recv()
    assert(resp['method'] == 'Tracing.tracingComplete')
    stream_handle = resp['params']['stream']

    # Read back the trace stream
    resp = {'eof': False}
    while not resp['eof']:
      resp = self.send('IO.read', {'handle': stream_handle})
      trace_fd.write(resp['data'].encode('utf-8'))

    self.send('IO.close', {'handle': stream_handle})
    trace_fd.close()


if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('--host', default='localhost')
  parser.add_argument('--port', '-p', default=9222)
  parser.add_argument('--output-trace', '-o', default=None)
  args = parser.parse_args()

  if args.output_trace is None:
    trace_fd = tempfile.NamedTemporaryFile(prefix='trace_', suffix='.json',
                                           delete=False)
  else:
    trace_fd = open(args.output_trace, 'wb')

  cli = TracingDevtoolsClient(args.host, args.port)
  cli.dump(trace_fd)
  print '\n%s: %d KB' % (trace_fd.name, os.stat(trace_fd.name).st_size / 1000)