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