Modules/_multiprocessing/multiprocessing.c

   1 /*
   2  * Extension module used by multiprocessing package
   3  *
   4  * multiprocessing.c
   5  *
   6  * Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
   7  */
   8
   9 #include "multiprocessing.h"
  10
  11 #ifdef SCM_RIGHTS
  12         #define HAVE_FD_TRANSFER 1
  13 #else
  14         #define HAVE_FD_TRANSFER 0
  15 #endif
  16
  17 PyObject *create_win32_namespace(void);
  18
  19 PyObject *pickle_dumps, *pickle_loads, *pickle_protocol;
  20 PyObject *ProcessError, *BufferTooShort;
  21
  22 /*
  23  * Function which raises exceptions based on error codes
  24  */
  25
  26 PyObject *
  27 mp_SetError(PyObject *Type, int num)
  28 {
  29         switch (num) {
  30 #ifdef MS_WINDOWS
  31         case MP_STANDARD_ERROR:
  32                 if (Type == NULL)
  33                         Type = PyExc_WindowsError;
  34                 PyErr_SetExcFromWindowsErr(Type, 0);
  35                 break;
  36         case MP_SOCKET_ERROR:
  37                 if (Type == NULL)
  38                         Type = PyExc_WindowsError;
  39                 PyErr_SetExcFromWindowsErr(Type, WSAGetLastError());
  40                 break;
  41 #else /* !MS_WINDOWS */
  42         case MP_STANDARD_ERROR:
  43         case MP_SOCKET_ERROR:
  44                 if (Type == NULL)
  45                         Type = PyExc_OSError;
  46                 PyErr_SetFromErrno(Type);
  47                 break;
  48 #endif /* !MS_WINDOWS */
  49         case MP_MEMORY_ERROR:
  50                 PyErr_NoMemory();
  51                 break;
  52         case MP_END_OF_FILE:
  53                 PyErr_SetNone(PyExc_EOFError);
  54                 break;
  55         case MP_EARLY_END_OF_FILE:
  56                 PyErr_SetString(PyExc_IOError,
  57                                 "got end of file during message");
  58                 break;
  59         case MP_BAD_MESSAGE_LENGTH:
  60                 PyErr_SetString(PyExc_IOError, "bad message length");
  61                 break;
  62         case MP_EXCEPTION_HAS_BEEN_SET:
  63                 break;
  64         default:
  65                 PyErr_Format(PyExc_RuntimeError,
  66                              "unkown error number %d", num);
  67         }
  68         return NULL;
  69 }
  70
  71
  72 /*
  73  * Windows only
  74  */
  75
  76 #ifdef MS_WINDOWS
  77
  78 /* On Windows we set an event to signal Ctrl-C; compare with timemodule.c */
  79
  80 HANDLE sigint_event = NULL;
  81
  82 static BOOL WINAPI
  83 ProcessingCtrlHandler(DWORD dwCtrlType)
  84 {
  85         SetEvent(sigint_event);
  86         return FALSE;
  87 }
  88
  89 /*
  90  * Unix only
  91  */
  92
  93 #else /* !MS_WINDOWS */
  94
  95 #if HAVE_FD_TRANSFER
  96
  97 /* Functions for transferring file descriptors between processes.
  98    Reimplements some of the functionality of the fdcred
  99    module at http://www.mca-ltd.com/resources/fdcred_1.tgz. */
 100
 101 static PyObject *
 102 multiprocessing_sendfd(PyObject *self, PyObject *args)
 103 {
 104         int conn, fd, res;
 105         char dummy_char;
 106         char buf[CMSG_SPACE(sizeof(int))];
 107         struct msghdr msg = {0};
 108         struct iovec dummy_iov;
 109         struct cmsghdr *cmsg;
 110
 111         if (!PyArg_ParseTuple(args, "ii", &conn, &fd))
 112                 return NULL;
 113
 114         dummy_iov.iov_base = &dummy_char;
 115         dummy_iov.iov_len = 1;
 116         msg.msg_control = buf;
 117         msg.msg_controllen = sizeof(buf);
 118         msg.msg_iov = &dummy_iov;
 119         msg.msg_iovlen = 1;
 120         cmsg = CMSG_FIRSTHDR(&msg);
 121         cmsg->cmsg_level = SOL_SOCKET;
 122         cmsg->cmsg_type = SCM_RIGHTS;
 123         cmsg->cmsg_len = CMSG_LEN(sizeof(int));
 124         msg.msg_controllen = cmsg->cmsg_len;
 125         *(int*)CMSG_DATA(cmsg) = fd;
 126
 127         Py_BEGIN_ALLOW_THREADS
 128         res = sendmsg(conn, &msg, 0);
 129         Py_END_ALLOW_THREADS
 130
 131         if (res < 0)
 132                 return PyErr_SetFromErrno(PyExc_OSError);
 133         Py_RETURN_NONE;
 134 }
 135
 136 static PyObject *
 137 multiprocessing_recvfd(PyObject *self, PyObject *args)
 138 {
 139         int conn, fd, res;
 140         char dummy_char;
 141         char buf[CMSG_SPACE(sizeof(int))];
 142         struct msghdr msg = {0};
 143         struct iovec dummy_iov;
 144         struct cmsghdr *cmsg;
 145
 146         if (!PyArg_ParseTuple(args, "i", &conn))
 147                 return NULL;
 148
 149         dummy_iov.iov_base = &dummy_char;
 150         dummy_iov.iov_len = 1;
 151         msg.msg_control = buf;
 152         msg.msg_controllen = sizeof(buf);
 153         msg.msg_iov = &dummy_iov;
 154         msg.msg_iovlen = 1;
 155         cmsg = CMSG_FIRSTHDR(&msg);
 156         cmsg->cmsg_level = SOL_SOCKET;
 157         cmsg->cmsg_type = SCM_RIGHTS;
 158         cmsg->cmsg_len = CMSG_LEN(sizeof(int));
 159         msg.msg_controllen = cmsg->cmsg_len;
 160
 161         Py_BEGIN_ALLOW_THREADS
 162         res = recvmsg(conn, &msg, 0);
 163         Py_END_ALLOW_THREADS
 164
 165         if (res < 0)
 166                 return PyErr_SetFromErrno(PyExc_OSError);
 167
 168         fd = *(int*)CMSG_DATA(cmsg);
 169         return Py_BuildValue("i", fd);
 170 }
 171
 172 #endif /* HAVE_FD_TRANSFER */
 173
 174 #endif /* !MS_WINDOWS */
 175
 176
 177 /*
 178  * All platforms
 179  */
 180
 181 static PyObject*
 182 multiprocessing_address_of_buffer(PyObject *self, PyObject *obj)
 183 {
 184         void *buffer;
 185         Py_ssize_t buffer_len;
 186
 187         if (PyObject_AsWriteBuffer(obj, &buffer, &buffer_len) < 0)
 188                 return NULL;
 189
 190         return Py_BuildValue("N" F_PY_SSIZE_T,
 191                              PyLong_FromVoidPtr(buffer), buffer_len);
 192 }
 193
 194
 195 /*
 196  * Function table
 197  */
 198
 199 static PyMethodDef module_methods[] = {
 200         {"address_of_buffer", multiprocessing_address_of_buffer, METH_O,
 201          "address_of_buffer(obj) -> int\n"
 202          "Return address of obj assuming obj supports buffer inteface"},
 203 #if HAVE_FD_TRANSFER
 204         {"sendfd", multiprocessing_sendfd, METH_VARARGS,
 205          "sendfd(sockfd, fd) -> None\n"
 206          "Send file descriptor given by fd over the unix domain socket\n"
 207          "whose file decriptor is sockfd"},
 208         {"recvfd", multiprocessing_recvfd, METH_VARARGS,
 209          "recvfd(sockfd) -> fd\n"
 210          "Receive a file descriptor over a unix domain socket\n"
 211          "whose file decriptor is sockfd"},
 212 #endif
 213         {NULL}
 214 };
 215
 216
 217 /*
 218  * Initialize
 219  */
 220
 221 PyMODINIT_FUNC
 222 init_multiprocessing(void)
 223 {
 224         PyObject *module, *temp, *value;
 225
 226         /* Initialize module */
 227         module = Py_InitModule("_multiprocessing", module_methods);
 228         if (!module)
 229                 return;
 230
 231         /* Get copy of objects from pickle */
 232         temp = PyImport_ImportModule(PICKLE_MODULE);
 233         if (!temp)
 234                 return;
 235         pickle_dumps = PyObject_GetAttrString(temp, "dumps");
 236         pickle_loads = PyObject_GetAttrString(temp, "loads");
 237         pickle_protocol = PyObject_GetAttrString(temp, "HIGHEST_PROTOCOL");
 238         Py_XDECREF(temp);
 239
 240         /* Get copy of BufferTooShort */
 241         temp = PyImport_ImportModule("multiprocessing");
 242         if (!temp)
 243                 return;
 244         BufferTooShort = PyObject_GetAttrString(temp, "BufferTooShort");
 245         Py_XDECREF(temp);
 246
 247         /* Add connection type to module */
 248         if (PyType_Ready(&ConnectionType) < 0)
 249                 return;
 250         Py_INCREF(&ConnectionType);
 251         PyModule_AddObject(module, "Connection", (PyObject*)&ConnectionType);
 252
 253 #if defined(MS_WINDOWS) ||                                              \
 254   (defined(HAVE_SEM_OPEN) && !defined(HAVE_BROKEN_POSIX_SEMAPHORES))
 255         /* Add SemLock type to module */
 256         if (PyType_Ready(&SemLockType) < 0)
 257                 return;
 258         Py_INCREF(&SemLockType);
 259         PyDict_SetItemString(SemLockType.tp_dict, "SEM_VALUE_MAX",
 260                              Py_BuildValue("i", SEM_VALUE_MAX));
 261         PyModule_AddObject(module, "SemLock", (PyObject*)&SemLockType);
 262 #endif
 263
 264 #ifdef MS_WINDOWS
 265         /* Add PipeConnection to module */
 266         if (PyType_Ready(&PipeConnectionType) < 0)
 267                 return;
 268         Py_INCREF(&PipeConnectionType);
 269         PyModule_AddObject(module, "PipeConnection",
 270                            (PyObject*)&PipeConnectionType);
 271
 272         /* Initialize win32 class and add to multiprocessing */
 273         temp = create_win32_namespace();
 274         if (!temp)
 275                 return;
 276         PyModule_AddObject(module, "win32", temp);
 277
 278         /* Initialize the event handle used to signal Ctrl-C */
 279         sigint_event = CreateEvent(NULL, TRUE, FALSE, NULL);
 280         if (!sigint_event) {
 281                 PyErr_SetFromWindowsErr(0);
 282                 return;
 283         }
 284         if (!SetConsoleCtrlHandler(ProcessingCtrlHandler, TRUE)) {
 285                 PyErr_SetFromWindowsErr(0);
 286                 return;
 287         }
 288 #endif
 289
 290         /* Add configuration macros */
 291         temp = PyDict_New();
 292         if (!temp)
 293                 return;
 294 #define ADD_FLAG(name)                                            \
 295         value = Py_BuildValue("i", name);                         \
 296         if (value == NULL) { Py_DECREF(temp); return; }           \
 297         if (PyDict_SetItemString(temp, #name, value) < 0) {       \
 298                 Py_DECREF(temp); Py_DECREF(value); return; }      \
 299         Py_DECREF(value)
 300
 301 #if defined(HAVE_SEM_OPEN) && !defined(HAVE_BROKEN_POSIX_SEMAPHORES)
 302         ADD_FLAG(HAVE_SEM_OPEN);
 303 #endif
 304 #ifdef HAVE_SEM_TIMEDWAIT
 305         ADD_FLAG(HAVE_SEM_TIMEDWAIT);
 306 #endif
 307 #ifdef HAVE_FD_TRANSFER
 308         ADD_FLAG(HAVE_FD_TRANSFER);
 309 #endif
 310 #ifdef HAVE_BROKEN_SEM_GETVALUE
 311         ADD_FLAG(HAVE_BROKEN_SEM_GETVALUE);
 312 #endif
 313 #ifdef HAVE_BROKEN_SEM_UNLINK
 314         ADD_FLAG(HAVE_BROKEN_SEM_UNLINK);
 315 #endif
 316         if (PyModule_AddObject(module, "flags", temp) < 0)
 317                 return;
 318 }