Changeset 3660

Timestamp:

12/11/08 03:22:46 (2 months ago)

Author:

tack

Message:

Signals.add for convenience (self.signals += ('foo', 'bar'));
Socket.write() now returns InProgress?, which is finished when the data passed
to that write() call is actually written to the socket.

In preparation for Process rewrite, factor out logic that will eventually be
shared between Socket and Process into a common base class called IODescriptor.
This isn't terribly well tested, hopefully I didn't break much.

Location:

trunk/base/src/notifier

Files:

• signals.py (modified) (1 diff)
• sockets.py (modified) (16 diffs)

trunk/base/src/notifier/signals.py

@@ -315 +315 @@
 
 
+    def __add__(self, signals):
+        return Signals(self, *signals)
+
+
     def add(self, *signals):
         """

trunk/base/src/notifier/sockets.py

@@ -38 +38 @@
 import errno
 import logging
+import time
 
 import nf_wrapper as notifier
@@ -47 +48 @@
 from kaa.tmpfile import tempfile
 
+# FIXME: this file is getting big enough: move IOMonitor and IODescriptor to
+# io.py, and leave Socket in this file.
+
 # get logging object
 log = logging.getLogger('notifier')
 
-IO_READ   = 0
-IO_WRITE  = 1
+IO_READ   = 1
+IO_WRITE  = 2
 
 class IOMonitor(notifier.NotifierCallback):
@@ -64 +68 @@
     def register(self, fd, condition = IO_READ):
         """
-        Register the IOMonitor to a specific socket
+        Register the IOMonitor to a specific file descriptor
         @param fd: File descriptor or Python socket object
         @param condition: IO_READ or IO_WRITE
@@ -74 +78 @@
         if not is_mainthread():
             return MainThreadCallback(self.register)(fd, condition)
-        notifier.socket_add(fd, self, condition)
+        notifier.socket_add(fd, self, condition-1)
         self._condition = condition
         # Must be called _id to correspond with base class.
@@ -88 +92 @@
         if not is_mainthread():
             return MainThreadCallback(self.unregister)()
-        notifier.socket_remove(self._id, self._condition)
+        notifier.socket_remove(self._id, self._condition-1)
         super(IOMonitor, self).unregister()
 
@@ -102 +106 @@
 # We need to import main for the signals dict (we add a handler to
 # shutdown to gracefully close sockets), but main itself imports us
-# for IOHandler.  So we must import main _after_ declaring IOHandler,
+# for IOMonitor.  So we must import main _after_ declaring IOMonitor,
 # instead of doing so at the top of the file.
 from kaa.notifier import main
 
-class SocketError(Exception):
-    pass
-
-class Socket(object):
+class IODescriptor(object):
     """
-    Notifier-aware socket class.
+    Base class for read-only, write-only or read-write descriptors such as
+    Socket and Process.  Implements logic common to communication over
+    descriptors such as async read/writes and read/write buffering.
     """
-
-    def __init__(self, buffer_size=None, chunk_size=1024*1024):
-        self.signals = Signals('closed', 'read', 'readline', 'new-client')
-        self._socket = None
-        self._connecting = False
-        self._write_buffer = []
-        self._addr = None
-        self._listening = False
+    def __init__(self, fd=None, mode=IO_READ|IO_WRITE, chunk_size=1024*1024):
+        self.signals = Signals('closed', 'read', 'readline', 'write')
+        self.wrap(fd, mode)
+        self._write_queue = []
         self._queue_close = False
-        self._buffer_size = buffer_size
         self._chunk_size = chunk_size
-
+    
         # Internal signals for read() and readline()  (these are different from
-        # the public signals 'read' and 'readline' as they get emitted even
-        # when data is None.  When these signals get updated, we call
-        # _update_read_monitor to register the read IOMonitor.
+        # the same-named public signals as they get emitted even when data is
+        # None.  When these signals get updated, we call _update_read_monitor
+        # to register the read IOMonitor.
         cb = WeakCallback(self._update_read_monitor)
         self._read_signal = Signal(cb)
@@ -136 +134 @@
 
         # These variables hold the IOMonitors for monitoring; we only allocate
-        # a monitor when the socket is connected to avoid a ref cycle so
-        # that disconnected sockets will get properly deleted when they are not
+        # a monitor when the fd is connected to avoid a ref cycle so that
+        # disconnected fds will get properly deleted when they are not
         # referenced.
-        self._rmon = self._wmon = None
+        self._rmon = None
+        self._wmon = None
+
+
+    @property
+    def alive(self):
+        """
+        Returns True if the fd exists and is open.
+        """
+        # If the fd is closed, self._fd will be None.
+        return self._fd != None
+
+
+    @property
+    def fileno(self):
+        """
+        Returns the integer for this file descriptor, or None if no descriptor
+        has been set.
+        """
+        try:
+            return self._fd.fileno()
+        except AttributeError:
+            return self._fd
+
+
+    @property
+    def chunk_size(self):
+        """
+        Number of bytes to attempt to read from the fd at a time.  The
+        default is 1M.  A 'read' signal is emitted for each chunk read from the
+        fd.  (The number of bytes read at a time may be less than the chunk
+        size, but will never be more.)
+        """
+        return self._chunk_size
+
+
+    @chunk_size.setter
+    def chunk_size(self, size):
+        self._chunk_size = size
+
+
+    # TODO: settable write_queue_max property
+
+    @property
+    def write_queue_size(self):
+        """
+        Returns the number of bytes queued in memory to be written to the
+        descriptor.
+        """
+        # XXX: this is not terribly efficient when the write queue has
+        # many elements.  We may decide to keep a separate counter.
+        return sum(len(data) for data, inprogress in self._write_queue)
+
+
+    def _is_read_connected(self):
+        """
+        Returns True if we're interested in read events.
+        """
+        return not len(self._read_signal) == len(self._readline_signal) == \
+                   len(self.signals['read']) == len(self.signals['readline']) == 0
+        
+
+    def _update_read_monitor(self, signal=None, change=None):
+        """
+        Update read IOMonitor to register or unregister based on if there are
+        any handlers attached to the read signals.  If there are no handlers,
+        there is no point in reading data from the descriptor since it will go 
+        nowhere.  This also allows us to push back the read buffer to the OS.
+
+        We must call this immediately after reading a block, and not defer
+        it until the end of the mainloop iteration via a timer in order not
+        to lose incoming data between read() calls.
+        """
+        if not (self._mode & IO_READ) or not self._rmon or change == Signal.SIGNAL_DISCONNECTED:
+            return
+        elif not self._is_read_connected():
+            self._rmon.unregister()
+        elif not self._rmon.active():
+            self._rmon.register(self.fileno, IO_READ)
+
+
+    def _set_non_blocking(self):
+        """
+        Low-level call to set the fd non-blocking.  Can be overridden by
+        subclasses.
+        """
+        flags = fcntl.fcntl(self.fileno, fcntl.F_GETFL)
+        fcntl.fcntl(self.fileno, fcntl.F_SETFL, flags | os.O_NONBLOCK)
+
+
+    def wrap(self, fd, mode):
+        """
+        Wraps an existing file descriptor.
+        """
+        self._fd = fd
+        self._mode = mode
+        if not fd:
+            return
+        self._set_non_blocking()
+
+        if self._mode & IO_READ:
+            if self._rmon:
+                self._rmon.unregister()
+            self._rmon = IOMonitor(self._handle_read)
+            self._update_read_monitor()
+
+        if self._mode & IO_WRITE:
+            if self._wmon:
+                self._wmon.unregister()
+            self._wmon = IOMonitor(self._handle_write)
+            if self._write_queue:
+                self._wmon.register(self.fileno, IO_WRITE)
+
+        # Disconnect socket and remove socket file (if unix socket) on shutdown
+        main.signals['shutdown'].connect_weak(self.close)
+
+
+    def _is_readable(self):
+        """
+        Low-level call to read from fd.  Can be overridden by subclasses.
+        """
+        return self._fd != None
+
+
+    def _async_read(self, signal):
+        """
+        Common implementation for read() and readline().
+        """
+        if not (self._mode & IO_READ):
+            raise IOError(9, 'Cannot read on a write-only descriptor')
+        if not self._is_readable():
+            # fd is not readable.  Return an InProgress pre-finished
+            # with None
+            return InProgress().finish(None)
+
+        return inprogress(signal)
+
+
+    def read(self):
+        """
+        Reads a chunk of data from the fd.  This function returns an 
+        InProgress object.  If the InProgress is finished with None, it
+        means that no data was collected and the fd closed.
+
+        It is therefore possible to busy-loop by reading on a closed
+        fd::
+
+            while True:
+                fd.read().wait()
+
+        So the return value of read() should be checked.  Alternatively,
+        fd.alive could be tested::
+
+            while fd.alive:
+                fd.read().wait()
+
+        """
+        return self._async_read(self._read_signal)
+
+
+    def readline(self):
+        """
+        Reads a line from the fd (with newline stripped).  The function
+        returns an InProgress object.  If the InProgress is finished with None
+        or the empty string, it means that no data was collected and the socket
+        closed.
+        """
+        return self._async_read(self._readline_signal)
+
+
+    def _read(self, size):
+        """
+        Low-level call to read from fd.  Can be overridden by subclasses.
+        """
+        return os.read(self.fileno, size)
+
+
+    def _handle_read(self):
+        """
+        IOMonitor callback when there is data to be read from the fd.
+        
+        This callback is only registered when we know the user is interested in
+        reading data (by connecting to the read or readline signals, or calling
+        read() or readline()).  This is necessary for flow control.
+        """
+        try:
+            data = self._read(self._chunk_size)
+        except (IOError, socket.error), (errno, msg):
+            if errno == 11:
+                # Resource temporarily unavailable -- we are trying to read
+                # data on a socket when none is available.
+                return
+            # If we're here, then the socket is likely disconnected.
+            data = None
+        except:
+            log.exception('%s._handle_read failed with unknown exception, closing socket', self.__class__.__name__)
+            data = None
+
+        # _read_signal is for InProgress objects waiting on the next read().
+        # For these we must emit even when data is None.
+        self._read_signal.emit(data)
+
+        if not data:
+            self._readline_signal.emit(data)
+            return self.close(immediate=True, expected=False)
+
+        self.signals['read'].emit(data)
+        # Update read monitor if necessary.  If there are no longer any
+        # callbacks left on any of the read signals (most likely _read_signal
+        # or _readline_signal), we want to prevent _handle_read() from being
+        # called, otherwise next time read() or readline() is called, we will
+        # have lost that data.
+        self._update_read_monitor()
+
+
+        # TODO: parse input into separate lines and emit readline.
+
+
+    def _write(self, data):
+        """
+        Low-level call to write to the fd  Can be overridden by subclasses.
+        """
+        return os.write(self.fileno, data)
+
+
+    def write(self, data):
+        """
+        Writes the given data to the fd.  This method returns an InProgress
+        object which is finished when the given data is fully written to the
+        file descriptor.
+
+        It is not required that the descriptor be open in order to write to it.
+        Written data is queued until the descriptor open and then flushed.  As
+        writes are asynchronous, all written data is queued.  It is the
+        caller's responsibility to ensure the internal write queue does not
+        exceed the desired size by waiting for past write() InProgress to
+        finish before writing more data.
+
+        If a write does not complete because the file descriptor was closed
+        prematurely, an IOError is thrown to the InProgress.
+        """
+        if not (self._mode & IO_WRITE):
+            raise IOError(9, 'Cannot write to a read-only descriptor')
+
+        inprogress = InProgress()
+        if data:
+            self._write_queue.append((data, inprogress))
+            if self._fd and self._wmon and not self._wmon.active():
+                self._wmon.register(self.fileno, IO_WRITE)
+        else:
+            # We're writing the null string, nothing really to do.  We're
+            # implicitly done.
+            inprogress.finish(0)
+        return inprogress
+
+
+    def _handle_write(self):
+        """
+        IOMonitor callback when the fd is writable.  This callback is not
+        registered then the write queue is empty, so we only get called when
+        there is something to write.
+        """
+        if not self._write_queue:
+            # Shouldn't happen; sanity check.
+            return
+
+        try:
+            while self._write_queue:
+                data, inprogress = self._write_queue.pop(0)
+                sent = self._write(data)
+                if sent != len(data):
+                    # Not all data was able to be sent; push remaining data
+                    # back onto the write buffer.
+                    self._write_queue.insert(0, (data[sent:], inprogress))
+                    break
+                else:
+                    # All data is written, finish the InProgress associated
+                    # with this write.
+                    inprogress.finish(sent)
+
+            if not self._write_queue:
+                if self._queue_close:
+                    return self.close(immediate=True)
+                self._wmon.unregister()
+
+        except (IOError, socket.error), (errno, msg):
+            if errno == 11:
+                # Resource temporarily unavailable -- we are trying to write
+                # data to a socket when none is available.  To prevent a busy
+                # loop (notifier loop will keep calling us back) we sleep a
+                # tiny bit.  It's admittedly a bit kludgy, but it's a simple
+                # solution to a condition which should not occur often.
+                time.sleep(0.001)
+                return
+
+            # If we're here, then the socket is likely disconnected.
+            self.close(immediate=True, expected=False)
+
+
+    def _close(self):
+        """
+        Low-level call to close the fd  Can be overridden by subclasses.
+        """
+        os.close(self.fileno)
+
+
+    def close(self, immediate=False, expected=True):
+        """
+        Closes the fd.  If immediate is False and there is data in the
+        write buffer, the fd is closed once the write buffer is emptied.
+        Otherwise the fd is closed immediately and the 'closed' signal
+        is emitted.
+        """
+        if not immediate and self._write_queue:
+            # Immediate close not requested and we have some data left
+            # to be written, so defer close until after write queue
+            # is empty.
+            self._queue_close = True
+            return
+
+        if not self._rmon and not self._wmon:
+            # already closed
+            return
+
+        if self._rmon:
+            self._rmon.unregister()
+        if self._wmon:
+            self._wmon.unregister()
+        self._rmon = None
+        self._wmon = None
+        self._queue_close = False
+
+        # Throw IOError to any pending InProgress in the write queue
+        for data, inprogress in self._write_queue:
+            if len(inprogress):
+                # Somebody cares about this InProgress, so we need to finish
+                # it.
+                inprogress.throw(IOError, IOError(9, "Descriptor closed prematurely"), None)
+        del self._write_queue[:]
+
+        self._close()
+        self._fd = None
+
+        self.signals['closed'].emit(expected)
+        main.signals['shutdown'].disconnect(self.close)
+
+
+class SocketError(Exception):
+    pass
+
+class Socket(IODescriptor):
+    """
+    Notifier-aware socket class, implementing fully asynchronous reads
+    and writes.
+    """
+
+    def __init__(self, buffer_size=None, chunk_size=1024*1024):
+        self._connecting = False
+        self._addr = None
+        self._listening = False
+        self._buffer_size = buffer_size
+
+        super(Socket, self).__init__(chunk_size=chunk_size)
+        self.signals += ('new-client',)
+
 
     def __repr__(self):
-        if not self._socket:
+        if not self._fd:
             return '<kaa.Socket - disconnected>'
         return '<kaa.Socket fd=%d>' % self.fileno
@@ -186 +548 @@
         try:
             # Will raise exception if socket is not connected.
-            self._socket.getpeername()
+            self._fd.getpeername()
             return True
         except:
@@ -216 +578 @@
     def buffer_size(self, size):
         self._buffer_size = size
-        if self._socket and size:
-            self._set_buffer_size(self._socket, size)
-
-
-    @property
-    def chunk_size(self):
-        """
-        Number of bytes to attempt to read from the socket at a time.  The
-        default is 1M.  A 'read' signal is emitted for each chunk read from the
-        socket.  (The number of bytes read at a time may be less than the chunk
-        size, but will never be more.)
-        """
-        return self._chunk_size
-
-
-    @chunk_size.setter
-    def chunk_size(self, size):
-        self._chunk_size = size
-
-
-    @property
-    def fileno(self):
-        """
-        Returns the file descriptor of the socket, or None if the socket is
-        not connected.
-        """
-        if not self._socket:
-            return None
-        return self._socket.fileno()
+        if self._fd and size:
+            self._set_buffer_size(self._fd, size)
 
 
@@ -254 +589 @@
         s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, size)
         
-
-    def _update_read_monitor(self, signal = None, change = None):
-        """
-        Update read IOMonitor to register or unregister based on if there are
-        any handlers attached to the read signals.  If there are no handlers,
-        there is no point in reading data from the socket since it will go 
-        nowhere.  This also allows us to push back the read buffer to the OS.
-
-        We must call this immediately after reading a block, and not defer
-        it until the end of the mainloop iteration via a timer in order not
-        to lose incoming data between read() calls.
-        """
-        if not self._rmon or change == Signal.SIGNAL_DISCONNECTED:
-            return
-        elif not self._listening and len(self._read_signal) == len(self._readline_signal) == \
-                                     len(self.signals['read']) == len(self.signals['readline']) == 0:
-            self._rmon.unregister()
-        elif not self._rmon.active():
-            self._rmon.register(self._socket, IO_READ)
-
 
     def _normalize_address(self, addr):
@@ -297 +612 @@
 
 
-    def _make_socket(self, addr = None, overwrite = False):
+    def _make_socket(self, addr=None, overwrite=False):
         """
         Constructs a socket based on the given addr.  Returns the socket and
@@ -337 +652 @@
 
 
-    def _replace_socket(self, sock, addr):
-        """
-        Replaces the existing socket and address spec with the ones supplied.
-        Any existing socket is closed.
-        """
-        if self._socket:
-            self._socket.close()
-
-        self._socket, self._addr = sock, addr
-
-
-    def listen(self, bind_info, qlen = 5):
+    def listen(self, bind_info, qlen=5):
         """
         Sets the socket to listen on bind_info, which is either an integer
@@ -413 +717 @@
 
 
-    def wrap(self, sock, addr = None):
+    def wrap(self, sock, addr=None):
         """
         Wraps an existing low-level socket object.  addr specifies the address
         corresponding to the socket.
         """
-        self._socket = sock or self._socket
+        super(Socket, self).wrap(sock, IO_READ|IO_WRITE)
         self._addr = addr or self._addr
-        self._queue_close = False
-
-        sock.setblocking(False)
+
         if self._buffer_size:
             self._set_buffer_size(sock, self._buffer_size)
 
-        if self._rmon:
-            self._rmon.unregister()
-            self._wmon.unregister()
-
-        self._rmon = IOMonitor(self._handle_read)
-        self._wmon = IOMonitor(self._handle_write)
-
-        self._update_read_monitor()
-        if self._write_buffer:
-            self._wmon.register(sock, IO_WRITE)
-
-        # Disconnect socket and remove socket file (if unix socket) on shutdown
-        main.signals['shutdown'].connect_weak(self.close)
-        
-
-    def _async_read(self, signal):
-        if not self._socket and not self._connecting:
-            # Socket is closed.  Return an InProgress pre-finished
-            # with None
-            return InProgress().finish(None)
-
-        if self._listening:
-            raise SocketError("Can't read on a listening socket.")
-
-        return inprogress(signal)
-
-
-    def read(self):
-        """
-        Reads a chunk of data from the socket.  This function returns an 
-        InProgress object.  If the InProgress is finished with None, it
-        means that no data was collected and the socket closed.
-
-        It is therefore possible to busy-loop by reading on a closed
-        socket::
-
-            while True:
-                socket.read().wait()
-
-        So the return value of read() should be checked.  Alternatively,
-        socket.alive could be tested::
-
-            while socket.alive:
-                socket.read().wait()
-
-        """
-        return self._async_read(self._read_signal)
-
-
-    def readline(self):
-        """
-        Reads a line from the socket (with newline stripped).  The function
-        returns an InProgress object.  If the InProgress is finished with
-        None, it means that no data was collected and the socket closed.
-        """
-        return self._async_read(self._readline_signal)
+
+    def _is_read_connected(self):
+        return self._listening or super(Socket, self)._is_read_connected()
+
+
+    def _set_non_blocking(self):
+        self._fd.setblocking(False)
+
+
+    def _is_readable(self):
+        return self._fd and not self._connecting
+
+
+    def _read(self, size):
+        return self._fd.recv(size)
+
+
+    def _write(self, data):
+        return self._fd.send(data)
+
+
+    def _close(self):
+        self._fd.close()
 
 
@@ -488 +757 @@
         Accept a new connection and return a new Socket object.
         """
-        sock, addr = self._socket.accept()
+        sock, addr = self._fd.accept()
         # create new Socket from the same class this object is
         client_socket = self.__class__()
@@ -499 +768 @@
             return self._accept()
 
-        try:
-            data = self._socket.recv(self._chunk_size)
-        except socket.error, (errno, msg):
-            if errno == 11:
-                # Resource temporarily unavailable -- we are trying to read