1# Copyright (c) 2014 The Chromium OS 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.
4
5
6import time
7import functools
8import logging
9
10
11_registry = {}
12
13
14def stat(category=None):
15    """
16    Decorator that registers the function as a function that, when called,
17    will submit a stat to statsd.
18
19    @param category: The set of servers to be run against.
20    @param f: A function that submits stats.
21    @returns: f
22    """
23    def curry(f):  # pylint: disable-msg=C0111
24        _registry.setdefault(category, []).append(f)
25        return f
26    return curry
27
28
29def loop_stat(category=None):
30    """
31    Decorator that registers the function as a function that, when called,
32    will submit a stat to statsd.  This function is then registered so that
33    it will be called periodically.
34
35    You probably want to use this one.
36
37    @param category: The set of servers to be run against.
38    @param f: A function that submits stats.
39    @returns: f
40    """
41    def curry(f):  # pylint: disable-msg=C0111
42        @functools.wraps(f)
43        def looped(*args, **kwargs):  # pylint: disable-msg=C0111
44            while True:
45                try:
46                    f(*args, **kwargs)
47                except Exception as e:
48                    logging.exception(e)
49                time.sleep(15)
50        _registry.setdefault(category, []).append(looped)
51        return f
52    return curry
53
54
55def registered_functions():
56    """
57    Return all functions registered as a stat.
58
59    returns: A list of 0-arity functions.
60    """
61    return _registry
62