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
  36 static PyObject *
  37 connection_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 %"
  51                              PY_FORMAT_SIZE_T "d", (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
  78 static void
  79 connection_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
  96 static PyObject *
  97 connection_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         Py_BEGIN_ALLOW_THREADS
 133         res = conn_send_string(self, buffer + offset, size);
 134         Py_END_ALLOW_THREADS
 135
 136         if (res < 0)
 137                 return mp_SetError(PyExc_IOError, res);
 138
 139         Py_RETURN_NONE;
 140 }
 141
 142 static PyObject *
 143 connection_recvbytes(ConnectionObject *self, PyObject *args)
 144 {
 145         char *freeme = NULL;
 146         Py_ssize_t res, maxlength = PY_SSIZE_T_MAX;
 147         PyObject *result = NULL;
 148
 149         if (!PyArg_ParseTuple(args, "|" F_PY_SSIZE_T, &maxlength))
 150                 return NULL;
 151
 152         CHECK_READABLE(self);
 153
 154         if (maxlength < 0) {
 155                 PyErr_SetString(PyExc_ValueError, "maxlength < 0");
 156                 return NULL;
 157         }
 158
 159         Py_BEGIN_ALLOW_THREADS
 160         res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
 161                                &freeme, maxlength);
 162         Py_END_ALLOW_THREADS
 163
 164         if (res < 0) {
 165                 if (res == MP_BAD_MESSAGE_LENGTH) {
 166                         if ((self->flags & WRITABLE) == 0) {
 167                                 Py_BEGIN_ALLOW_THREADS
 168                                 CLOSE(self->handle);
 169                                 Py_END_ALLOW_THREADS
 170                                 self->handle = INVALID_HANDLE_VALUE;
 171                         } else {
 172                                 self->flags = WRITABLE;
 173                         }
 174                 }
 175                 mp_SetError(PyExc_IOError, res);
 176         } else {
 177                 if (freeme == NULL) {
 178                         result = PyString_FromStringAndSize(self->buffer, res);
 179                 } else {
 180                         result = PyString_FromStringAndSize(freeme, res);
 181                         PyMem_Free(freeme);
 182                 }
 183         }
 184
 185         return result;
 186 }
 187
 188 static PyObject *
 189 connection_recvbytes_into(ConnectionObject *self, PyObject *args)
 190 {
 191         char *freeme = NULL, *buffer = NULL;
 192         Py_ssize_t res, length, offset = 0;
 193         PyObject *result = NULL;
 194
 195         if (!PyArg_ParseTuple(args, "w#|" F_PY_SSIZE_T,
 196                               &buffer, &length, &offset))
 197                 return NULL;
 198
 199         CHECK_READABLE(self);
 200
 201         if (offset < 0) {
 202                 PyErr_SetString(PyExc_ValueError, "negative offset");
 203                 return NULL;
 204         }
 205
 206         if (offset > length) {
 207                 PyErr_SetString(PyExc_ValueError, "offset too large");
 208                 return NULL;
 209         }
 210
 211         Py_BEGIN_ALLOW_THREADS
 212         res = conn_recv_string(self, buffer+offset, length-offset,
 213                                &freeme, PY_SSIZE_T_MAX);
 214         Py_END_ALLOW_THREADS
 215
 216         if (res < 0) {
 217                 if (res == MP_BAD_MESSAGE_LENGTH) {
 218                         if ((self->flags & WRITABLE) == 0) {
 219                                 Py_BEGIN_ALLOW_THREADS
 220                                 CLOSE(self->handle);
 221                                 Py_END_ALLOW_THREADS
 222                                 self->handle = INVALID_HANDLE_VALUE;
 223                         } else {
 224                                 self->flags = WRITABLE;
 225                         }
 226                 }
 227                 mp_SetError(PyExc_IOError, res);
 228         } else {
 229                 if (freeme == NULL) {
 230                         result = PyInt_FromSsize_t(res);
 231                 } else {
 232                         result = PyObject_CallFunction(BufferTooShort,
 233                                                        F_RBUFFER "#",
 234                                                        freeme, res);
 235                         PyMem_Free(freeme);
 236                         if (result) {
 237                                 PyErr_SetObject(BufferTooShort, result);
 238                                 Py_DECREF(result);
 239                         }
 240                         return NULL;
 241                 }
 242         }
 243
 244         return result;
 245 }
 246
 247 /*
 248  * Functions for transferring objects
 249  */
 250
 251 static PyObject *
 252 connection_send_obj(ConnectionObject *self, PyObject *obj)
 253 {
 254         char *buffer;
 255         int res;
 256         Py_ssize_t length;
 257         PyObject *pickled_string = NULL;
 258
 259         CHECK_WRITABLE(self);
 260
 261         pickled_string = PyObject_CallFunctionObjArgs(pickle_dumps, obj,
 262                                                       pickle_protocol, NULL);
 263         if (!pickled_string)
 264                 goto failure;
 265
 266         if (PyString_AsStringAndSize(pickled_string, &buffer, &length) < 0)
 267                 goto failure;
 268
 269         Py_BEGIN_ALLOW_THREADS
 270         res = conn_send_string(self, buffer, (int)length);
 271         Py_END_ALLOW_THREADS
 272
 273         if (res < 0) {
 274                 mp_SetError(PyExc_IOError, res);
 275                 goto failure;
 276         }
 277
 278         Py_XDECREF(pickled_string);
 279         Py_RETURN_NONE;
 280
 281   failure:
 282         Py_XDECREF(pickled_string);
 283         return NULL;
 284 }
 285
 286 static PyObject *
 287 connection_recv_obj(ConnectionObject *self)
 288 {
 289         char *freeme = NULL;
 290         Py_ssize_t res;
 291         PyObject *temp = NULL, *result = NULL;
 292
 293         CHECK_READABLE(self);
 294
 295         Py_BEGIN_ALLOW_THREADS
 296         res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
 297                                &freeme, PY_SSIZE_T_MAX);
 298         Py_END_ALLOW_THREADS
 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
 332 static PyObject *
 333 connection_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);
 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
 370 static PyObject *
 371 connection_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
 380 static PyObject *
 381 connection_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
 393 static PyObject *
 394 connection_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 %" PY_FORMAT_SIZE_T "d>",
 400                            conn_type[self->flags - 1],
 401                            CONNECTION_NAME, (Py_ssize_t)self->handle);
 402 }
 403
 404 /*
 405  * Getters and setters
 406  */
 407
 408 static PyObject *
 409 connection_closed(ConnectionObject *self, void *closure)
 410 {
 411         return PyBool_FromLong((long)(self->handle == INVALID_HANDLE_VALUE));
 412 }
 413
 414 static PyObject *
 415 connection_readable(ConnectionObject *self, void *closure)
 416 {
 417         return PyBool_FromLong((long)(self->flags & READABLE));
 418 }
 419
 420 static PyObject *
 421 connection_writable(ConnectionObject *self, void *closure)
 422 {
 423         return PyBool_FromLong((long)(self->flags & WRITABLE));
 424 }
 425
 426 /*
 427  * Tables
 428  */
 429
 430 static 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
 454 static 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
 468 PyDoc_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
 473 PyTypeObject 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 */