Home | History | Annotate | Download | only in _multiprocessing 
      1 /*
      2  * A type which wraps a socket
      3  *
      4  * socket_connection.c
      5  *
      6  * Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
      7  */
      8 
      9 #include "multiprocessing.h"
     10 
     11 #if defined(HAVE_POLL) && !defined(HAVE_BROKEN_POLL)
     12 #  include "poll.h"
     13 #endif
     14 
     15 #ifdef MS_WINDOWS
     16 #  define WRITE(h, buffer, length) send((SOCKET)h, buffer, length, 0)
     17 #  define READ(h, buffer, length) recv((SOCKET)h, buffer, length, 0)
     18 #  define CLOSE(h) closesocket((SOCKET)h)
     19 #else
     20 #  define WRITE(h, buffer, length) write(h, buffer, length)
     21 #  define READ(h, buffer, length) read(h, buffer, length)
     22 #  define CLOSE(h) close(h)
     23 #endif
     24 
     25 /*
     26  * Wrapper for PyErr_CheckSignals() which can be called without the GIL
     27  */
     28 
     29 static int
     30 check_signals(void)
     31 {
     32     PyGILState_STATE state;
     33     int res;
     34     state = PyGILState_Ensure();
     35     res = PyErr_CheckSignals();
     36     PyGILState_Release(state);
     37     return res;
     38 }
     39 
     40 /*
     41  * Send string to file descriptor
     42  */
     43 
     44 static Py_ssize_t
     45 _conn_sendall(HANDLE h, char *string, size_t length)
     46 {
     47     char *p = string;
     48     Py_ssize_t res;
     49 
     50     while (length > 0) {
     51         res = WRITE(h, p, length);
     52         if (res < 0) {
     53             if (errno == EINTR) {
     54                 if (check_signals() < 0)
     55                     return MP_EXCEPTION_HAS_BEEN_SET;
     56                 continue;
     57             }
     58             return MP_SOCKET_ERROR;
     59         }
     60         length -= res;
     61         p += res;
     62     }
     63 
     64     return MP_SUCCESS;
     65 }
     66 
     67 /*
     68  * Receive string of exact length from file descriptor
     69  */
     70 
     71 static Py_ssize_t
     72 _conn_recvall(HANDLE h, char *buffer, size_t length)
     73 {
     74     size_t remaining = length;
     75     Py_ssize_t temp;
     76     char *p = buffer;
     77 
     78     while (remaining > 0) {
     79         temp = READ(h, p, remaining);
     80         if (temp < 0) {
     81             if (errno == EINTR) {
     82                 if (check_signals() < 0)
     83                     return MP_EXCEPTION_HAS_BEEN_SET;
     84                 continue;
     85             }
     86             return temp;
     87         }
     88         else if (temp == 0) {
     89             return remaining == length ? MP_END_OF_FILE : MP_EARLY_END_OF_FILE;
     90         }
     91         remaining -= temp;
     92         p += temp;
     93     }
     94 
     95     return MP_SUCCESS;
     96 }
     97 
     98 /*
     99  * Send a string prepended by the string length in network byte order
    100  */
    101 
    102 static Py_ssize_t
    103 conn_send_string(ConnectionObject *conn, char *string, size_t length)
    104 {
    105     Py_ssize_t res;
    106     /* The "header" of the message is a 32 bit unsigned number (in
    107        network order) which specifies the length of the "body".  If
    108        the message is shorter than about 16kb then it is quicker to
    109        combine the "header" and the "body" of the message and send
    110        them at once. */
    111     if (length < (16*1024)) {
    112         char *message;
    113 
    114         message = PyMem_Malloc(length+4);
    115         if (message == NULL)
    116             return MP_MEMORY_ERROR;
    117 
    118         *(UINT32*)message = htonl((UINT32)length);
    119         memcpy(message+4, string, length);
    120         Py_BEGIN_ALLOW_THREADS
    121         res = _conn_sendall(conn->handle, message, length+4);
    122         Py_END_ALLOW_THREADS
    123         PyMem_Free(message);
    124     } else {
    125         UINT32 lenbuff;
    126 
    127         if (length > MAX_MESSAGE_LENGTH)
    128             return MP_BAD_MESSAGE_LENGTH;
    129 
    130         lenbuff = htonl((UINT32)length);
    131         Py_BEGIN_ALLOW_THREADS
    132         res = _conn_sendall(conn->handle, (char*)&lenbuff, 4) ||
    133             _conn_sendall(conn->handle, string, length);
    134         Py_END_ALLOW_THREADS
    135     }
    136     return res;
    137 }
    138 
    139 /*
    140  * Attempts to read into buffer, or failing that into *newbuffer
    141  *
    142  * Returns number of bytes read.
    143  */
    144 
    145 static Py_ssize_t
    146 conn_recv_string(ConnectionObject *conn, char *buffer,
    147                  size_t buflength, char **newbuffer, size_t maxlength)
    148 {
    149     Py_ssize_t res;
    150     UINT32 ulength;
    151 
    152     *newbuffer = NULL;
    153 
    154     Py_BEGIN_ALLOW_THREADS
    155     res = _conn_recvall(conn->handle, (char*)&ulength, 4);
    156     Py_END_ALLOW_THREADS
    157     if (res < 0)
    158         return res;
    159 
    160     ulength = ntohl(ulength);
    161     if (ulength > maxlength)
    162         return MP_BAD_MESSAGE_LENGTH;
    163 
    164     if (ulength > buflength) {
    165         *newbuffer = buffer = PyMem_Malloc((size_t)ulength);
    166         if (buffer == NULL)
    167             return MP_MEMORY_ERROR;
    168     }
    169 
    170     Py_BEGIN_ALLOW_THREADS
    171     res = _conn_recvall(conn->handle, buffer, (size_t)ulength);
    172     Py_END_ALLOW_THREADS
    173 
    174     if (res >= 0) {
    175         res = (Py_ssize_t)ulength;
    176     } else if (*newbuffer != NULL) {
    177         PyMem_Free(*newbuffer);
    178         *newbuffer = NULL;
    179     }
    180     return res;
    181 }
    182 
    183 /*
    184  * Check whether any data is available for reading -- neg timeout blocks
    185  */
    186 
    187 static int
    188 conn_poll(ConnectionObject *conn, double timeout, PyThreadState *_save)
    189 {
    190 #if defined(HAVE_POLL) && !defined(HAVE_BROKEN_POLL)
    191     int res;
    192     struct pollfd p;
    193 
    194     p.fd = (int)conn->handle;
    195     p.events = POLLIN | POLLPRI;
    196     p.revents = 0;
    197 
    198     if (timeout < 0) {
    199         do {
    200             res = poll(&p, 1, -1);
    201         } while (res < 0 && errno == EINTR);
    202     } else {
    203         res = poll(&p, 1, (int)(timeout * 1000 + 0.5));
    204         if (res < 0 && errno == EINTR) {
    205             /* We were interrupted by a signal.  Just indicate a
    206                timeout even though we are early. */
    207             return FALSE;
    208         }
    209     }
    210 
    211     if (res < 0) {
    212         return MP_SOCKET_ERROR;
    213     } else if (p.revents & (POLLNVAL|POLLERR)) {
    214         Py_BLOCK_THREADS
    215         PyErr_SetString(PyExc_IOError, "poll() gave POLLNVAL or POLLERR");
    216         Py_UNBLOCK_THREADS
    217         return MP_EXCEPTION_HAS_BEEN_SET;
    218     } else if (p.revents != 0) {
    219         return TRUE;
    220     } else {
    221         assert(res == 0);
    222         return FALSE;
    223     }
    224 #else
    225     int res;
    226     fd_set rfds;
    227 
    228     /*
    229      * Verify the handle, issue 3321. Not required for windows.
    230      */
    231     #ifndef MS_WINDOWS
    232         if (((int)conn->handle) < 0 || ((int)conn->handle) >= FD_SETSIZE) {
    233             Py_BLOCK_THREADS
    234             PyErr_SetString(PyExc_IOError, "handle out of range in select()");
    235             Py_UNBLOCK_THREADS
    236             return MP_EXCEPTION_HAS_BEEN_SET;
    237         }
    238     #endif
    239 
    240     FD_ZERO(&rfds);
    241     FD_SET((SOCKET)conn->handle, &rfds);
    242 
    243     if (timeout < 0.0) {
    244         do {
    245             res = select((int)conn->handle+1, &rfds, NULL, NULL, NULL);
    246         } while (res < 0 && errno == EINTR);
    247     } else {
    248         struct timeval tv;
    249         tv.tv_sec = (long)timeout;
    250         tv.tv_usec = (long)((timeout - tv.tv_sec) * 1e6 + 0.5);
    251         res = select((int)conn->handle+1, &rfds, NULL, NULL, &tv);
    252         if (res < 0 && errno == EINTR) {
    253             /* We were interrupted by a signal.  Just indicate a
    254                timeout even though we are early. */
    255             return FALSE;
    256         }
    257     }
    258 
    259     if (res < 0) {
    260         return MP_SOCKET_ERROR;
    261     } else if (FD_ISSET(conn->handle, &rfds)) {
    262         return TRUE;
    263     } else {
    264         assert(res == 0);
    265         return FALSE;
    266     }
    267 #endif
    268 }
    269 
    270 /*
    271  * "connection.h" defines the Connection type using defs above
    272  */
    273 
    274 #define CONNECTION_NAME "Connection"
    275 #define CONNECTION_TYPE ConnectionType
    276 
    277 #include "connection.h"
    278