/Modules/_multiprocessing/multiprocessing.c

http://unladen-swallow.googlecode.com/ · C · 311 lines · 224 code · 49 blank · 38 comment · 27 complexity · 98d5e4f8bff6ec15f349c9ae31dfca89 MD5 · raw file

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