/Modules/_multiprocessing/connection.h

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