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 %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
  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         res = conn_send_string(self, buffer + offset, size);
 133
 134         if (res < 0) {
 135                 if (PyErr_Occurred())
 136                         return NULL;
 137                 else
 138                         return mp_SetError(PyExc_IOError, res);
 139         }
 140
 141         Py_RETURN_NONE;
 142 }
 143
 144 static PyObject *
 145 connection_recvbytes(ConnectionObject *self, PyObject *args)
 146 {
 147         char *freeme = NULL;
 148         Py_ssize_t res, maxlength = PY_SSIZE_T_MAX;
 149         PyObject *result = NULL;
 150
 151         if (!PyArg_ParseTuple(args, "|" F_PY_SSIZE_T, &maxlength))
 152                 return NULL;
 153
 154         CHECK_READABLE(self);
 155
 156         if (maxlength < 0) {
 157                 PyErr_SetString(PyExc_ValueError, "maxlength < 0");
 158                 return NULL;
 159         }
 160
 161         res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
 162                                &freeme, maxlength);
 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         Py_buffer pbuf;
 195
 196         CHECK_READABLE(self);
 197
 198         if (!PyArg_ParseTuple(args, "w*|" F_PY_SSIZE_T,
 199                               &pbuf, &offset))
 200                 return NULL;
 201
 202         buffer = pbuf.buf;
 203         length = pbuf.len;
 204
 205         if (offset < 0) {
 206                 PyErr_SetString(PyExc_ValueError, "negative offset");
 207                 goto _error;
 208         }
 209
 210         if (offset > length) {
 211                 PyErr_SetString(PyExc_ValueError, "offset too large");
 212                 goto _error;
 213         }
 214
 215         res = conn_recv_string(self, buffer+offset, length-offset,
 216                                &freeme, PY_SSIZE_T_MAX);
 217
 218         if (res < 0) {
 219                 if (res == MP_BAD_MESSAGE_LENGTH) {
 220                         if ((self->flags & WRITABLE) == 0) {
 221                                 Py_BEGIN_ALLOW_THREADS
 222                                 CLOSE(self->handle);
 223                                 Py_END_ALLOW_THREADS
 224                                 self->handle = INVALID_HANDLE_VALUE;
 225                         } else {
 226                                 self->flags = WRITABLE;
 227                         }
 228                 }
 229                 mp_SetError(PyExc_IOError, res);
 230         } else {
 231                 if (freeme == NULL) {
 232                         result = PyInt_FromSsize_t(res);
 233                 } else {
 234                         result = PyObject_CallFunction(BufferTooShort,
 235                                                        F_RBUFFER "#",
 236                                                        freeme, res);
 237                         PyMem_Free(freeme);
 238                         if (result) {
 239                                 PyErr_SetObject(BufferTooShort, result);
 240                                 Py_DECREF(result);
 241                         }
 242                         goto _error;
 243                 }
 244         }
 245
 246 _cleanup:
 247         PyBuffer_Release(&pbuf);
 248         return result;
 249
 250 _error:
 251         result = NULL;
 252         goto _cleanup;
 253 }
 254
 255 /*
 256  * Functions for transferring objects
 257  */
 258
 259 static PyObject *
 260 connection_send_obj(ConnectionObject *self, PyObject *obj)
 261 {
 262         char *buffer;
 263         int res;
 264         Py_ssize_t length;
 265         PyObject *pickled_string = NULL;
 266
 267         CHECK_WRITABLE(self);
 268
 269         pickled_string = PyObject_CallFunctionObjArgs(pickle_dumps, obj,
 270                                                       pickle_protocol, NULL);
 271         if (!pickled_string)
 272                 goto failure;
 273
 274         if (PyString_AsStringAndSize(pickled_string, &buffer, &length) < 0)
 275                 goto failure;
 276
 277         res = conn_send_string(self, buffer, (int)length);
 278
 279         if (res < 0) {
 280                 mp_SetError(PyExc_IOError, res);
 281                 goto failure;
 282         }
 283
 284         Py_XDECREF(pickled_string);
 285         Py_RETURN_NONE;
 286
 287   failure:
 288         Py_XDECREF(pickled_string);
 289         return NULL;
 290 }
 291
 292 static PyObject *
 293 connection_recv_obj(ConnectionObject *self)
 294 {
 295         char *freeme = NULL;
 296         Py_ssize_t res;
 297         PyObject *temp = NULL, *result = NULL;
 298
 299         CHECK_READABLE(self);
 300
 301         res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
 302                                &freeme, PY_SSIZE_T_MAX);
 303
 304         if (res < 0) {
 305                 if (res == MP_BAD_MESSAGE_LENGTH) {
 306                         if ((self->flags & WRITABLE) == 0) {
 307                                 Py_BEGIN_ALLOW_THREADS
 308                                 CLOSE(self->handle);
 309                                 Py_END_ALLOW_THREADS
 310                                 self->handle = INVALID_HANDLE_VALUE;
 311                         } else {
 312                                 self->flags = WRITABLE;
 313                         }
 314                 }
 315                 mp_SetError(PyExc_IOError, res);
 316         } else {
 317                 if (freeme == NULL) {
 318                         temp = PyString_FromStringAndSize(self->buffer, res);
 319                 } else {
 320                         temp = PyString_FromStringAndSize(freeme, res);
 321                         PyMem_Free(freeme);
 322                 }
 323         }
 324
 325         if (temp)
 326                 result = PyObject_CallFunctionObjArgs(pickle_loads,
 327                                                       temp, NULL);
 328         Py_XDECREF(temp);
 329         return result;
 330 }
 331
 332 /*
 333  * Other functions
 334  */
 335
 336 static PyObject *
 337 connection_poll(ConnectionObject *self, PyObject *args)
 338 {
 339         PyObject *timeout_obj = NULL;
 340         double timeout = 0.0;
 341         int res;
 342
 343         CHECK_READABLE(self);
 344
 345         if (!PyArg_ParseTuple(args, "|O", &timeout_obj))
 346                 return NULL;
 347
 348         if (timeout_obj == NULL) {
 349                 timeout = 0.0;
 350         } else if (timeout_obj == Py_None) {
 351                 timeout = -1.0;                         /* block forever */
 352         } else {
 353                 timeout = PyFloat_AsDouble(timeout_obj);
 354                 if (PyErr_Occurred())
 355                         return NULL;
 356                 if (timeout < 0.0)
 357                         timeout = 0.0;
 358         }
 359
 360         Py_BEGIN_ALLOW_THREADS
 361         res = conn_poll(self, timeout, _save);
 362         Py_END_ALLOW_THREADS
 363
 364         switch (res) {
 365         case TRUE:
 366                 Py_RETURN_TRUE;
 367         case FALSE:
 368                 Py_RETURN_FALSE;
 369         default:
 370                 return mp_SetError(PyExc_IOError, res);
 371         }
 372 }
 373
 374 static PyObject *
 375 connection_fileno(ConnectionObject* self)
 376 {
 377         if (self->handle == INVALID_HANDLE_VALUE) {
 378                 PyErr_SetString(PyExc_IOError, "handle is invalid");
 379                 return NULL;
 380         }
 381         return PyInt_FromLong((long)self->handle);
 382 }
 383
 384 static PyObject *
 385 connection_close(ConnectionObject *self)
 386 {
 387         if (self->handle != INVALID_HANDLE_VALUE) {
 388                 Py_BEGIN_ALLOW_THREADS
 389                 CLOSE(self->handle);
 390                 Py_END_ALLOW_THREADS
 391                 self->handle = INVALID_HANDLE_VALUE;
 392         }
 393
 394         Py_RETURN_NONE;
 395 }
 396
 397 static PyObject *
 398 connection_repr(ConnectionObject *self)
 399 {
 400         static char *conn_type[] = {"read-only", "write-only", "read-write"};
 401
 402         assert(self->flags >= 1 && self->flags <= 3);
 403         return FROM_FORMAT("<%s %s, handle %zd>",
 404                            conn_type[self->flags - 1],
 405                            CONNECTION_NAME, (Py_ssize_t)self->handle);
 406 }
 407
 408 /*
 409  * Getters and setters
 410  */
 411
 412 static PyObject *
 413 connection_closed(ConnectionObject *self, void *closure)
 414 {
 415         return PyBool_FromLong((long)(self->handle == INVALID_HANDLE_VALUE));
 416 }
 417
 418 static PyObject *
 419 connection_readable(ConnectionObject *self, void *closure)
 420 {
 421         return PyBool_FromLong((long)(self->flags & READABLE));
 422 }
 423
 424 static PyObject *
 425 connection_writable(ConnectionObject *self, void *closure)
 426 {
 427         return PyBool_FromLong((long)(self->flags & WRITABLE));
 428 }
 429
 430 /*
 431  * Tables
 432  */
 433
 434 static PyMethodDef connection_methods[] = {
 435         {"send_bytes", (PyCFunction)connection_sendbytes, METH_VARARGS,
 436          "send the byte data from a readable buffer-like object"},
 437         {"recv_bytes", (PyCFunction)connection_recvbytes, METH_VARARGS,
 438          "receive byte data as a string"},
 439         {"recv_bytes_into",(PyCFunction)connection_recvbytes_into,METH_VARARGS,
 440          "receive byte data into a writeable buffer-like object\n"
 441          "returns the number of bytes read"},
 442
 443         {"send", (PyCFunction)connection_send_obj, METH_O,
 444          "send a (picklable) object"},
 445         {"recv", (PyCFunction)connection_recv_obj, METH_NOARGS,
 446          "receive a (picklable) object"},
 447
 448         {"poll", (PyCFunction)connection_poll, METH_VARARGS,
 449          "whether there is any input available to be read"},
 450         {"fileno", (PyCFunction)connection_fileno, METH_NOARGS,
 451          "file descriptor or handle of the connection"},
 452         {"close", (PyCFunction)connection_close, METH_NOARGS,
 453          "close the connection"},
 454
 455         {NULL}  /* Sentinel */
 456 };
 457
 458 static PyGetSetDef connection_getset[] = {
 459         {"closed", (getter)connection_closed, NULL,
 460          "True if the connection is closed", NULL},
 461         {"readable", (getter)connection_readable, NULL,
 462          "True if the connection is readable", NULL},
 463         {"writable", (getter)connection_writable, NULL,
 464          "True if the connection is writable", NULL},
 465         {NULL}
 466 };
 467
 468 /*
 469  * Connection type
 470  */
 471
 472 PyDoc_STRVAR(connection_doc,
 473              "Connection type whose constructor signature is\n\n"
 474              "    Connection(handle, readable=True, writable=True).\n\n"
 475              "The constructor does *not* duplicate the handle.");
 476
 477 PyTypeObject CONNECTION_TYPE = {
 478         PyVarObject_HEAD_INIT(NULL, 0)
 479         /* tp_name           */ "_multiprocessing." CONNECTION_NAME,
 480         /* tp_basicsize      */ sizeof(ConnectionObject),
 481         /* tp_itemsize       */ 0,
 482         /* tp_dealloc        */ (destructor)connection_dealloc,
 483         /* tp_print          */ 0,
 484         /* tp_getattr        */ 0,
 485         /* tp_setattr        */ 0,
 486         /* tp_compare        */ 0,
 487         /* tp_repr           */ (reprfunc)connection_repr,
 488         /* tp_as_number      */ 0,
 489         /* tp_as_sequence    */ 0,
 490         /* tp_as_mapping     */ 0,
 491         /* tp_hash           */ 0,
 492         /* tp_call           */ 0,
 493         /* tp_str            */ 0,
 494         /* tp_getattro       */ 0,
 495         /* tp_setattro       */ 0,
 496         /* tp_as_buffer      */ 0,
 497         /* tp_flags          */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
 498                                 Py_TPFLAGS_HAVE_WEAKREFS,
 499         /* tp_doc            */ connection_doc,
 500         /* tp_traverse       */ 0,
 501         /* tp_clear          */ 0,
 502         /* tp_richcompare    */ 0,
 503         /* tp_weaklistoffset */ offsetof(ConnectionObject, weakreflist),
 504         /* tp_iter           */ 0,
 505         /* tp_iternext       */ 0,
 506         /* tp_methods        */ connection_methods,
 507         /* tp_members        */ 0,
 508         /* tp_getset         */ connection_getset,
 509         /* tp_base           */ 0,
 510         /* tp_dict           */ 0,
 511         /* tp_descr_get      */ 0,
 512         /* tp_descr_set      */ 0,
 513         /* tp_dictoffset     */ 0,
 514         /* tp_init           */ 0,
 515         /* tp_alloc          */ 0,
 516         /* tp_new            */ connection_new,
 517 };
 518
 519 #endif /* CONNECTION_H */