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