Allow generating thumbnails of graph images
[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         def graph_info(self, template_name, object_id="default",
259                         timezone=None, locale=None, **kwargs):
260                 template = self.get_template(template_name, object_id=object_id,
261                         timezone=timezone, locale=locale)
262                 if not template:
263                         raise RuntimeError("Could not find template %s" % template_name)
264
265                 return template.graph_info()
266
267         def last_update(self, object_id="default"):
268                 object = self.get_object(object_id)
269                 if not object:
270                         raise RuntimeError("Could not find object %s" % object_id)
271
272                 return object.last_update()
273
274
275 class Object(object):
276         # The schema of the RRD database.
277         rrd_schema = None
278
279         # RRA properties.
280         rra_types     = ("AVERAGE", "MIN", "MAX")
281         rra_timespans = (
282                 ("1m", "10d"),
283                 ("1h", "18M"),
284                 ("1d",  "5y"),
285         )
286
287         def __init__(self, plugin, *args, **kwargs):
288                 self.plugin = plugin
289
290                 # Indicates if this object has collected its data
291                 self.collected = False
292
293                 # Initialise this object
294                 self.init(*args, **kwargs)
295
296                 # Create the database file.
297                 self.create()
298
299         def __repr__(self):
300                 return "<%s>" % self.__class__.__name__
301
302         @property
303         def collecty(self):
304                 return self.plugin.collecty
305
306         @property
307         def log(self):
308                 return self.plugin.log
309
310         @property
311         def id(self):
312                 """
313                         Returns a UNIQUE identifier for this object. As this is incorporated
314                         into the path of RRD file, it must only contain ASCII characters.
315                 """
316                 raise NotImplementedError
317
318         @property
319         def file(self):
320                 """
321                         The absolute path to the RRD file of this plugin.
322                 """
323                 filename = self._normalise_filename("%s.rrd" % self.id)
324
325                 return os.path.join(DATABASE_DIR, self.plugin.path, filename)
326
327         @staticmethod
328         def _normalise_filename(filename):
329                 # Convert the filename into ASCII characters only
330                 filename = unicodedata.normalize("NFKC", filename)
331
332                 # Replace any spaces by dashes
333                 filename = filename.replace(" ", "-")
334
335                 return filename
336
337         ### Basic methods
338
339         def init(self, *args, **kwargs):
340                 """
341                         Do some custom initialization stuff here.
342                 """
343                 pass
344
345         def create(self):
346                 """
347                         Creates an empty RRD file with the desired data structures.
348                 """
349                 # Skip if the file does already exist.
350                 if os.path.exists(self.file):
351                         return
352
353                 dirname = os.path.dirname(self.file)
354                 if not os.path.exists(dirname):
355                         os.makedirs(dirname)
356
357                 # Create argument list.
358                 args = self.get_rrd_schema()
359
360                 rrdtool.create(self.file, *args)
361
362                 self.log.debug(_("Created RRD file %s.") % self.file)
363                 for arg in args:
364                         self.log.debug("  %s" % arg)
365
366         def info(self):
367                 return rrdtool.info(self.file)
368
369         def last_update(self):
370                 """
371                         Returns a dictionary with the timestamp and
372                         data set of the last database update.
373                 """
374                 return {
375                         "dataset"   : self.last_dataset,
376                         "timestamp" : self.last_updated,
377                 }
378
379         def _last_update(self):
380                 return rrdtool.lastupdate(self.file)
381
382         @property
383         def last_updated(self):
384                 """
385                         Returns the timestamp when this database was last updated
386                 """
387                 lu = self._last_update()
388
389                 if lu:
390                         return lu.get("date")
391
392         @property
393         def last_dataset(self):
394                 """
395                         Returns the latest dataset in the database
396                 """
397                 lu = self._last_update()
398
399                 if lu:
400                         return lu.get("ds")
401
402         @property
403         def stepsize(self):
404                 return self.plugin.interval
405
406         @property
407         def heartbeat(self):
408                 return self.stepsize * 2
409
410         def get_rrd_schema(self):
411                 schema = [
412                         "--step", "%s" % self.stepsize,
413                 ]
414                 for line in self.rrd_schema:
415                         if line.startswith("DS:"):
416                                 try:
417                                         (prefix, name, type, lower_limit, upper_limit) = line.split(":")
418
419                                         line = ":".join((
420                                                 prefix,
421                                                 name,
422                                                 type,
423                                                 "%s" % self.heartbeat,
424                                                 lower_limit,
425                                                 upper_limit
426                                         ))
427                                 except ValueError:
428                                         pass
429
430                         schema.append(line)
431
432                 xff = 0.1
433
434                 for steps, rows in self.rra_timespans:
435                         for type in self.rra_types:
436                                 schema.append("RRA:%s:%s:%s:%s" % (type, xff, steps, rows))
437
438                 return schema
439
440         def execute(self):
441                 if self.collected:
442                         raise RuntimeError("This object has already collected its data")
443
444                 self.collected = True
445                 self.now = datetime.datetime.utcnow()
446
447                 # Call the collect
448                 result = self.collect()
449
450         def commit(self):
451                 """
452                         Will commit the collected data to the database.
453                 """
454                 # Make sure that the RRD database has been created
455                 self.create()
456
457                 # Write everything to disk that is in the write queue
458                 self.collecty.write_queue.commit_file(self.file)
459
460
461 class GraphTemplate(object):
462         # A unique name to identify this graph template.
463         name = None
464
465         # Headline of the graph image
466         graph_title = None
467
468         # Vertical label of the graph
469         graph_vertical_label = None
470
471         # Limits
472         lower_limit = None
473         upper_limit = None
474
475         # Instructions how to create the graph.
476         rrd_graph = None
477
478         # Extra arguments passed to rrdgraph.
479         rrd_graph_args = []
480
481         intervals = {
482                 None   : "-3h",
483                 "hour" : "-1h",
484                 "day"  : "-25h",
485                 "month": "-30d",
486                 "week" : "-360h",
487                 "year" : "-365d",
488         }
489
490         # Default dimensions for this graph
491         height = GRAPH_DEFAULT_HEIGHT
492         width  = GRAPH_DEFAULT_WIDTH
493
494         def __init__(self, plugin, object_id, locale=None, timezone=None):
495                 self.plugin = plugin
496
497                 # Save localisation parameters
498                 self.locale = locales.get(locale)
499                 self.timezone = timezone
500
501                 # Get all required RRD objects
502                 self.object_id = object_id
503
504                 # Get the main object
505                 self.object = self.get_object(self.object_id)
506
507         def __repr__(self):
508                 return "<%s>" % self.__class__.__name__
509
510         @property
511         def collecty(self):
512                 return self.plugin.collecty
513
514         @property
515         def log(self):
516                 return self.plugin.log
517
518         def _make_command_line(self, interval, format=DEFAULT_IMAGE_FORMAT,
519                         width=None, height=None, with_title=True, thumbnail=False):
520                 args = [e for e in GRAPH_DEFAULT_ARGUMENTS]
521
522                 # Set the default dimensions
523                 default_height, default_width = GRAPH_DEFAULT_HEIGHT, GRAPH_DEFAULT_WIDTH
524
525                 # A thumbnail doesn't have a legend and other labels
526                 if thumbnail:
527                         args.append("--only-graph")
528
529                         default_height = THUMBNAIL_DEFAULT_HEIGHT
530                         default_width = THUMBNAIL_DEFAULT_WIDTH
531
532                 args += [
533                         "--imgformat", format,
534                         "--height", "%s" % (height or default_height),
535                         "--width", "%s" % (width or default_width),
536                 ]
537
538                 args += self.rrd_graph_args
539
540                 # Graph title
541                 if with_title and self.graph_title:
542                         args += ["--title", self.graph_title]
543
544                 # Vertical label
545                 if self.graph_vertical_label:
546                         args += ["--vertical-label", self.graph_vertical_label]
547
548                 if self.lower_limit is not None or self.upper_limit is not None:
549                         # Force to honour the set limits
550                         args.append("--rigid")
551
552                         if self.lower_limit is not None:
553                                 args += ["--lower-limit", self.lower_limit]
554
555                         if self.upper_limit is not None:
556                                 args += ["--upper-limit", self.upper_limit]
557
558                 try:
559                         interval = self.intervals[interval]
560                 except KeyError:
561                         interval = "end-%s" % interval
562
563                 # Add interval
564                 args += ["--start", interval]
565
566                 return args
567
568         def get_object(self, *args, **kwargs):
569                 return self.plugin.get_object(*args, **kwargs)
570
571         def get_object_table(self):
572                 return {
573                         "file" : self.object,
574                 }
575
576         @property
577         def object_table(self):
578                 if not hasattr(self, "_object_table"):
579                         self._object_table = self.get_object_table()
580
581                 return self._object_table
582
583         def get_object_files(self):
584                 files = {}
585
586                 for id, obj in self.object_table.items():
587                         files[id] = obj.file
588
589                 return files
590
591         def generate_graph(self, interval=None, **kwargs):
592                 # Make sure that all collected data is in the database
593                 # to get a recent graph image
594                 if self.object:
595                         self.object.commit()
596
597                 args = self._make_command_line(interval, **kwargs)
598
599                 self.log.info(_("Generating graph %s") % self)
600                 self.log.debug("  args: %s" % args)
601
602                 object_files = self.get_object_files()
603
604                 for item in self.rrd_graph:
605                         try:
606                                 args.append(item % object_files)
607                         except TypeError:
608                                 args.append(item)
609
610                         self.log.debug("  %s" % args[-1])
611
612                 # Convert arguments to string
613                 args = [str(e) for e in args]
614
615                 with Environment(self.timezone, self.locale.lang):
616                         graph = rrdtool.graphv("-", *args)
617
618                 return {
619                         "image"        : graph.get("image"),
620                         "image_height" : graph.get("image_height"),
621                         "image_width"  : graph.get("image_width"),
622                 }
623
624         def graph_info(self):
625                 """
626                         Returns a dictionary with useful information
627                         about this graph.
628                 """
629                 return {
630                         "title"        : self.graph_title,
631                         "object_id"    : self.object_id or "",
632                         "template"     : self.name,
633                 }