61a4682564e9855ac745ed410fea33562e234a84
[collecty.git] / src / collecty / plugins / base.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 from __future__ import division
23
24 import datetime
25 import logging
26 import math
27 import os
28 import rrdtool
29 import tempfile
30 import threading
31 import time
32
33 from ..constants import *
34 from ..i18n import _
35
36 _plugins = {}
37
38 def get():
39         """
40                 Returns a list with all automatically registered plugins.
41         """
42         return _plugins.values()
43
44 class Timer(object):
45         def __init__(self, timeout, heartbeat=1):
46                 self.timeout = timeout
47                 self.heartbeat = heartbeat
48
49                 self.delay = 0
50
51                 self.reset()
52
53         def reset(self, delay=0):
54                 # Save start time.
55                 self.start = time.time()
56
57                 self.delay = delay
58
59                 # Has this timer been killed?
60                 self.killed = False
61
62         @property
63         def elapsed(self):
64                 return time.time() - self.start - self.delay
65
66         def cancel(self):
67                 self.killed = True
68
69         def wait(self):
70                 while self.elapsed < self.timeout and not self.killed:
71                         time.sleep(self.heartbeat)
72
73                 return self.elapsed > self.timeout
74
75
76 class Plugin(threading.Thread):
77         # The name of this plugin.
78         name = None
79
80         # A description for this plugin.
81         description = None
82
83         # Templates which can be used to generate a graph out of
84         # the data from this data source.
85         templates = []
86
87         # The default interval for all plugins
88         interval = 60
89
90         # Automatically register all providers.
91         class __metaclass__(type):
92                 def __init__(plugin, name, bases, dict):
93                         type.__init__(plugin, name, bases, dict)
94
95                         # The main class from which is inherited is not registered
96                         # as a plugin.
97                         if name == "Plugin":
98                                 return
99
100                         if not all((plugin.name, plugin.description)):
101                                 raise RuntimeError(_("Plugin is not properly configured: %s") \
102                                         % plugin)
103
104                         _plugins[plugin.name] = plugin
105
106         def __init__(self, collecty, **kwargs):
107                 threading.Thread.__init__(self, name=self.description)
108                 self.daemon = True
109
110                 self.collecty = collecty
111
112                 # Check if this plugin was configured correctly.
113                 assert self.name, "Name of the plugin is not set: %s" % self.name
114                 assert self.description, "Description of the plugin is not set: %s" % self.description
115
116                 # Initialize the logger.
117                 self.log = logging.getLogger("collecty.plugins.%s" % self.name)
118                 self.log.propagate = 1
119
120                 self.data = []
121
122                 # Run some custom initialization.
123                 self.init(**kwargs)
124
125                 # Keepalive options
126                 self.running = True
127                 self.timer = Timer(self.interval)
128
129                 self.log.debug(_("Successfully initialized %s") % self.__class__.__name__)
130
131         @property
132         def path(self):
133                 """
134                         Returns the name of the sub directory in which all RRD files
135                         for this plugin should be stored in.
136                 """
137                 return self.name
138
139         ### Basic methods
140
141         def init(self, **kwargs):
142                 """
143                         Do some custom initialization stuff here.
144                 """
145                 pass
146
147         def collect(self):
148                 """
149                         Gathers the statistical data, this plugin collects.
150                 """
151                 time_start = time.time()
152
153                 # Run through all objects of this plugin and call the collect method.
154                 for o in self.objects:
155                         now = datetime.datetime.utcnow()
156                         try:
157                                 result = o.collect()
158
159                                 if isinstance(result, tuple) or isinstance(result, list):
160                                         result = ":".join(("%s" % e for e in result))
161                         except:
162                                 self.log.warning(_("Unhandled exception in %s.collect()") % o, exc_info=True)
163                                 continue
164
165                         if not result:
166                                 self.log.warning(_("Received empty result: %s") % o)
167                                 continue
168
169                         self.log.debug(_("Collected %s: %s") % (o, result))
170
171                         # Add the object to the write queue so that the data is written
172                         # to the databases later.
173                         self.collecty.write_queue.add(o, now, result)
174
175                 # Returns the time this function took to complete.
176                 return (time.time() - time_start)
177
178         def run(self):
179                 self.log.debug(_("%s plugin has started") % self.name)
180
181                 # Initially collect everything
182                 self.collect()
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                                 delay = self.collect()
191                                 self.timer.reset(delay)
192
193                 self.log.debug(_("%s plugin has stopped") % self.name)
194
195         def shutdown(self):
196                 self.log.debug(_("Received shutdown signal."))
197                 self.running = False
198
199                 # Kill any running timers.
200                 if self.timer:
201                         self.timer.cancel()
202
203         def get_object(self, id):
204                 for object in self.objects:
205                         if not object.id == id:
206                                 continue
207
208                         return object
209
210         def get_template(self, template_name, object_id):
211                 for template in self.templates:
212                         if not template.name == template_name:
213                                 continue
214
215                         return template(self, object_id)
216
217         def generate_graph(self, template_name, object_id="default", **kwargs):
218                 template = self.get_template(template_name, object_id=object_id)
219                 if not template:
220                         raise RuntimeError("Could not find template %s" % template_name)
221
222                 time_start = time.time()
223
224                 graph = template.generate_graph(**kwargs)
225
226                 duration = time.time() - time_start
227                 self.log.debug(_("Generated graph %s in %.1fms") \
228                         % (template, duration * 1000))
229
230                 return graph
231
232
233 class Object(object):
234         # The schema of the RRD database.
235         rrd_schema = None
236
237         # RRA properties.
238         rra_types     = ["AVERAGE", "MIN", "MAX"]
239         rra_timespans = [3600, 86400, 604800, 2678400, 31622400]
240         rra_rows      = 2880
241
242         def __init__(self, plugin, *args, **kwargs):
243                 self.plugin = plugin
244
245                 # Indicates if this object has collected its data
246                 self.collected = False
247
248                 # Initialise this object
249                 self.init(*args, **kwargs)
250
251                 # Create the database file.
252                 self.create()
253
254         def __repr__(self):
255                 return "<%s>" % self.__class__.__name__
256
257         @property
258         def collecty(self):
259                 return self.plugin.collecty
260
261         @property
262         def log(self):
263                 return self.plugin.log
264
265         @property
266         def id(self):
267                 """
268                         Returns a UNIQUE identifier for this object. As this is incorporated
269                         into the path of RRD file, it must only contain ASCII characters.
270                 """
271                 raise NotImplementedError
272
273         @property
274         def file(self):
275                 """
276                         The absolute path to the RRD file of this plugin.
277                 """
278                 filename = self._normalise_filename("%s.rrd" % self.id)
279
280                 return os.path.join(DATABASE_DIR, self.plugin.path, filename)
281
282         @staticmethod
283         def _normalise_filename(filename):
284                 # Convert the filename into ASCII characters only
285                 filename = filename.encode("ascii", "ignore")
286
287                 # Replace any spaces by dashes
288                 filename = filename.replace(" ", "-")
289
290                 return filename
291
292         ### Basic methods
293
294         def init(self, *args, **kwargs):
295                 """
296                         Do some custom initialization stuff here.
297                 """
298                 pass
299
300         def create(self):
301                 """
302                         Creates an empty RRD file with the desired data structures.
303                 """
304                 # Skip if the file does already exist.
305                 if os.path.exists(self.file):
306                         return
307
308                 dirname = os.path.dirname(self.file)
309                 if not os.path.exists(dirname):
310                         os.makedirs(dirname)
311
312                 # Create argument list.
313                 args = self.get_rrd_schema()
314
315                 rrdtool.create(self.file, *args)
316
317                 self.log.debug(_("Created RRD file %s.") % self.file)
318                 for arg in args:
319                         self.log.debug("  %s" % arg)
320
321         def info(self):
322                 return rrdtool.info(self.file)
323
324         @property
325         def stepsize(self):
326                 return self.plugin.interval
327
328         @property
329         def heartbeat(self):
330                 return self.stepsize * 2
331
332         def get_rrd_schema(self):
333                 schema = [
334                         "--step", "%s" % self.stepsize,
335                 ]
336                 for line in self.rrd_schema:
337                         if line.startswith("DS:"):
338                                 try:
339                                         (prefix, name, type, lower_limit, upper_limit) = line.split(":")
340
341                                         line = ":".join((
342                                                 prefix,
343                                                 name,
344                                                 type,
345                                                 "%s" % self.heartbeat,
346                                                 lower_limit,
347                                                 upper_limit
348                                         ))
349                                 except ValueError:
350                                         pass
351
352                         schema.append(line)
353
354                 xff = 0.1
355
356                 cdp_length = 0
357                 for rra_timespan in self.rra_timespans:
358                         if (rra_timespan / self.stepsize) < self.rra_rows:
359                                 rra_timespan = self.stepsize * self.rra_rows
360
361                         if cdp_length == 0:
362                                 cdp_length = 1
363                         else:
364                                 cdp_length = rra_timespan // (self.rra_rows * self.stepsize)
365
366                         cdp_number = math.ceil(rra_timespan / (cdp_length * self.stepsize))
367
368                         for rra_type in self.rra_types:
369                                 schema.append("RRA:%s:%.10f:%d:%d" % \
370                                         (rra_type, xff, cdp_length, cdp_number))
371
372                 return schema
373
374         def execute(self):
375                 if self.collected:
376                         raise RuntimeError("This object has already collected its data")
377
378                 self.collected = True
379                 self.now = datetime.datetime.utcnow()
380
381                 # Call the collect
382                 result = self.collect()
383
384         def commit(self):
385                 """
386                         Will commit the collected data to the database.
387                 """
388                 # Make sure that the RRD database has been created
389                 self.create()
390
391
392 class GraphTemplate(object):
393         # A unique name to identify this graph template.
394         name = None
395
396         # Headline of the graph image
397         graph_title = None
398
399         # Vertical label of the graph
400         graph_vertical_label = None
401
402         # Limits
403         lower_limit = None
404         upper_limit = None
405
406         # Instructions how to create the graph.
407         rrd_graph = None
408
409         # Extra arguments passed to rrdgraph.
410         rrd_graph_args = []
411
412         intervals = {
413                 None   : "-3h",
414                 "hour" : "-1h",
415                 "day"  : "-25h",
416                 "week" : "-360h",
417                 "year" : "-365d",
418         }
419
420         # Default dimensions for this graph
421         height = GRAPH_DEFAULT_HEIGHT
422         width  = GRAPH_DEFAULT_WIDTH
423
424         def __init__(self, plugin, object_id):
425                 self.plugin = plugin
426
427                 # Get all required RRD objects
428                 self.object_id = object_id
429
430                 # Get the main object
431                 self.object = self.get_object(self.object_id)
432
433         def __repr__(self):
434                 return "<%s>" % self.__class__.__name__
435
436         @property
437         def collecty(self):
438                 return self.plugin.collecty
439
440         @property
441         def log(self):
442                 return self.plugin.log
443
444         def _make_command_line(self, interval, format=DEFAULT_IMAGE_FORMAT,
445                         width=None, height=None):
446                 args = []
447
448                 args += GRAPH_DEFAULT_ARGUMENTS
449
450                 args += [
451                         "--imgformat", format,
452                         "--height", "%s" % (height or self.height),
453                         "--width", "%s" % (width or self.width),
454                 ]
455
456                 args += self.rrd_graph_args
457
458                 # Graph title
459                 if self.graph_title:
460                         args += ["--title", self.graph_title]
461
462                 # Vertical label
463                 if self.graph_vertical_label:
464                         args += ["--vertical-label", self.graph_vertical_label]
465
466                 if self.lower_limit is not None or self.upper_limit is not None:
467                         # Force to honour the set limits
468                         args.append("--rigid")
469
470                         if self.lower_limit is not None:
471                                 args += ["--lower-limit", self.lower_limit]
472
473                         if self.upper_limit is not None:
474                                 args += ["--upper-limit", self.upper_limit]
475
476                 # Add interval
477                 args.append("--start")
478
479                 try:
480                         args.append(self.intervals[interval])
481                 except KeyError:
482                         args.append(str(interval))
483
484                 return args
485
486         def get_object(self, *args, **kwargs):
487                 return self.plugin.get_object(*args, **kwargs)
488
489         def get_object_table(self):
490                 return {
491                         "file" : self.object,
492                 }
493
494         @property
495         def object_table(self):
496                 if not hasattr(self, "_object_table"):
497                         self._object_table = self.get_object_table()
498
499                 return self._object_table
500
501         def get_object_files(self):
502                 files = {}
503
504                 for id, obj in self.object_table.items():
505                         files[id] = obj.file
506
507                 return files
508
509         def generate_graph(self, interval=None, **kwargs):
510                 args = self._make_command_line(interval, **kwargs)
511
512                 self.log.info(_("Generating graph %s") % self)
513                 self.log.debug("  args: %s" % args)
514
515                 object_files = self.get_object_files()
516
517                 for item in self.rrd_graph:
518                         try:
519                                 args.append(item % object_files)
520                         except TypeError:
521                                 args.append(item)
522
523                         self.log.debug("  %s" % args[-1])
524
525                 return self.write_graph(*args)
526
527         def write_graph(self, *args):
528                 # Convert all arguments to string
529                 args = [str(e) for e in args]
530
531                 with tempfile.NamedTemporaryFile() as f:
532                         rrdtool.graph(f.name, *args)
533
534                         # Get back to the beginning of the file
535                         f.seek(0)
536
537                         # Return all the content
538                         return f.read()