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