1#!/usr/bin/env python
2# Copyright (c) 2018 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6from __future__ import print_function
7
8import contextlib
9import functools
10import json
11import os
12import sys
13import tempfile
14import threading
15import time
16import traceback
17
18try:
19  import urllib2 as urllib
20except ImportError:  # For Py3 compatibility
21  import urllib.request as urllib
22
23import detect_host_arch
24import gclient_utils
25import metrics_utils
26import subprocess2
27
28
29DEPOT_TOOLS = os.path.dirname(os.path.abspath(__file__))
30CONFIG_FILE = os.path.join(DEPOT_TOOLS, 'metrics.cfg')
31UPLOAD_SCRIPT = os.path.join(DEPOT_TOOLS, 'upload_metrics.py')
32
33DISABLE_METRICS_COLLECTION = os.environ.get('DEPOT_TOOLS_METRICS') == '0'
34DEFAULT_COUNTDOWN = 10
35
36INVALID_CONFIG_WARNING = (
37    'WARNING: Your metrics.cfg file was invalid or nonexistent. A new one will '
38    'be created.'
39)
40PERMISSION_DENIED_WARNING = (
41    'Could not write the metrics collection config:\n\t%s\n'
42    'Metrics collection will be disabled.'
43)
44
45
46class _Config(object):
47  def __init__(self):
48    self._initialized = False
49    self._config = {}
50
51  def _ensure_initialized(self):
52    if self._initialized:
53      return
54
55    try:
56      config = json.loads(gclient_utils.FileRead(CONFIG_FILE))
57    except (IOError, ValueError):
58      config = {}
59
60    self._config = config.copy()
61
62    if 'is-googler' not in self._config:
63      # /should-upload is only accessible from Google IPs, so we only need to
64      # check if we can reach the page. An external developer would get access
65      # denied.
66      try:
67        req = urllib.urlopen(metrics_utils.APP_URL + '/should-upload')
68        self._config['is-googler'] = req.getcode() == 200
69      except (urllib.URLError, urllib.HTTPError):
70        self._config['is-googler'] = False
71
72    # Make sure the config variables we need are present, and initialize them to
73    # safe values otherwise.
74    self._config.setdefault('countdown', DEFAULT_COUNTDOWN)
75    self._config.setdefault('opt-in', None)
76    self._config.setdefault('version', metrics_utils.CURRENT_VERSION)
77
78    if config != self._config:
79      print(INVALID_CONFIG_WARNING, file=sys.stderr)
80      self._write_config()
81
82    self._initialized = True
83
84  def _write_config(self):
85    try:
86      gclient_utils.FileWrite(CONFIG_FILE, json.dumps(self._config))
87    except IOError as e:
88      print(PERMISSION_DENIED_WARNING % e, file=sys.stderr)
89      self._config['opt-in'] = False
90
91  @property
92  def version(self):
93    self._ensure_initialized()
94    return self._config['version']
95
96  @property
97  def is_googler(self):
98    self._ensure_initialized()
99    return self._config['is-googler']
100
101  @property
102  def opted_in(self):
103    self._ensure_initialized()
104    return self._config['opt-in']
105
106  @opted_in.setter
107  def opted_in(self, value):
108    self._ensure_initialized()
109    self._config['opt-in'] = value
110    self._config['version'] = metrics_utils.CURRENT_VERSION
111    self._write_config()
112
113  @property
114  def countdown(self):
115    self._ensure_initialized()
116    return self._config['countdown']
117
118  @property
119  def should_collect_metrics(self):
120    # Don't collect the metrics unless the user is a googler, the user has opted
121    # in, or the countdown has expired.
122    if not self.is_googler:
123      return False
124    if self.opted_in is False:
125      return False
126    if self.opted_in is None and self.countdown > 0:
127      return False
128    return True
129
130  def decrease_countdown(self):
131    self._ensure_initialized()
132    if self.countdown == 0:
133      return
134    self._config['countdown'] -= 1
135    if self.countdown == 0:
136      self._config['version'] = metrics_utils.CURRENT_VERSION
137    self._write_config()
138
139  def reset_config(self):
140    # Only reset countdown if we're already collecting metrics.
141    if self.should_collect_metrics:
142      self._ensure_initialized()
143      self._config['countdown'] = DEFAULT_COUNTDOWN
144      self._config['opt-in'] = None
145
146
147class MetricsCollector(object):
148  def __init__(self):
149    self._metrics_lock = threading.Lock()
150    self._reported_metrics = {}
151    self._config = _Config()
152    self._collecting_metrics = False
153    self._collect_custom_metrics = True
154
155  @property
156  def config(self):
157    return self._config
158
159  @property
160  def collecting_metrics(self):
161    return self._collecting_metrics
162
163  def add(self, name, value):
164    if self._collect_custom_metrics:
165      with self._metrics_lock:
166        self._reported_metrics[name] = value
167
168  def add_repeated(self, name, value):
169    if self._collect_custom_metrics:
170      with self._metrics_lock:
171        self._reported_metrics.setdefault(name, []).append(value)
172
173  @contextlib.contextmanager
174  def pause_metrics_collection(self):
175    collect_custom_metrics = self._collect_custom_metrics
176    self._collect_custom_metrics = False
177    try:
178      yield
179    finally:
180      self._collect_custom_metrics = collect_custom_metrics
181
182  def _upload_metrics_data(self):
183    """Upload the metrics data to the AppEngine app."""
184    # We invoke a subprocess, and use stdin.write instead of communicate(),
185    # so that we are able to return immediately, leaving the upload running in
186    # the background.
187    p = subprocess2.Popen(['vpython3', UPLOAD_SCRIPT], stdin=subprocess2.PIPE)
188    p.stdin.write(json.dumps(self._reported_metrics).encode('utf-8'))
189
190  def _collect_metrics(self, func, command_name, *args, **kwargs):
191    # If we're already collecting metrics, just execute the function.
192    # e.g. git-cl split invokes git-cl upload several times to upload each
193    # split CL.
194    if self.collecting_metrics:
195      # Don't collect metrics for this function.
196      # e.g. Don't record the arguments git-cl split passes to git-cl upload.
197      with self.pause_metrics_collection():
198        return func(*args, **kwargs)
199
200    self._collecting_metrics = True
201    self.add('metrics_version', metrics_utils.CURRENT_VERSION)
202    self.add('command', command_name)
203    try:
204      start = time.time()
205      result = func(*args, **kwargs)
206      exception = None
207    # pylint: disable=bare-except
208    except:
209      exception = sys.exc_info()
210    finally:
211      self.add('execution_time', time.time() - start)
212
213    exit_code = metrics_utils.return_code_from_exception(exception)
214    self.add('exit_code', exit_code)
215
216    # Add metrics regarding environment information.
217    self.add('timestamp', int(time.time()))
218    self.add('python_version', metrics_utils.get_python_version())
219    self.add('host_os', gclient_utils.GetMacWinAixOrLinux())
220    self.add('host_arch', detect_host_arch.HostArch())
221
222    depot_tools_age = metrics_utils.get_repo_timestamp(DEPOT_TOOLS)
223    if depot_tools_age is not None:
224      self.add('depot_tools_age', int(depot_tools_age))
225
226    git_version = metrics_utils.get_git_version()
227    if git_version:
228      self.add('git_version', git_version)
229
230    self._upload_metrics_data()
231    if exception:
232      gclient_utils.reraise(exception[0], exception[1], exception[2])
233    return result
234
235  def collect_metrics(self, command_name):
236    """A decorator used to collect metrics over the life of a function.
237
238    This decorator executes the function and collects metrics about the system
239    environment and the function performance.
240    """
241    def _decorator(func):
242      # Do this first so we don't have to read, and possibly create a config
243      # file.
244      if DISABLE_METRICS_COLLECTION:
245        return func
246      if not self.config.should_collect_metrics:
247        return func
248      # Otherwise, collect the metrics.
249      # Needed to preserve the __name__ and __doc__ attributes of func.
250      @functools.wraps(func)
251      def _inner(*args, **kwargs):
252        return self._collect_metrics(func, command_name, *args, **kwargs)
253      return _inner
254    return _decorator
255
256  @contextlib.contextmanager
257  def print_notice_and_exit(self):
258    """A context manager used to print the notice and terminate execution.
259
260    This decorator executes the function and prints the monitoring notice if
261    necessary. If an exception is raised, we will catch it, and print it before
262    printing the metrics collection notice.
263    This will call sys.exit() with an appropriate exit code to ensure the notice
264    is the last thing printed."""
265    # Needed to preserve the __name__ and __doc__ attributes of func.
266    try:
267      yield
268      exception = None
269    # pylint: disable=bare-except
270    except:
271      exception = sys.exc_info()
272
273    # Print the exception before the metrics notice, so that the notice is
274    # clearly visible even if gclient fails.
275    if exception:
276      if isinstance(exception[1], KeyboardInterrupt):
277        sys.stderr.write('Interrupted\n')
278      elif not isinstance(exception[1], SystemExit):
279        traceback.print_exception(*exception)
280
281    # Check if the version has changed
282    if (not DISABLE_METRICS_COLLECTION and self.config.is_googler
283        and self.config.opted_in is not False
284        and self.config.version != metrics_utils.CURRENT_VERSION):
285      metrics_utils.print_version_change(self.config.version)
286      self.config.reset_config()
287
288    # Print the notice
289    if (not DISABLE_METRICS_COLLECTION and self.config.is_googler
290        and self.config.opted_in is None):
291      metrics_utils.print_notice(self.config.countdown)
292      self.config.decrease_countdown()
293
294    sys.exit(metrics_utils.return_code_from_exception(exception))
295
296
297collector = MetricsCollector()
298