Add "variants" field to LUDiagnoseOS
[ganeti-local] / lib / daemon.py
index e363be3..d357349 100644 (file)
 """Module with helper classes and functions for daemons"""
 
 
+import asyncore
 import os
 import select
 import signal
 import errno
 import logging
+import sched
+import time
+import socket
+import sys
 
 from ganeti import utils
 from ganeti import constants
+from ganeti import errors
+
+
+class SchedulerBreakout(Exception):
+  """Exception used to get out of the scheduler loop
+
+  """
+
+
+def AsyncoreDelayFunction(timeout):
+  """Asyncore-compatible scheduler delay function.
+
+  This is a delay function for sched that, rather than actually sleeping,
+  executes asyncore events happening in the meantime.
+
+  After an event has occurred, rather than returning, it raises a
+  SchedulerBreakout exception, which will force the current scheduler.run()
+  invocation to terminate, so that we can also check for signals. The main loop
+  will then call the scheduler run again, which will allow it to actually
+  process any due events.
+
+  This is needed because scheduler.run() doesn't support a count=..., as
+  asyncore loop, and the scheduler module documents throwing exceptions from
+  inside the delay function as an allowed usage model.
+
+  """
+  asyncore.loop(timeout=timeout, count=1, use_poll=True)
+  raise SchedulerBreakout()
+
+
+class AsyncoreScheduler(sched.scheduler):
+  """Event scheduler integrated with asyncore
+
+  """
+  def __init__(self, timefunc):
+    sched.scheduler.__init__(self, timefunc, AsyncoreDelayFunction)
+
+
+class AsyncUDPSocket(asyncore.dispatcher):
+  """An improved asyncore udp socket.
+
+  """
+  def __init__(self):
+    """Constructor for AsyncUDPSocket
+
+    """
+    asyncore.dispatcher.__init__(self)
+    self._out_queue = []
+    self.create_socket(socket.AF_INET, socket.SOCK_DGRAM)
+
+  # this method is overriding an asyncore.dispatcher method
+  def handle_connect(self):
+    # Python thinks that the first udp message from a source qualifies as a
+    # "connect" and further ones are part of the same connection. We beg to
+    # differ and treat all messages equally.
+    pass
+
+  # this method is overriding an asyncore.dispatcher method
+  def handle_read(self):
+    try:
+      try:
+        payload, address = self.recvfrom(constants.MAX_UDP_DATA_SIZE)
+      except socket.error, err:
+        if err.errno == errno.EINTR:
+          # we got a signal while trying to read. no need to do anything,
+          # handle_read will be called again if there is data on the socket.
+          return
+        else:
+          raise
+      ip, port = address
+      self.handle_datagram(payload, ip, port)
+    except:
+      # we need to catch any exception here, log it, but proceed, because even
+      # if we failed handling a single request, we still want to continue.
+      logging.error("Unexpected exception", exc_info=True)
+
+  def handle_datagram(self, payload, ip, port):
+    """Handle an already read udp datagram
+
+    """
+    raise NotImplementedError
+
+  # this method is overriding an asyncore.dispatcher method
+  def writable(self):
+    # We should check whether we can write to the socket only if we have
+    # something scheduled to be written
+    return bool(self._out_queue)
+
+  def handle_write(self):
+    try:
+      if not self._out_queue:
+        logging.error("handle_write called with empty output queue")
+        return
+      (ip, port, payload) = self._out_queue[0]
+      try:
+        self.sendto(payload, 0, (ip, port))
+      except socket.error, err:
+        if err.errno == errno.EINTR:
+          # we got a signal while trying to write. no need to do anything,
+          # handle_write will be called again because we haven't emptied the
+          # _out_queue, and we'll try again
+          return
+        else:
+          raise
+      self._out_queue.pop(0)
+    except:
+      # we need to catch any exception here, log it, but proceed, because even
+      # if we failed sending a single datagram we still want to continue.
+      logging.error("Unexpected exception", exc_info=True)
+
+  def enqueue_send(self, ip, port, payload):
+    """Enqueue a datagram to be sent when possible
+
+    """
+    if len(payload) > constants.MAX_UDP_DATA_SIZE:
+      raise errors.UdpDataSizeError('Packet too big: %s > %s' % (len(payload),
+                                    constants.MAX_UDP_DATA_SIZE))
+    self._out_queue.append((ip, port, payload))
 
 
 class Mainloop(object):
@@ -39,99 +162,50 @@ class Mainloop(object):
   def __init__(self):
     """Constructs a new Mainloop instance.
 
+    @ivar scheduler: A L{sched.scheduler} object, which can be used to register
+    timed events
+
     """
-    self._io_wait = {}
-    self._io_wait_add = []
-    self._io_wait_remove = []
     self._signal_wait = []
+    self.scheduler = AsyncoreScheduler(time.time)
 
-  def Run(self, handle_sigchld=True, handle_sigterm=True, stop_on_empty=False):
+  @utils.SignalHandled([signal.SIGCHLD])
+  @utils.SignalHandled([signal.SIGTERM])
+  def Run(self, stop_on_empty=False, signal_handlers=None):
     """Runs the mainloop.
 
-    @type handle_sigchld: bool
-    @param handle_sigchld: Whether to install handler for SIGCHLD
-    @type handle_sigterm: bool
-    @param handle_sigterm: Whether to install handler for SIGTERM
     @type stop_on_empty: bool
     @param stop_on_empty: Whether to stop mainloop once all I/O waiters
                           unregistered
+    @type signal_handlers: dict
+    @param signal_handlers: signal->L{utils.SignalHandler} passed by decorator
 
     """
-    poller = select.poll()
-
-    # Setup signal handlers
-    if handle_sigchld:
-      sigchld_handler = utils.SignalHandler([signal.SIGCHLD])
-    else:
-      sigchld_handler = None
-    try:
-      if handle_sigterm:
-        sigterm_handler = utils.SignalHandler([signal.SIGTERM])
+    assert isinstance(signal_handlers, dict) and \
+           len(signal_handlers) > 0, \
+           "Broken SignalHandled decorator"
+    running = True
+    # Start actual main loop
+    while running:
+      # Stop if nothing is listening anymore
+      if stop_on_empty and not (self._io_wait):
+        break
+
+      if not self.scheduler.empty():
+        try:
+          self.scheduler.run()
+        except SchedulerBreakout:
+          pass
       else:
-        sigterm_handler = None
+        asyncore.loop(count=1, use_poll=True)
 
-      try:
-        running = True
-
-        # Start actual main loop
-        while running:
-          # Entries could be added again afterwards, hence removing first
-          if self._io_wait_remove:
-            for fd in self._io_wait_remove:
-              try:
-                poller.unregister(fd)
-              except KeyError:
-                pass
-              try:
-                del self._io_wait[fd]
-              except KeyError:
-                pass
-            self._io_wait_remove = []
-
-          # Add new entries
-          if self._io_wait_add:
-            for (owner, fd, conditions) in self._io_wait_add:
-              self._io_wait[fd] = owner
-              poller.register(fd, conditions)
-            self._io_wait_add = []
-
-          # Stop if nothing is listening anymore
-          if stop_on_empty and not (self._io_wait):
-            break
-
-          # Wait for I/O events
-          try:
-            io_events = poller.poll(None)
-          except select.error, err:
-            # EINTR can happen when signals are sent
-            if err.args and err.args[0] in (errno.EINTR,):
-              io_events = None
-            else:
-              raise
-
-          if io_events:
-            # Check for I/O events
-            for (evfd, evcond) in io_events:
-              owner = self._io_wait.get(evfd, None)
-              if owner:
-                owner.OnIO(evfd, evcond)
-
-          # Check whether signal was raised
-          if sigchld_handler and sigchld_handler.called:
-            self._CallSignalWaiters(signal.SIGCHLD)
-            sigchld_handler.Clear()
-
-          if sigterm_handler and sigterm_handler.called:
-            self._CallSignalWaiters(signal.SIGTERM)
-            running = False
-            sigterm_handler.Clear()
-      finally:
-        # Restore signal handlers
-        if sigterm_handler:
-          sigterm_handler.Reset()
-    finally:
-      if sigchld_handler:
-        sigchld_handler.Reset()
+      # Check whether a signal was raised
+      for sig in signal_handlers:
+        handler = signal_handlers[sig]
+        if handler.called:
+          self._CallSignalWaiters(sig)
+          running = (sig != signal.SIGTERM)
+          handler.Clear()
 
   def _CallSignalWaiters(self, signum):
     """Calls all signal waiters for a certain signal.
@@ -141,42 +215,7 @@ class Mainloop(object):
 
     """
     for owner in self._signal_wait:
-      owner.OnSignal(signal.SIGCHLD)
-
-  def RegisterIO(self, owner, fd, condition):
-    """Registers a receiver for I/O notifications
-
-    The receiver must support a "OnIO(self, fd, conditions)" function.
-
-    @type owner: instance
-    @param owner: Receiver
-    @type fd: int
-    @param fd: File descriptor
-    @type condition: int
-    @param condition: ORed field of conditions to be notified
-                      (see select module)
-
-    """
-    # select.Poller also supports file() like objects, but we don't.
-    assert isinstance(fd, (int, long)), \
-      "Only integers are supported for file descriptors"
-
-    self._io_wait_add.append((owner, fd, condition))
-
-  def UnregisterIO(self, fd):
-    """Unregister a file descriptor.
-
-    It'll be unregistered the next time the mainloop checks for it.
-
-    @type fd: int
-    @param fd: File descriptor
-
-    """
-    # select.Poller also supports file() like objects, but we don't.
-    assert isinstance(fd, (int, long)), \
-      "Only integers are supported for file descriptors"
-
-    self._io_wait_remove.append(fd)
+      owner.OnSignal(signum)
 
   def RegisterSignal(self, owner):
     """Registers a receiver for signal notifications
@@ -272,4 +311,3 @@ def GenericMain(daemon_name, optionparser, dirs, check_fn, exec_fn):
     exec_fn(options, args)
   finally:
     utils.RemovePidFile(daemon_name)
-