Add dbus interface
[collecty.git] / src / collecty / daemon.py
1 #!/usr/bin/python
2 ###############################################################################
3 #                                                                             #
4 # collecty - A system statistics collection daemon for IPFire                 #
5 # Copyright (C) 2012 IPFire development team                                  #
6 #                                                                             #
7 # This program is free software: you can redistribute it and/or modify        #
8 # it under the terms of the GNU General Public License as published by        #
9 # the Free Software Foundation, either version 3 of the License, or           #
10 # (at your option) any later version.                                         #
11 #                                                                             #
12 # This program is distributed in the hope that it will be useful,             #
13 # but WITHOUT ANY WARRANTY; without even the implied warranty of              #
14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
15 # GNU General Public License for more details.                                #
16 #                                                                             #
17 # You should have received a copy of the GNU General Public License           #
18 # along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
19 #                                                                             #
20 ###############################################################################
21
22 import Queue as queue
23 import rrdtool
24 import signal
25 import threading
26 import time
27
28 import bus
29 import plugins
30
31 from constants import *
32 from i18n import _
33
34 import logging
35 log = logging.getLogger("collecty")
36
37 class Collecty(object):
38         # The default interval, when all data is written to disk.
39         SUBMIT_INTERVAL = 300
40
41         HEARTBEAT = 5
42
43         def __init__(self, debug=False):
44                 self.debug = debug
45
46                 # Enable debug logging when running in debug mode
47                 if self.debug:
48                         log.setLevel(logging.DEBUG)
49
50                 self.plugins = []
51
52                 # Indicates whether this process should be running or not.
53                 self.running = True
54
55                 # The write queue holds all collected pieces of data which
56                 # will be written to disk later.
57                 self.write_queue = WriteQueue(self, self.SUBMIT_INTERVAL)
58
59                 # Create a thread that connects to dbus and processes requests we
60                 # get from there.
61                 self.bus = bus.Bus(self)
62
63                 # Add all plugins
64                 for plugin in plugins.get():
65                         self.add_plugin(plugin)
66
67                 log.info(_("Collecty successfully initialized with %s plugins") \
68                         % len(self.plugins))
69
70         def add_plugin(self, plugin_class):
71                 # Try initialising a new plugin. If that fails, we will log the
72                 # error and try to go on.
73                 try:
74                         plugin = plugin_class(self)
75                 except:
76                         log.critical(_("Plugin %s could not be initialised") % plugin_class, exc_info=True)
77                         return
78
79                 self.plugins.append(plugin)
80
81         @property
82         def templates(self):
83                 for plugin in self.plugins:
84                         for template in plugin.templates:
85                                 yield template
86
87         def run(self):
88                 # Register signal handlers.
89                 self.register_signal_handler()
90
91                 # Start the bus
92                 self.bus.start()
93
94                 # Start all data source threads.
95                 for p in self.plugins:
96                         p.start()
97
98                 # Run the write queue thread
99                 self.write_queue.start()
100
101                 # Regularly submit all data to disk.
102                 while self.running:
103                         try:
104                                 time.sleep(self.HEARTBEAT)
105                         except KeyboardInterrupt:
106                                 self.shutdown()
107                                 break
108
109                 # Wait until all plugins are finished.
110                 for p in self.plugins:
111                         p.join()
112
113                 # Stop the bus thread
114                 self.bus.shutdown()
115
116                 # Write all collected data to disk before ending the main thread
117                 self.write_queue.shutdown()
118
119                 log.debug(_("Main thread exited"))
120
121         def shutdown(self):
122                 if not self.running:
123                         return
124
125                 log.debug(_("Received shutdown signal"))
126                 self.running = False
127
128                 # Propagating shutdown to all threads.
129                 for p in self.plugins:
130                         p.shutdown()
131
132         def register_signal_handler(self):
133                 for s in (signal.SIGTERM, signal.SIGINT, signal.SIGUSR1):
134                         log.debug(_("Registering signal %d") % s)
135
136                 signal.signal(s, self.signal_handler)
137
138         def signal_handler(self, sig, *args, **kwargs):
139                 log.info(_("Caught signal %d") % sig)
140
141                 if sig in (signal.SIGTERM, signal.SIGINT):
142                         # Shutdown this application.
143                         self.shutdown()
144
145                 elif sig == signal.SIGUSR1:
146                         # Commit all data.
147                         self.write_queue.commit()
148
149         def get_plugin_from_template(self, template_name):
150                 for plugin in self.plugins:
151                         if not template_name in [t.name for t in plugin.templates]:
152                                 continue
153
154                         return plugin
155
156         def generate_graph(self, template_name, *args, **kwargs):
157                 plugin = self.get_plugin_from_template(template_name)
158                 if not plugin:
159                         raise RuntimeError("Could not find template %s" % template_name)
160
161                 return plugin.generate_graph(template_name, *args, **kwargs)
162
163
164 class WriteQueue(threading.Thread):
165         def __init__(self, collecty, submit_interval):
166                 threading.Thread.__init__(self)
167                 self.daemon = True
168
169                 self.collecty = collecty
170
171                 self.log = logging.getLogger("collecty.queue")
172                 self.log.setLevel(logging.DEBUG)
173                 self.log.propagate = 1
174
175                 self.timer = plugins.Timer(submit_interval)
176                 self._queue = queue.PriorityQueue()
177
178                 self.log.debug(_("Initialised write queue"))
179
180         def run(self):
181                 self.log.debug(_("Write queue process started"))
182                 self.running = True
183
184                 while self.running:
185                         # Reset the timer.
186                         self.timer.reset()
187
188                         # Wait until the timer has successfully elapsed.
189                         if self.timer.wait():
190                                 self.commit()
191
192                 self.commit()
193                 self.log.debug(_("Write queue process stopped"))
194
195         def shutdown(self):
196                 self.running = False
197                 self.timer.cancel()
198
199                 # Wait until all data has been written.
200                 self.join()
201
202         def add(self, object, time, data):
203                 result = QueueObject(object.file, time, data)
204                 self._queue.put(result)
205
206         def commit(self):
207                 """
208                         Flushes the read data to disk.
209                 """
210                 # There is nothing to do if the queue is empty
211                 if self._queue.empty():
212                         self.log.debug(_("No data to commit"))
213                         return
214
215                 time_start = time.time()
216
217                 self.log.debug(_("Submitting data to the databases..."))
218
219                 # Get all objects from the queue and group them by the RRD file
220                 # to commit them all at once
221                 results = {}
222                 while not self._queue.empty():
223                         result = self._queue.get()
224
225                         try:
226                                 results[result.file].append(result)
227                         except KeyError:
228                                 results[result.file] = [result]
229
230                 # Write the collected data to disk
231                 for filename, results in results.items():
232                         self._commit_file(filename, results)
233
234                 duration = time.time() - time_start
235                 self.log.debug(_("Emptied write queue in %.2fs") % duration)
236
237         def _commit_file(self, filename, results):
238                 self.log.debug(_("Committing %(counter)s entries to %(filename)s:") \
239                         % { "counter" : len(results), "filename" : filename })
240
241                 if self.collecty.debug:
242                         for result in results:
243                                 self.log.debug("  %s: %s" % (result.time, result.data))
244
245                 rrdtool.update(filename, *["%s" % r for r in results])
246
247
248 class QueueObject(object):
249         def __init__(self, file, time, data):
250                 self.file = file
251                 self.time = time
252                 self.data = data
253
254         def __str__(self):
255                 return "%s:%s" % (self.time.strftime("%s"), self.data)
256
257         def __cmp__(self, other):
258                 return cmp(self.time, other.time)