1# Copyright 2016 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4import atexit
5import json
6import os
7import sys
8import time
9import threading
10
11from catapult_base import lock
12
13
14_lock = threading.Lock()
15
16_enabled = False
17_log_file = None
18
19_cur_events = [] # events that have yet to be buffered
20
21_tls = threading.local() # tls used to detect forking/etc
22_atexit_regsitered_for_pid = None
23
24_control_allowed = True
25
26
27class TraceException(Exception):
28  pass
29
30def _note(msg, *args):
31  pass
32#  print "%i: %s" % (os.getpid(), msg)
33
34
35def _locked(fn):
36  def locked_fn(*args,**kwargs):
37    _lock.acquire()
38    try:
39      ret = fn(*args,**kwargs)
40    finally:
41      _lock.release()
42    return ret
43  return locked_fn
44
45def _disallow_tracing_control():
46  global _control_allowed
47  _control_allowed = False
48
49def trace_enable(log_file=None):
50  _trace_enable(log_file)
51
52@_locked
53def _trace_enable(log_file=None):
54  global _enabled
55  if _enabled:
56    raise TraceException("Already enabled")
57  if not _control_allowed:
58    raise TraceException("Tracing control not allowed in child processes.")
59  _enabled = True
60  global _log_file
61  if log_file == None:
62    if sys.argv[0] == '':
63      n = 'trace_event'
64    else:
65      n = sys.argv[0]
66    log_file = open("%s.json" % n, "ab", False)
67    _note("trace_event: tracelog name is %s.json" % n)
68  elif isinstance(log_file, basestring):
69    _note("trace_event: tracelog name is %s" % log_file)
70    log_file = open("%s" % log_file, "ab", False)
71  elif not hasattr(log_file, 'fileno'):
72    raise TraceException(
73        "Log file must be None, a string, or file-like object with a fileno()")
74
75  _log_file = log_file
76  with lock.FileLock(_log_file, lock.LOCK_EX):
77    _log_file.seek(0, os.SEEK_END)
78
79    lastpos = _log_file.tell()
80    creator = lastpos == 0
81    if creator:
82      _note("trace_event: Opened new tracelog, lastpos=%i", lastpos)
83      _log_file.write('[')
84
85      tid = threading.current_thread().ident
86      if not tid:
87        tid = os.getpid()
88      x = {"ph": "M", "category": "process_argv",
89           "pid": os.getpid(), "tid": threading.current_thread().ident,
90           "ts": time.time(),
91           "name": "process_argv", "args": {"argv": sys.argv}}
92      _log_file.write("%s\n" % json.dumps(x))
93    else:
94      _note("trace_event: Opened existing tracelog")
95    _log_file.flush()
96
97@_locked
98def trace_flush():
99  if _enabled:
100    _flush()
101
102@_locked
103def trace_disable():
104  global _enabled
105  if not _control_allowed:
106    raise TraceException("Tracing control not allowed in child processes.")
107  if not _enabled:
108    return
109  _enabled = False
110  _flush(close=True)
111
112def _flush(close=False):
113  global _log_file
114  with lock.FileLock(_log_file, lock.LOCK_EX):
115    _log_file.seek(0, os.SEEK_END)
116    if len(_cur_events):
117      _log_file.write(",\n")
118      _log_file.write(",\n".join([json.dumps(e) for e in _cur_events]))
119      del _cur_events[:]
120
121    if close:
122      # We might not be the only process writing to this logfile. So,
123      # we will simply close the file rather than writign the trailing ] that
124      # it technically requires. The trace viewer understands that this may
125      # happen and will insert a trailing ] during loading.
126      pass
127    _log_file.flush()
128
129  if close:
130    _note("trace_event: Closed")
131    _log_file.close()
132    _log_file = None
133  else:
134    _note("trace_event: Flushed")
135
136@_locked
137def trace_is_enabled():
138  return _enabled
139
140@_locked
141def add_trace_event(ph, ts, category, name, args=None):
142  global _enabled
143  if not _enabled:
144    return
145  if not hasattr(_tls, 'pid') or _tls.pid != os.getpid():
146    _tls.pid = os.getpid()
147    global _atexit_regsitered_for_pid
148    if _tls.pid != _atexit_regsitered_for_pid:
149      _atexit_regsitered_for_pid = _tls.pid
150      atexit.register(_trace_disable_atexit)
151      _tls.pid = os.getpid()
152      del _cur_events[:] # we forked, clear the event buffer!
153    tid = threading.current_thread().ident
154    if not tid:
155      tid = os.getpid()
156    _tls.tid = tid
157
158  _cur_events.append({"ph": ph,
159                      "category": category,
160                      "pid": _tls.pid,
161                      "tid": _tls.tid,
162                      "ts": ts,
163                      "name": name,
164                      "args": args or {}});
165
166def trace_begin(name, args=None):
167  add_trace_event("B", time.time(), "python", name, args)
168
169def trace_end(name, args=None):
170  add_trace_event("E", time.time(), "python", name, args)
171
172def _trace_disable_atexit():
173  trace_disable()
174
175def is_tracing_controllable():
176  global _control_allowed
177  return _control_allowed
178