src/collecty/daemon.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 multiprocessing
  25 import os
  26 import queue
  27 import rrdtool
  28 import signal
  29 import threading
  30 import time
  31
  32 from . import bus
  33 from . import locales
  34 from . import plugins
  35
  36 from .constants import *
  37 from .i18n import _
  38
  39 log = logging.getLogger("collecty")
  40
  41 class Collecty(object):
  42         # The default interval, when all data is written to disk.
  43         SUBMIT_INTERVAL = 300
  44
  45         HEARTBEAT = 1
  46
  47         def __init__(self, debug=False):
  48                 self.debug = debug
  49
  50                 # Reset timezone to UTC
  51                 # rrdtool is reading that from the environment
  52                 os.environ["TZ"] = "UTC"
  53
  54                 # Enable debug logging when running in debug mode
  55                 if self.debug:
  56                         log.setLevel(logging.DEBUG)
  57
  58                 self.plugins = []
  59
  60                 # Indicates whether this process should be running or not.
  61                 self.running = True
  62
  63                 # The write queue holds all collected pieces of data which
  64                 # will be written to disk later.
  65                 self.write_queue = WriteQueue(self, self.SUBMIT_INTERVAL)
  66
  67                 # Create worker threads
  68                 self.worker_threads = self.create_worker_threads()
  69
  70                 self._timer_queue = queue.PriorityQueue()
  71                 self._worker_queue = queue.Queue()
  72
  73                 # Create a thread that connects to dbus and processes requests we
  74                 # get from there.
  75                 self.bus = bus.Bus(self)
  76
  77                 # Add all plugins
  78                 for plugin in plugins.get():
  79                         self.add_plugin(plugin)
  80
  81                 log.debug(_("Collecty successfully initialized with %s plugins") \
  82                         % len(self.plugins))
  83
  84                 log.debug(_("Supported locales: %s") % ", ".join(locales.get_supported_locales()))
  85
  86         def add_plugin(self, plugin_class):
  87                 # Try initialising a new plugin. If that fails, we will log the
  88                 # error and try to go on.
  89                 try:
  90                         plugin = plugin_class(self)
  91                 except:
  92                         log.critical(_("Plugin %s could not be initialised") % plugin_class, exc_info=True)
  93                         return
  94
  95                 self.plugins.append(plugin)
  96
  97         @property
  98         def templates(self):
  99                 for plugin in self.plugins:
 100                         for template in plugin.templates:
 101                                 yield template
 102
 103         def run(self):
 104                 # Register signal handlers.
 105                 self.register_signal_handler()
 106
 107                 # Cannot do anything if no plugins have been initialised
 108                 if not self.plugins:
 109                         log.critical(_("No plugins have been initialised"))
 110                         return
 111
 112                 # Start the bus
 113                 self.bus.start()
 114
 115                 # Initialise the timer queue
 116                 self.initialise_timer_queue()
 117
 118                 # Start worker threads
 119                 for w in self.worker_threads:
 120                         w.start()
 121
 122                 # Run the write queue thread
 123                 self.write_queue.start()
 124
 125                 # Regularly submit all data to disk.
 126                 while self.running:
 127                         try:
 128                                 # Try processing one event from the queue. If that succeeded
 129                                 # we will retry immediately.
 130                                 if self.process_timer_queue():
 131                                         continue
 132
 133                                 # Otherwise we will sleep for a bit
 134                                 time.sleep(self.HEARTBEAT)
 135
 136                                 # Log warnings if the worker queue is filling up
 137                                 queue_size = self._worker_queue.qsize()
 138                                 if queue_size >= 5:
 139                                         log.warning(_("Worker queue is filling up with %s events") % queue_size)
 140
 141                         except KeyboardInterrupt:
 142                                 self.shutdown()
 143                                 break
 144
 145                 # Wait until all worker threads are finished
 146                 for w in self.worker_threads:
 147                         w.join()
 148
 149                 # Stop the bus thread
 150                 self.bus.shutdown()
 151
 152                 # Write all collected data to disk before ending the main thread
 153                 self.write_queue.shutdown()
 154
 155                 log.debug(_("Main thread exited"))
 156
 157         def shutdown(self):
 158                 if not self.running:
 159                         return
 160
 161                 log.info(_("Received shutdown signal"))
 162                 self.running = False
 163
 164                 # Propagating shutdown to all threads.
 165                 for w in self.worker_threads:
 166                         w.shutdown()
 167
 168         def register_signal_handler(self):
 169                 for s in (signal.SIGTERM, signal.SIGINT, signal.SIGUSR1):
 170                         log.debug(_("Registering signal %d") % s)
 171
 172                 signal.signal(s, self.signal_handler)
 173
 174         def signal_handler(self, sig, *args, **kwargs):
 175                 log.info(_("Caught signal %d") % sig)
 176
 177                 if sig in (signal.SIGTERM, signal.SIGINT):
 178                         # Shutdown this application.
 179                         self.shutdown()
 180
 181                 elif sig == signal.SIGUSR1:
 182                         # Commit all data.
 183                         self.write_queue.commit()
 184
 185         def get_plugin_from_template(self, template_name):
 186                 for plugin in self.plugins:
 187                         if not template_name in [t.name for t in plugin.templates]:
 188                                 continue
 189
 190                         return plugin
 191
 192         def generate_graph(self, template_name, *args, **kwargs):
 193                 plugin = self.get_plugin_from_template(template_name)
 194                 if not plugin:
 195                         raise RuntimeError("Could not find template %s" % template_name)
 196
 197                 return plugin.generate_graph(template_name, *args, **kwargs)
 198
 199         def graph_info(self, template_name, *args, **kwargs):
 200                 plugin = self.get_plugin_from_template(template_name)
 201                 if not plugin:
 202                         raise RuntimeError("Could not find template %s" % template_name)
 203
 204                 return plugin.graph_info(template_name, *args, **kwargs)
 205
 206         def last_update(self, template_name, *args, **kwargs):
 207                 plugin = self.get_plugin_from_template(template_name)
 208                 if not plugin:
 209                         raise RuntimeError("Could not find template %s" % template_name)
 210
 211                 return plugin.last_update(*args, **kwargs)
 212
 213         def create_worker_threads(self, num=None):
 214                 """
 215                         Creates a number of worker threads
 216                 """
 217                 # If no number of threads is given, we will create as many as we have
 218                 # active processor cores but never less than two.
 219                 if num is None:
 220                         num = max(multiprocessing.cpu_count(), 2)
 221
 222                 worker_threads = []
 223
 224                 for id in range(num):
 225                         worker_thread = WorkerThread(self, id)
 226                         worker_threads.append(worker_thread)
 227
 228                 return worker_threads
 229
 230         def initialise_timer_queue(self):
 231                 for p in self.plugins:
 232                         timer = PluginTimer(p)
 233
 234                         self._timer_queue.put(timer)
 235
 236         def process_timer_queue(self):
 237                 # Take the item from the timer queue that is to be due first
 238                 timer = self._timer_queue.get()
 239
 240                 try:
 241                         # If the timer event is to be executed, we will put the plugin
 242                         # into the worker queue and reset the timer
 243                         if timer.is_due():
 244                                 self._worker_queue.put(timer.plugin)
 245                                 timer.reset_deadline()
 246
 247                                 return timer
 248                 finally:
 249                         # Put the timer back into the timer queue.
 250                         self._timer_queue.put(timer)
 251
 252
 253 class WorkerThread(threading.Thread):
 254         HEARTBEAT = 2.5
 255
 256         def __init__(self, collecty, id):
 257                 threading.Thread.__init__(self)
 258                 self.daemon = True
 259
 260                 self.log = logging.getLogger("collecty.worker")
 261
 262                 self.collecty = collecty
 263                 self.id = id
 264
 265                 self.log.debug(_("Worker thread %s has been initialised") % self.id)
 266
 267         @property
 268         def queue(self):
 269                 """
 270                         The queue this thread is getting events from
 271                 """
 272                 return self.collecty._worker_queue
 273
 274         def run(self):
 275                 self.log.debug(_("Worker thread %s has been started") % self.id)
 276                 self.running = True
 277
 278                 while self.running:
 279                         try:
 280                                 plugin = self.queue.get(block=True, timeout=self.HEARTBEAT)
 281
 282                         # If the queue has been empty we just retry
 283                         except queue.Empty:
 284                                 continue
 285
 286                         # Execute the collect operation for this plugin
 287                         plugin.collect()
 288
 289                 self.log.debug(_("Worker thread %s has been terminated") % self.id)
 290
 291         def shutdown(self):
 292                 self.running = False
 293
 294
 295 class WriteQueue(threading.Thread):
 296         def __init__(self, collecty, submit_interval):
 297                 threading.Thread.__init__(self)
 298                 self.daemon = True
 299
 300                 self.collecty = collecty
 301
 302                 self.log = logging.getLogger("collecty.queue")
 303
 304                 self.timer = plugins.Timer(submit_interval)
 305                 self._queue = queue.PriorityQueue()
 306
 307                 self.log.debug(_("Initialised write queue"))
 308
 309         def run(self):
 310                 self.log.debug(_("Write queue process started"))
 311                 self.running = True
 312
 313                 while self.running:
 314                         # Reset the timer.
 315                         self.timer.reset()
 316
 317                         # Wait until the timer has successfully elapsed.
 318                         if self.timer.wait():
 319                                 self.commit()
 320
 321                 self.commit()
 322                 self.log.debug(_("Write queue process stopped"))
 323
 324         def shutdown(self):
 325                 self.running = False
 326                 self.timer.cancel()
 327
 328                 # Wait until all data has been written.
 329                 self.join()
 330
 331         def add(self, object, time, data):
 332                 result = QueueObject(object.file, time, data)
 333                 self._queue.put(result)
 334
 335         def commit(self):
 336                 """
 337                         Flushes the read data to disk.
 338                 """
 339                 # There is nothing to do if the queue is empty
 340                 if self._queue.empty():
 341                         self.log.debug(_("No data to commit"))
 342                         return
 343
 344                 time_start = time.time()
 345
 346                 self.log.debug(_("Submitting data to the databases..."))
 347
 348                 # Get all objects from the queue and group them by the RRD file
 349                 # to commit them all at once
 350                 results = {}
 351                 while not self._queue.empty():
 352                         result = self._queue.get()
 353
 354                         try:
 355                                 results[result.file].append(result)
 356                         except KeyError:
 357                                 results[result.file] = [result]
 358
 359                 # Write the collected data to disk
 360                 for filename, results in list(results.items()):
 361                         self._commit_file(filename, results)
 362
 363                 duration = time.time() - time_start
 364                 self.log.debug(_("Emptied write queue in %.2fs") % duration)
 365
 366         def _commit_file(self, filename, results):
 367                 self.log.debug(_("Committing %(counter)s entries to %(filename)s") \
 368                         % { "counter" : len(results), "filename" : filename })
 369
 370                 for result in results:
 371                         self.log.debug("  %s: %s" % (result.time, result.data))
 372
 373                 try:
 374                         rrdtool.update(filename, *["%s" % r for r in results])
 375
 376                 # Catch operational errors like unreadable/unwritable RRD databases
 377                 # or those where the format has changed. The collected data will be lost.
 378                 except rrdtool.OperationalError as e:
 379                         self.log.critical(_("Could not update RRD database %s: %s") \
 380                                 % (filename, e))
 381
 382         def commit_file(self, filename):
 383                 """
 384                         Commits all data that is in the write queue for the given
 385                         RRD database.
 386                 """
 387                 results, others = [], []
 388
 389                 # We will have to walk through the entire queue since we cannot
 390                 # ready any items selectively. Everything that belongs to our
 391                 # transaction is kept. Everything else will be put back into the
 392                 # queue.
 393                 while not self._queue.empty():
 394                         result = self._queue.get()
 395
 396                         if result.file == filename:
 397                                 results.append(result)
 398                         else:
 399                                 others.append(result)
 400
 401                 # Put back all items that did not match
 402                 for result in others:
 403                         self._queue.put(result)
 404
 405                 # Write everything else to disk
 406                 if results:
 407                         self._commit_file(filename, results)
 408
 409
 410 class QueueObject(object):
 411         def __init__(self, file, time, data):
 412                 self.file = file
 413                 self.time = time
 414                 self.data = data
 415
 416         def __str__(self):
 417                 return "%s:%s" % (self.time.strftime("%s"), self.data)
 418
 419         def __lt__(self, other):
 420                 return self.time < other.time
 421
 422
 423 class PluginTimer(object):
 424         def __init__(self, plugin):
 425                 self.plugin = plugin
 426
 427                 self.deadline = datetime.datetime.utcnow()
 428
 429         def __repr__(self):
 430                 return "<%s %s>" % (self.__class__.__name__, self.deadline)
 431
 432         def __lt__(self, other):
 433                 return self.deadline < other.deadline
 434
 435         def reset_deadline(self):
 436                 self.deadline = datetime.datetime.utcnow() \
 437                         + datetime.timedelta(seconds=self.plugin.interval)
 438
 439         def is_due(self):
 440                 return datetime.datetime.utcnow() >= self.deadline