/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
 11PyObject *create_win32_namespace(void);
 12
 13PyObject *pickle_dumps, *pickle_loads, *pickle_protocol;
 14PyObject *ProcessError, *BufferTooShort;
 15
 16/*
 17 * Function which raises exceptions based on error codes
 18 */
 19
 20PyObject *
 21mp_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
 74HANDLE sigint_event = NULL;
 75
 76static BOOL WINAPI
 77ProcessingCtrlHandler(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
 95static PyObject *
 96multiprocessing_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
130static PyObject *
131multiprocessing_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
175static PyObject*
176multiprocessing_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
193static 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
215PyMODINIT_FUNC
216init_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}