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 #if (defined(CMSG_LEN) && defined(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 *
mp_SetError(PyObject * Type,int num)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 "unknown 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
ProcessingCtrlHandler(DWORD dwCtrlType)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 /* Based in http://resin.csoft.net/cgi-bin/man.cgi?section=3&topic=CMSG_DATA */
101
102 static PyObject *
multiprocessing_sendfd(PyObject * self,PyObject * args)103 multiprocessing_sendfd(PyObject *self, PyObject *args)
104 {
105 int conn, fd, res;
106 struct iovec dummy_iov;
107 char dummy_char;
108 struct msghdr msg;
109 struct cmsghdr *cmsg;
110 union {
111 struct cmsghdr hdr;
112 unsigned char buf[CMSG_SPACE(sizeof(int))];
113 } cmsgbuf;
114
115 if (!PyArg_ParseTuple(args, "ii", &conn, &fd))
116 return NULL;
117
118 dummy_iov.iov_base = &dummy_char;
119 dummy_iov.iov_len = 1;
120
121 memset(&msg, 0, sizeof(msg));
122 msg.msg_control = &cmsgbuf.buf;
123 msg.msg_controllen = sizeof(cmsgbuf.buf);
124 msg.msg_iov = &dummy_iov;
125 msg.msg_iovlen = 1;
126
127 cmsg = CMSG_FIRSTHDR(&msg);
128 cmsg->cmsg_len = CMSG_LEN(sizeof(int));
129 cmsg->cmsg_level = SOL_SOCKET;
130 cmsg->cmsg_type = SCM_RIGHTS;
131 * (int *) CMSG_DATA(cmsg) = fd;
132
133 Py_BEGIN_ALLOW_THREADS
134 res = sendmsg(conn, &msg, 0);
135 Py_END_ALLOW_THREADS
136
137 if (res < 0)
138 return PyErr_SetFromErrno(PyExc_OSError);
139 Py_RETURN_NONE;
140 }
141
142 static PyObject *
multiprocessing_recvfd(PyObject * self,PyObject * args)143 multiprocessing_recvfd(PyObject *self, PyObject *args)
144 {
145 int conn, fd, res;
146 char dummy_char;
147 struct iovec dummy_iov;
148 struct msghdr msg = {0};
149 struct cmsghdr *cmsg;
150 union {
151 struct cmsghdr hdr;
152 unsigned char buf[CMSG_SPACE(sizeof(int))];
153 } cmsgbuf;
154
155 if (!PyArg_ParseTuple(args, "i", &conn))
156 return NULL;
157
158 dummy_iov.iov_base = &dummy_char;
159 dummy_iov.iov_len = 1;
160
161 memset(&msg, 0, sizeof(msg));
162 msg.msg_control = &cmsgbuf.buf;
163 msg.msg_controllen = sizeof(cmsgbuf.buf);
164 msg.msg_iov = &dummy_iov;
165 msg.msg_iovlen = 1;
166
167 cmsg = CMSG_FIRSTHDR(&msg);
168 cmsg->cmsg_level = SOL_SOCKET;
169 cmsg->cmsg_type = SCM_RIGHTS;
170 cmsg->cmsg_len = CMSG_LEN(sizeof(int));
171 msg.msg_controllen = cmsg->cmsg_len;
172
173 Py_BEGIN_ALLOW_THREADS
174 res = recvmsg(conn, &msg, 0);
175 Py_END_ALLOW_THREADS
176
177 if (res < 0)
178 return PyErr_SetFromErrno(PyExc_OSError);
179
180 if (msg.msg_controllen < CMSG_LEN(sizeof(int)) ||
181 (cmsg = CMSG_FIRSTHDR(&msg)) == NULL ||
182 cmsg->cmsg_level != SOL_SOCKET ||
183 cmsg->cmsg_type != SCM_RIGHTS ||
184 cmsg->cmsg_len < CMSG_LEN(sizeof(int))) {
185 /* If at least one control message is present, there should be
186 no room for any further data in the buffer. */
187 PyErr_SetString(PyExc_RuntimeError, "No file descriptor received");
188 return NULL;
189 }
190
191 fd = * (int *) CMSG_DATA(cmsg);
192 return Py_BuildValue("i", fd);
193 }
194
195 #endif /* HAVE_FD_TRANSFER */
196
197 #endif /* !MS_WINDOWS */
198
199
200 /*
201 * All platforms
202 */
203
204 static PyObject*
multiprocessing_address_of_buffer(PyObject * self,PyObject * obj)205 multiprocessing_address_of_buffer(PyObject *self, PyObject *obj)
206 {
207 void *buffer;
208 Py_ssize_t buffer_len;
209
210 if (PyObject_AsWriteBuffer(obj, &buffer, &buffer_len) < 0)
211 return NULL;
212
213 return Py_BuildValue("N" F_PY_SSIZE_T,
214 PyLong_FromVoidPtr(buffer), buffer_len);
215 }
216
217
218 /*
219 * Function table
220 */
221
222 static PyMethodDef module_methods[] = {
223 {"address_of_buffer", multiprocessing_address_of_buffer, METH_O,
224 "address_of_buffer(obj) -> int\n"
225 "Return address of obj assuming obj supports buffer inteface"},
226 #if HAVE_FD_TRANSFER
227 {"sendfd", multiprocessing_sendfd, METH_VARARGS,
228 "sendfd(sockfd, fd) -> None\n"
229 "Send file descriptor given by fd over the unix domain socket\n"
230 "whose file decriptor is sockfd"},
231 {"recvfd", multiprocessing_recvfd, METH_VARARGS,
232 "recvfd(sockfd) -> fd\n"
233 "Receive a file descriptor over a unix domain socket\n"
234 "whose file decriptor is sockfd"},
235 #endif
236 {NULL}
237 };
238
239
240 /*
241 * Initialize
242 */
243
244 PyMODINIT_FUNC
init_multiprocessing(void)245 init_multiprocessing(void)
246 {
247 PyObject *module, *temp, *value;
248
249 /* Initialize module */
250 module = Py_InitModule("_multiprocessing", module_methods);
251 if (!module)
252 return;
253
254 /* Get copy of objects from pickle */
255 temp = PyImport_ImportModule(PICKLE_MODULE);
256 if (!temp)
257 return;
258 pickle_dumps = PyObject_GetAttrString(temp, "dumps");
259 pickle_loads = PyObject_GetAttrString(temp, "loads");
260 pickle_protocol = PyObject_GetAttrString(temp, "HIGHEST_PROTOCOL");
261 Py_XDECREF(temp);
262
263 /* Get copy of BufferTooShort */
264 temp = PyImport_ImportModule("multiprocessing");
265 if (!temp)
266 return;
267 BufferTooShort = PyObject_GetAttrString(temp, "BufferTooShort");
268 Py_XDECREF(temp);
269
270 /* Add connection type to module */
271 if (PyType_Ready(&ConnectionType) < 0)
272 return;
273 Py_INCREF(&ConnectionType);
274 PyModule_AddObject(module, "Connection", (PyObject*)&ConnectionType);
275
276 #if defined(MS_WINDOWS) || \
277 (defined(HAVE_SEM_OPEN) && !defined(POSIX_SEMAPHORES_NOT_ENABLED))
278 /* Add SemLock type to module */
279 if (PyType_Ready(&SemLockType) < 0)
280 return;
281 Py_INCREF(&SemLockType);
282 {
283 PyObject *py_sem_value_max;
284 /* Some systems define SEM_VALUE_MAX as an unsigned value that
285 * causes it to be negative when used as an int (NetBSD). */
286 if ((int)(SEM_VALUE_MAX) < 0)
287 py_sem_value_max = PyLong_FromLong(INT_MAX);
288 else
289 py_sem_value_max = PyLong_FromLong(SEM_VALUE_MAX);
290 if (py_sem_value_max == NULL)
291 return;
292 PyDict_SetItemString(SemLockType.tp_dict, "SEM_VALUE_MAX",
293 py_sem_value_max);
294 }
295 PyModule_AddObject(module, "SemLock", (PyObject*)&SemLockType);
296 #endif
297
298 #ifdef MS_WINDOWS
299 /* Add PipeConnection to module */
300 if (PyType_Ready(&PipeConnectionType) < 0)
301 return;
302 Py_INCREF(&PipeConnectionType);
303 PyModule_AddObject(module, "PipeConnection",
304 (PyObject*)&PipeConnectionType);
305
306 /* Initialize win32 class and add to multiprocessing */
307 temp = create_win32_namespace();
308 if (!temp)
309 return;
310 PyModule_AddObject(module, "win32", temp);
311
312 /* Initialize the event handle used to signal Ctrl-C */
313 sigint_event = CreateEvent(NULL, TRUE, FALSE, NULL);
314 if (!sigint_event) {
315 PyErr_SetFromWindowsErr(0);
316 return;
317 }
318 if (!SetConsoleCtrlHandler(ProcessingCtrlHandler, TRUE)) {
319 PyErr_SetFromWindowsErr(0);
320 return;
321 }
322 #endif
323
324 /* Add configuration macros */
325 temp = PyDict_New();
326 if (!temp)
327 return;
328 #define ADD_FLAG(name) \
329 value = Py_BuildValue("i", name); \
330 if (value == NULL) { Py_DECREF(temp); return; } \
331 if (PyDict_SetItemString(temp, #name, value) < 0) { \
332 Py_DECREF(temp); Py_DECREF(value); return; } \
333 Py_DECREF(value)
334
335 #if defined(HAVE_SEM_OPEN) && !defined(POSIX_SEMAPHORES_NOT_ENABLED)
336 ADD_FLAG(HAVE_SEM_OPEN);
337 #endif
338 #ifdef HAVE_SEM_TIMEDWAIT
339 ADD_FLAG(HAVE_SEM_TIMEDWAIT);
340 #endif
341 #ifdef HAVE_FD_TRANSFER
342 ADD_FLAG(HAVE_FD_TRANSFER);
343 #endif
344 #ifdef HAVE_BROKEN_SEM_GETVALUE
345 ADD_FLAG(HAVE_BROKEN_SEM_GETVALUE);
346 #endif
347 #ifdef HAVE_BROKEN_SEM_UNLINK
348 ADD_FLAG(HAVE_BROKEN_SEM_UNLINK);
349 #endif
350 if (PyModule_AddObject(module, "flags", temp) < 0)
351 return;
352 }
353