Allow atrace over adb tcp/ip

systrace/systrace/tracing_agents/atrace_agent.py

Index: systrace/systrace/tracing_agents/atrace_agent.py
diff --git a/systrace/systrace/tracing_agents/atrace_agent.py b/systrace/systrace/tracing_agents/atrace_agent.py
index a22a6aa59d02a150052ef0f49cc37a7422a7c776..9e5a87dc0f2bd79be49720091e838e49f0812241 100644
--- a/systrace/systrace/tracing_agents/atrace_agent.py
+++ b/systrace/systrace/tracing_agents/atrace_agent.py
@@ -48,7 +48,8 @@ def list_categories(config):
   Args:
       config: Tracing config.
   """
-  devutils = device_utils.DeviceUtils(config.device_serial_number)
+  devutils = device_utils.DeviceUtils(config.device_serial_number,
+                                      disable_gce=config.disable_gce)
   categories = devutils.RunShellCommand(
       LIST_CATEGORIES_ARGS, check_return=True)
 
@@ -68,7 +69,8 @@ def get_available_categories(config, device_sdk_version):
       config: Tracing config.
       device_sdk_version: Sdk version int of device to be queried.
   """
-  devutils = device_utils.DeviceUtils(config.device_serial_number)
+  devutils = device_utils.DeviceUtils(config.device_serial_number,
+                                      disable_gce=config.disable_gce)
   categories_output = devutils.RunShellCommand(
       LIST_CATEGORIES_ARGS, check_return=True)
   categories = [c.split('-')[0].strip() for c in categories_output]
@@ -178,7 +180,8 @@ class AtraceAgent(tracing_agents.TracingAgent):
         x in avail_cats]
     if unavailable:
       print 'These categories are unavailable: ' + ' '.join(unavailable)
-    self._device_utils = device_utils.DeviceUtils(config.device_serial_number)
+    self._device_utils = device_utils.DeviceUtils(config.device_serial_number,
+                           disable_gce=config.disable_gce)
     self._device_serial_number = config.device_serial_number
     self._tracer_args = _construct_atrace_args(config,
                                                self._categories)
@@ -479,7 +482,7 @@ def fix_circular_traces(out):
 class AtraceConfig(tracing_agents.TracingConfig):
   def __init__(self, atrace_categories, trace_buf_size, kfuncs,
                app_name, compress_trace_data, from_file,
-               device_serial_number, trace_time, target):
+               device_serial_number, trace_time, target, disable_gce):
     tracing_agents.TracingConfig.__init__(self)
     self.atrace_categories = atrace_categories
     self.trace_buf_size = trace_buf_size
@@ -490,6 +493,7 @@ class AtraceConfig(tracing_agents.TracingConfig):
     self.device_serial_number = device_serial_number
     self.trace_time = trace_time
     self.target = target
+    self.disable_gce = disable_gce
 
 
 def add_options(parser):
@@ -512,6 +516,10 @@ def add_options(parser):
                      action='store', help='read the trace from a '
                      'file (compressed) rather than running a '
                      'live trace')
+  options.add_option('--disable-gce', dest='disable_gce',
+                     default=False, action='store_true',
+                     help='treat devices connected by tcp-ip as '
+                     'local instead of in GCE')
   return options
 
 def get_config(options):
@@ -519,4 +527,4 @@ def get_config(options):
                       options.trace_buf_size, options.kfuncs,
                       options.app_name, options.compress_trace_data,
                       options.from_file, options.device_serial_number,
-                      options.trace_time, options.target)
+                      options.trace_time, options.target, options.disable_gce)