chromecast/tools/trace.py

   1 #!/usr/bin/env python
   2 #
   3 # Copyright 2015 The Chromium Authors. All rights reserved.
   4 # Use of this source code is governed by a BSD-style license that can be
   5 # found in the LICENSE file.
   6 #
   7 # This script was originally written by Alok Priyadarshi (alokp@)
   8 # with some minor local modifications.
   9
  10 import contextlib
  11 import json
  12 import logging
  13 import math
  14 import optparse
  15 import os
  16 import sys
  17 import websocket
  18
  19
  20 class TracingClient(object):
  21   def BufferUsage(self, buffer_usage):
  22     percent = int(math.floor(buffer_usage * 100))
  23     logging.debug('Buffer Usage: %i', percent)
  24
  25
  26 class TracingBackend(object):
  27   def __init__(self, devtools_port):
  28     self._socket = None
  29     self._next_request_id = 0
  30     self._tracing_client = None
  31     self._tracing_data = []
  32
  33   def Connect(self, device_ip, devtools_port, timeout=10):
  34     assert not self._socket
  35     url = 'ws://%s:%i/devtools/browser' % (device_ip, devtools_port)
  36     print('Connect to %s ...' % url)
  37     self._socket = websocket.create_connection(url, timeout=timeout)
  38     self._next_request_id = 0
  39
  40   def Disconnect(self):
  41     if self._socket:
  42       self._socket.close()
  43       self._socket = None
  44
  45   def StartTracing(self,
  46                    tracing_client=None,
  47                    custom_categories=None,
  48                    record_continuously=False,
  49                    buffer_usage_reporting_interval=0,
  50                    timeout=10):
  51     self._tracing_client = tracing_client
  52     self._socket.settimeout(timeout)
  53     req = {
  54       'method': 'Tracing.start',
  55       'params': {
  56         'categories': custom_categories,
  57         'bufferUsageReportingInterval': buffer_usage_reporting_interval,
  58         'options': 'record-continuously' if record_continuously else
  59                    'record-until-full'
  60       }
  61     }
  62     self._SendRequest(req)
  63
  64   def StopTracing(self, timeout=30):
  65     self._socket.settimeout(timeout)
  66     req = {'method': 'Tracing.end'}
  67     self._SendRequest(req)
  68     while self._socket:
  69       res = self._ReceiveResponse()
  70       if 'method' in res and self._HandleResponse(res):
  71         self._tracing_client = None
  72         result = self._tracing_data
  73         self._tracing_data = []
  74         return result
  75
  76   def _SendRequest(self, req):
  77     req['id'] = self._next_request_id
  78     self._next_request_id += 1
  79     data = json.dumps(req)
  80     self._socket.send(data)
  81
  82   def _ReceiveResponse(self):
  83     while self._socket:
  84       data = self._socket.recv()
  85       res = json.loads(data)
  86       return res
  87
  88   def _HandleResponse(self, res):
  89     method = res.get('method')
  90     value = res.get('params', {}).get('value')
  91     if 'Tracing.dataCollected' == method:
  92       if type(value) in [str, unicode]:
  93         self._tracing_data.append(value)
  94       elif type(value) is list:
  95         self._tracing_data.extend(value)
  96       else:
  97         logging.warning('Unexpected type in tracing data')
  98     elif 'Tracing.bufferUsage' == method and self._tracing_client:
  99       self._tracing_client.BufferUsage(value)
 100     elif 'Tracing.tracingComplete' == method:
 101       return True
 102
 103
 104 @contextlib.contextmanager
 105 def Connect(device_ip, devtools_port):
 106   backend = TracingBackend(devtools_port)
 107   try:
 108     backend.Connect(device_ip, devtools_port)
 109     yield backend
 110   finally:
 111     backend.Disconnect()
 112
 113
 114 def DumpTrace(trace, options):
 115   filepath = os.path.expanduser(options.output) if options.output \
 116       else os.path.join(os.getcwd(), 'trace.json')
 117
 118   dirname = os.path.dirname(filepath)
 119   if dirname:
 120     if not os.path.exists(dirname):
 121       os.makedirs(dirname)
 122   else:
 123     filepath = os.path.join(os.getcwd(), filepath)
 124
 125   with open(filepath, "w") as f:
 126    json.dump(trace, f)
 127   return filepath
 128
 129
 130 def _CreateOptionParser():
 131   parser = optparse.OptionParser(description='Record about://tracing profiles '
 132                                  'from any running instance of Chrome.')
 133   parser.add_option(
 134       '-v', '--verbose', help='Verbose logging.', action='store_true')
 135   parser.add_option(
 136       '-p', '--port', help='Remote debugging port.', type="int", default=9222)
 137   parser.add_option(
 138       '-d', '--device', help='Device ip address.', type='string',
 139       default='127.0.0.1')
 140
 141   tracing_opts = optparse.OptionGroup(parser, 'Tracing options')
 142   tracing_opts.add_option(
 143       '-c', '--category-filter',
 144       help='Apply filter to control what category groups should be traced.',
 145       type='string')
 146   tracing_opts.add_option(
 147       '--record-continuously',
 148       help='Keep recording until stopped. The trace buffer is of fixed size '
 149            'and used as a ring buffer. If this option is omitted then '
 150            'recording stops when the trace buffer is full.',
 151       action='store_true')
 152   parser.add_option_group(tracing_opts)
 153
 154   output_options = optparse.OptionGroup(parser, 'Output options')
 155   output_options.add_option(
 156       '-o', '--output',
 157       help='Save trace output to file.')
 158   parser.add_option_group(output_options)
 159
 160   return parser
 161
 162
 163 def _ProcessOptions(options):
 164   websocket.enableTrace(options.verbose)
 165
 166
 167 def main():
 168   parser = _CreateOptionParser()
 169   options, _args = parser.parse_args()
 170   _ProcessOptions(options)
 171
 172   with Connect(options.device, options.port) as tracing_backend:
 173     tracing_backend.StartTracing(TracingClient(),
 174                                  options.category_filter,
 175                                  options.record_continuously)
 176     raw_input('Capturing trace. Press Enter to stop...')
 177     trace = tracing_backend.StopTracing()
 178
 179   filepath = DumpTrace(trace, options)
 180   print('Done')
 181   print('Trace written to file://%s' % filepath)
 182
 183
 184 if __name__ == '__main__':
 185   sys.exit(main())