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