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 }