1# pylint: disable-msg=C0111
2# TODO: get rid of above, fix docstrings. crbug.com/273903
3# Copyright (c) 2013 The Chromium OS 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
7import logging
8
9
10try:
11    import statsd
12except ImportError:
13    import statsd_mock as statsd
14
15
16# This is _type for all metadata logged to elasticsearch from here.
17STATS_ES_TYPE = 'stats_metadata'
18
19
20# statsd logs details about what its sending at the DEBUG level, which I really
21# don't want to see tons of stats in logs, so all of these are silenced by
22# setting the logging level for all of statsdto WARNING.
23logging.getLogger('statsd').setLevel(logging.WARNING)
24
25
26def _prepend_init(_es, _conn, _prefix):
27    def wrapper(original):
28        """Decorator to override __init__."""
29
30        class _Derived(original):
31            def __init__(self, name, connection=None, bare=False,
32                         metadata=None):
33                name = self._add_prefix(name, _prefix, bare)
34                conn = connection if connection else _conn
35                super(_Derived, self).__init__(name, conn)
36                self.metadata = metadata
37                self.es = _es
38
39            def _add_prefix(self, name, prefix, bare=False):
40                """
41                Since many people run their own local AFE, stats from a local
42                setup shouldn't get mixed into stats from prod.  Therefore,
43                this function exists to add a prefix, nominally the name of
44                the local server, if |name| doesn't already start with the
45                server name, so that each person has their own "folder" of
46                stats that they can look at.
47
48                However, this functionality might not always be wanted, so we
49                allow one to pass in |bare=True| to force us to not prepend
50                the local server name. (I'm not sure when one would use this,
51                but I don't see why I should disallow it...)
52
53                >>> prefix = 'potato_nyc'
54                >>> _add_prefix('rpc.create_job', bare=False)
55                'potato_nyc.rpc.create_job'
56                >>> _add_prefix('rpc.create_job', bare=True)
57                'rpc.create_job'
58
59                @param name The name to append to the server name if it
60                            doesn't start with the server name.
61                @param bare If True, |name| will be returned un-altered.
62                @return A string to use as the stat name.
63
64                """
65                if not bare and not name.startswith(prefix):
66                    name = '%s.%s' % (prefix, name)
67                return name
68
69        return _Derived
70    return wrapper
71
72
73class Statsd(object):
74    def __init__(self, es, host, port, prefix):
75        # This is the connection that we're going to reuse for every client
76        # that gets created. This should maximally reduce overhead of stats
77        # logging.
78        self.conn = statsd.Connection(host=host, port=port)
79
80        @_prepend_init(es, self.conn, prefix)
81        class Average(statsd.Average):
82            """Wrapper around statsd.Average."""
83
84            def send(self, subname, value):
85                """Sends time-series data to graphite and metadata (if any)
86                to es.
87
88                @param subname: The subname to report the data to (i.e.
89                    'daisy.reboot')
90                @param value: Value to be sent.
91                """
92                statsd.Average.send(self, subname, value)
93                self.es.post(type_str=STATS_ES_TYPE, metadata=self.metadata,
94                             subname=subname, value=value)
95
96        self.Average = Average
97
98        @_prepend_init(es, self.conn, prefix)
99        class Counter(statsd.Counter):
100            """Wrapper around statsd.Counter."""
101
102            def _send(self, subname, value):
103                """Sends time-series data to graphite and metadata (if any)
104                to es.
105
106                @param subname: The subname to report the data to (i.e.
107                    'daisy.reboot')
108                @param value: Value to be sent.
109                """
110                statsd.Counter._send(self, subname, value)
111                self.es.post(type_str=STATS_ES_TYPE, metadata=self.metadata,
112                             subname=subname, value=value)
113
114        self.Counter = Counter
115
116        @_prepend_init(es, self.conn, prefix)
117        class Gauge(statsd.Gauge):
118            """Wrapper around statsd.Gauge."""
119
120            def send(self, subname, value):
121                """Sends time-series data to graphite and metadata (if any)
122                to es.
123
124                @param subname: The subname to report the data to (i.e.
125                    'daisy.reboot')
126                @param value: Value to be sent.
127                """
128                statsd.Gauge.send(self, subname, value)
129                self.es.post(type_str=STATS_ES_TYPE, metadata=self.metadata,
130                             subname=subname, value=value)
131
132        self.Gauge = Gauge
133
134        @_prepend_init(es, self.conn, prefix)
135        class Timer(statsd.Timer):
136            """Wrapper around statsd.Timer."""
137
138            # To override subname to not implicitly append 'total'.
139            def stop(self, subname=''):
140                statsd.Timer.stop(self, subname)
141
142
143            def send(self, subname, value):
144                """Sends time-series data to graphite and metadata (if any)
145                to es.
146
147                @param subname: The subname to report the data to (i.e.
148                    'daisy.reboot')
149                @param value: Value to be sent.
150                """
151                statsd.Timer.send(self, subname, value)
152                self.es.post(type_str=STATS_ES_TYPE, metadata=self.metadata,
153                             subname=self.name, value=value)
154
155
156            def __enter__(self):
157                self.start()
158                return self
159
160
161            def __exit__(self, exn_type, exn_value, traceback):
162                if exn_type is None:
163                    self.stop()
164
165        self.Timer = Timer
166
167        @_prepend_init(es, self.conn, prefix)
168        class Raw(statsd.Raw):
169            """Wrapper around statsd.Raw."""
170
171            def send(self, subname, value, timestamp=None):
172                """Sends time-series data to graphite and metadata (if any)
173                to es.
174
175                The datapoint we send is pretty much unchanged (will not be
176                aggregated)
177
178                @param subname: The subname to report the data to (i.e.
179                    'daisy.reboot')
180                @param value: Value to be sent.
181                @param timestamp: Time associated with when this stat was sent.
182                """
183                statsd.Raw.send(self, subname, value, timestamp)
184                self.es.post(type_str=STATS_ES_TYPE, metadata=self.metadata,
185                             subname=subname, value=value, timestamp=timestamp)
186
187        self.Raw = Raw
188