/Modules/_multiprocessing/connection.h

http://unladen-swallow.googlecode.com/ · C Header · 515 lines · 400 code · 83 blank · 32 comment · 82 complexity · 473bd75c8c226fb880f3a852ec114e25 MD5 · raw file

  1. /*
  2. * Definition of a `Connection` type.
  3. * Used by `socket_connection.c` and `pipe_connection.c`.
  4. *
  5. * connection.h
  6. *
  7. * Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
  8. */
  9. #ifndef CONNECTION_H
  10. #define CONNECTION_H
  11. /*
  12. * Read/write flags
  13. */
  14. #define READABLE 1
  15. #define WRITABLE 2
  16. #define CHECK_READABLE(self) \
  17. if (!(self->flags & READABLE)) { \
  18. PyErr_SetString(PyExc_IOError, "connection is write-only"); \
  19. return NULL; \
  20. }
  21. #define CHECK_WRITABLE(self) \
  22. if (!(self->flags & WRITABLE)) { \
  23. PyErr_SetString(PyExc_IOError, "connection is read-only"); \
  24. return NULL; \
  25. }
  26. /*
  27. * Allocation and deallocation
  28. */
  29. static PyObject *
  30. connection_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  31. {
  32. ConnectionObject *self;
  33. HANDLE handle;
  34. BOOL readable = TRUE, writable = TRUE;
  35. static char *kwlist[] = {"handle", "readable", "writable", NULL};
  36. if (!PyArg_ParseTupleAndKeywords(args, kwds, F_HANDLE "|ii", kwlist,
  37. &handle, &readable, &writable))
  38. return NULL;
  39. if (handle == INVALID_HANDLE_VALUE || (Py_ssize_t)handle < 0) {
  40. PyErr_Format(PyExc_IOError, "invalid handle %zd",
  41. (Py_ssize_t)handle);
  42. return NULL;
  43. }
  44. if (!readable && !writable) {
  45. PyErr_SetString(PyExc_ValueError,
  46. "either readable or writable must be true");
  47. return NULL;
  48. }
  49. self = PyObject_New(ConnectionObject, type);
  50. if (self == NULL)
  51. return NULL;
  52. self->weakreflist = NULL;
  53. self->handle = handle;
  54. self->flags = 0;
  55. if (readable)
  56. self->flags |= READABLE;
  57. if (writable)
  58. self->flags |= WRITABLE;
  59. assert(self->flags >= 1 && self->flags <= 3);
  60. return (PyObject*)self;
  61. }
  62. static void
  63. connection_dealloc(ConnectionObject* self)
  64. {
  65. if (self->weakreflist != NULL)
  66. PyObject_ClearWeakRefs((PyObject*)self);
  67. if (self->handle != INVALID_HANDLE_VALUE) {
  68. Py_BEGIN_ALLOW_THREADS
  69. CLOSE(self->handle);
  70. Py_END_ALLOW_THREADS
  71. }
  72. PyObject_Del(self);
  73. }
  74. /*
  75. * Functions for transferring buffers
  76. */
  77. static PyObject *
  78. connection_sendbytes(ConnectionObject *self, PyObject *args)
  79. {
  80. char *buffer;
  81. Py_ssize_t length, offset=0, size=PY_SSIZE_T_MIN;
  82. int res;
  83. if (!PyArg_ParseTuple(args, F_RBUFFER "#|" F_PY_SSIZE_T F_PY_SSIZE_T,
  84. &buffer, &length, &offset, &size))
  85. return NULL;
  86. CHECK_WRITABLE(self);
  87. if (offset < 0) {
  88. PyErr_SetString(PyExc_ValueError, "offset is negative");
  89. return NULL;
  90. }
  91. if (length < offset) {
  92. PyErr_SetString(PyExc_ValueError, "buffer length < offset");
  93. return NULL;
  94. }
  95. if (size == PY_SSIZE_T_MIN) {
  96. size = length - offset;
  97. } else {
  98. if (size < 0) {
  99. PyErr_SetString(PyExc_ValueError, "size is negative");
  100. return NULL;
  101. }
  102. if (offset + size > length) {
  103. PyErr_SetString(PyExc_ValueError,
  104. "buffer length < offset + size");
  105. return NULL;
  106. }
  107. }
  108. res = conn_send_string(self, buffer + offset, size);
  109. if (res < 0)
  110. return mp_SetError(PyExc_IOError, res);
  111. Py_RETURN_NONE;
  112. }
  113. static PyObject *
  114. connection_recvbytes(ConnectionObject *self, PyObject *args)
  115. {
  116. char *freeme = NULL;
  117. Py_ssize_t res, maxlength = PY_SSIZE_T_MAX;
  118. PyObject *result = NULL;
  119. if (!PyArg_ParseTuple(args, "|" F_PY_SSIZE_T, &maxlength))
  120. return NULL;
  121. CHECK_READABLE(self);
  122. if (maxlength < 0) {
  123. PyErr_SetString(PyExc_ValueError, "maxlength < 0");
  124. return NULL;
  125. }
  126. res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
  127. &freeme, maxlength);
  128. if (res < 0) {
  129. if (res == MP_BAD_MESSAGE_LENGTH) {
  130. if ((self->flags & WRITABLE) == 0) {
  131. Py_BEGIN_ALLOW_THREADS
  132. CLOSE(self->handle);
  133. Py_END_ALLOW_THREADS
  134. self->handle = INVALID_HANDLE_VALUE;
  135. } else {
  136. self->flags = WRITABLE;
  137. }
  138. }
  139. mp_SetError(PyExc_IOError, res);
  140. } else {
  141. if (freeme == NULL) {
  142. result = PyString_FromStringAndSize(self->buffer, res);
  143. } else {
  144. result = PyString_FromStringAndSize(freeme, res);
  145. PyMem_Free(freeme);
  146. }
  147. }
  148. return result;
  149. }
  150. static PyObject *
  151. connection_recvbytes_into(ConnectionObject *self, PyObject *args)
  152. {
  153. char *freeme = NULL, *buffer = NULL;
  154. Py_ssize_t res, length, offset = 0;
  155. PyObject *result = NULL;
  156. Py_buffer pbuf;
  157. CHECK_READABLE(self);
  158. if (!PyArg_ParseTuple(args, "w*|" F_PY_SSIZE_T,
  159. &pbuf, &offset))
  160. return NULL;
  161. buffer = pbuf.buf;
  162. length = pbuf.len;
  163. if (offset < 0) {
  164. PyErr_SetString(PyExc_ValueError, "negative offset");
  165. goto _error;
  166. }
  167. if (offset > length) {
  168. PyErr_SetString(PyExc_ValueError, "offset too large");
  169. goto _error;
  170. }
  171. res = conn_recv_string(self, buffer+offset, length-offset,
  172. &freeme, PY_SSIZE_T_MAX);
  173. if (res < 0) {
  174. if (res == MP_BAD_MESSAGE_LENGTH) {
  175. if ((self->flags & WRITABLE) == 0) {
  176. Py_BEGIN_ALLOW_THREADS
  177. CLOSE(self->handle);
  178. Py_END_ALLOW_THREADS
  179. self->handle = INVALID_HANDLE_VALUE;
  180. } else {
  181. self->flags = WRITABLE;
  182. }
  183. }
  184. mp_SetError(PyExc_IOError, res);
  185. } else {
  186. if (freeme == NULL) {
  187. result = PyInt_FromSsize_t(res);
  188. } else {
  189. result = PyObject_CallFunction(BufferTooShort,
  190. F_RBUFFER "#",
  191. freeme, res);
  192. PyMem_Free(freeme);
  193. if (result) {
  194. PyErr_SetObject(BufferTooShort, result);
  195. Py_DECREF(result);
  196. }
  197. goto _error;
  198. }
  199. }
  200. _cleanup:
  201. PyBuffer_Release(&pbuf);
  202. return result;
  203. _error:
  204. result = NULL;
  205. goto _cleanup;
  206. }
  207. /*
  208. * Functions for transferring objects
  209. */
  210. static PyObject *
  211. connection_send_obj(ConnectionObject *self, PyObject *obj)
  212. {
  213. char *buffer;
  214. int res;
  215. Py_ssize_t length;
  216. PyObject *pickled_string = NULL;
  217. CHECK_WRITABLE(self);
  218. pickled_string = PyObject_CallFunctionObjArgs(pickle_dumps, obj,
  219. pickle_protocol, NULL);
  220. if (!pickled_string)
  221. goto failure;
  222. if (PyString_AsStringAndSize(pickled_string, &buffer, &length) < 0)
  223. goto failure;
  224. res = conn_send_string(self, buffer, (int)length);
  225. if (res < 0) {
  226. mp_SetError(PyExc_IOError, res);
  227. goto failure;
  228. }
  229. Py_XDECREF(pickled_string);
  230. Py_RETURN_NONE;
  231. failure:
  232. Py_XDECREF(pickled_string);
  233. return NULL;
  234. }
  235. static PyObject *
  236. connection_recv_obj(ConnectionObject *self)
  237. {
  238. char *freeme = NULL;
  239. Py_ssize_t res;
  240. PyObject *temp = NULL, *result = NULL;
  241. CHECK_READABLE(self);
  242. res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
  243. &freeme, PY_SSIZE_T_MAX);
  244. if (res < 0) {
  245. if (res == MP_BAD_MESSAGE_LENGTH) {
  246. if ((self->flags & WRITABLE) == 0) {
  247. Py_BEGIN_ALLOW_THREADS
  248. CLOSE(self->handle);
  249. Py_END_ALLOW_THREADS
  250. self->handle = INVALID_HANDLE_VALUE;
  251. } else {
  252. self->flags = WRITABLE;
  253. }
  254. }
  255. mp_SetError(PyExc_IOError, res);
  256. } else {
  257. if (freeme == NULL) {
  258. temp = PyString_FromStringAndSize(self->buffer, res);
  259. } else {
  260. temp = PyString_FromStringAndSize(freeme, res);
  261. PyMem_Free(freeme);
  262. }
  263. }
  264. if (temp)
  265. result = PyObject_CallFunctionObjArgs(pickle_loads,
  266. temp, NULL);
  267. Py_XDECREF(temp);
  268. return result;
  269. }
  270. /*
  271. * Other functions
  272. */
  273. static PyObject *
  274. connection_poll(ConnectionObject *self, PyObject *args)
  275. {
  276. PyObject *timeout_obj = NULL;
  277. double timeout = 0.0;
  278. int res;
  279. CHECK_READABLE(self);
  280. if (!PyArg_ParseTuple(args, "|O", &timeout_obj))
  281. return NULL;
  282. if (timeout_obj == NULL) {
  283. timeout = 0.0;
  284. } else if (timeout_obj == Py_None) {
  285. timeout = -1.0; /* block forever */
  286. } else {
  287. timeout = PyFloat_AsDouble(timeout_obj);
  288. if (PyErr_Occurred())
  289. return NULL;
  290. if (timeout < 0.0)
  291. timeout = 0.0;
  292. }
  293. Py_BEGIN_ALLOW_THREADS
  294. res = conn_poll(self, timeout, _save);
  295. Py_END_ALLOW_THREADS
  296. switch (res) {
  297. case TRUE:
  298. Py_RETURN_TRUE;
  299. case FALSE:
  300. Py_RETURN_FALSE;
  301. default:
  302. return mp_SetError(PyExc_IOError, res);
  303. }
  304. }
  305. static PyObject *
  306. connection_fileno(ConnectionObject* self)
  307. {
  308. if (self->handle == INVALID_HANDLE_VALUE) {
  309. PyErr_SetString(PyExc_IOError, "handle is invalid");
  310. return NULL;
  311. }
  312. return PyInt_FromLong((long)self->handle);
  313. }
  314. static PyObject *
  315. connection_close(ConnectionObject *self)
  316. {
  317. if (self->handle != INVALID_HANDLE_VALUE) {
  318. Py_BEGIN_ALLOW_THREADS
  319. CLOSE(self->handle);
  320. Py_END_ALLOW_THREADS
  321. self->handle = INVALID_HANDLE_VALUE;
  322. }
  323. Py_RETURN_NONE;
  324. }
  325. static PyObject *
  326. connection_repr(ConnectionObject *self)
  327. {
  328. static char *conn_type[] = {"read-only", "write-only", "read-write"};
  329. assert(self->flags >= 1 && self->flags <= 3);
  330. return FROM_FORMAT("<%s %s, handle %zd>",
  331. conn_type[self->flags - 1],
  332. CONNECTION_NAME, (Py_ssize_t)self->handle);
  333. }
  334. /*
  335. * Getters and setters
  336. */
  337. static PyObject *
  338. connection_closed(ConnectionObject *self, void *closure)
  339. {
  340. return PyBool_FromLong((long)(self->handle == INVALID_HANDLE_VALUE));
  341. }
  342. static PyObject *
  343. connection_readable(ConnectionObject *self, void *closure)
  344. {
  345. return PyBool_FromLong((long)(self->flags & READABLE));
  346. }
  347. static PyObject *
  348. connection_writable(ConnectionObject *self, void *closure)
  349. {
  350. return PyBool_FromLong((long)(self->flags & WRITABLE));
  351. }
  352. /*
  353. * Tables
  354. */
  355. static PyMethodDef connection_methods[] = {
  356. {"send_bytes", (PyCFunction)connection_sendbytes, METH_VARARGS,
  357. "send the byte data from a readable buffer-like object"},
  358. {"recv_bytes", (PyCFunction)connection_recvbytes, METH_VARARGS,
  359. "receive byte data as a string"},
  360. {"recv_bytes_into",(PyCFunction)connection_recvbytes_into,METH_VARARGS,
  361. "receive byte data into a writeable buffer-like object\n"
  362. "returns the number of bytes read"},
  363. {"send", (PyCFunction)connection_send_obj, METH_O,
  364. "send a (picklable) object"},
  365. {"recv", (PyCFunction)connection_recv_obj, METH_NOARGS,
  366. "receive a (picklable) object"},
  367. {"poll", (PyCFunction)connection_poll, METH_VARARGS,
  368. "whether there is any input available to be read"},
  369. {"fileno", (PyCFunction)connection_fileno, METH_NOARGS,
  370. "file descriptor or handle of the connection"},
  371. {"close", (PyCFunction)connection_close, METH_NOARGS,
  372. "close the connection"},
  373. {NULL} /* Sentinel */
  374. };
  375. static PyGetSetDef connection_getset[] = {
  376. {"closed", (getter)connection_closed, NULL,
  377. "True if the connection is closed", NULL},
  378. {"readable", (getter)connection_readable, NULL,
  379. "True if the connection is readable", NULL},
  380. {"writable", (getter)connection_writable, NULL,
  381. "True if the connection is writable", NULL},
  382. {NULL}
  383. };
  384. /*
  385. * Connection type
  386. */
  387. PyDoc_STRVAR(connection_doc,
  388. "Connection type whose constructor signature is\n\n"
  389. " Connection(handle, readable=True, writable=True).\n\n"
  390. "The constructor does *not* duplicate the handle.");
  391. PyTypeObject CONNECTION_TYPE = {
  392. PyVarObject_HEAD_INIT(NULL, 0)
  393. /* tp_name */ "_multiprocessing." CONNECTION_NAME,
  394. /* tp_basicsize */ sizeof(ConnectionObject),
  395. /* tp_itemsize */ 0,
  396. /* tp_dealloc */ (destructor)connection_dealloc,
  397. /* tp_print */ 0,
  398. /* tp_getattr */ 0,
  399. /* tp_setattr */ 0,
  400. /* tp_compare */ 0,
  401. /* tp_repr */ (reprfunc)connection_repr,
  402. /* tp_as_number */ 0,
  403. /* tp_as_sequence */ 0,
  404. /* tp_as_mapping */ 0,
  405. /* tp_hash */ 0,
  406. /* tp_call */ 0,
  407. /* tp_str */ 0,
  408. /* tp_getattro */ 0,
  409. /* tp_setattro */ 0,
  410. /* tp_as_buffer */ 0,
  411. /* tp_flags */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
  412. Py_TPFLAGS_HAVE_WEAKREFS,
  413. /* tp_doc */ connection_doc,
  414. /* tp_traverse */ 0,
  415. /* tp_clear */ 0,
  416. /* tp_richcompare */ 0,
  417. /* tp_weaklistoffset */ offsetof(ConnectionObject, weakreflist),
  418. /* tp_iter */ 0,
  419. /* tp_iternext */ 0,
  420. /* tp_methods */ connection_methods,
  421. /* tp_members */ 0,
  422. /* tp_getset */ connection_getset,
  423. /* tp_base */ 0,
  424. /* tp_dict */ 0,
  425. /* tp_descr_get */ 0,
  426. /* tp_descr_set */ 0,
  427. /* tp_dictoffset */ 0,
  428. /* tp_init */ 0,
  429. /* tp_alloc */ 0,
  430. /* tp_new */ connection_new,
  431. };
  432. #endif /* CONNECTION_H */