ca98ef721df25085098e1a744e2b860a15d8cd18
[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 }