~ubuntu-branches/ubuntu/maverick/python3.1/maverick

« back to all changes in this revision

Viewing changes to Modules/_multiprocessing/socket_connection.c

  • Committer: Bazaar Package Importer
  • Author(s): Matthias Klose
  • Date: 2009-03-23 00:01:27 UTC
  • Revision ID: james.westby@ubuntu.com-20090323000127-5fstfxju4ufrhthq
Tags: upstream-3.1~a1+20090322
ImportĀ upstreamĀ versionĀ 3.1~a1+20090322

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
/*
 
2
 * A type which wraps a socket
 
3
 *
 
4
 * socket_connection.c
 
5
 *
 
6
 * Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
 
7
 */
 
8
 
 
9
#include "multiprocessing.h"
 
10
 
 
11
#ifdef MS_WINDOWS
 
12
#  define WRITE(h, buffer, length) send((SOCKET)h, buffer, length, 0)
 
13
#  define READ(h, buffer, length) recv((SOCKET)h, buffer, length, 0)
 
14
#  define CLOSE(h) closesocket((SOCKET)h)
 
15
#else
 
16
#  define WRITE(h, buffer, length) write(h, buffer, length)
 
17
#  define READ(h, buffer, length) read(h, buffer, length)
 
18
#  define CLOSE(h) close(h)
 
19
#endif
 
20
 
 
21
/*
 
22
 * Send string to file descriptor
 
23
 */
 
24
 
 
25
static Py_ssize_t
 
26
_conn_sendall(HANDLE h, char *string, size_t length)
 
27
{
 
28
        char *p = string;
 
29
        Py_ssize_t res;
 
30
 
 
31
        while (length > 0) {
 
32
                res = WRITE(h, p, length);
 
33
                if (res < 0)
 
34
                        return MP_SOCKET_ERROR;
 
35
                length -= res;
 
36
                p += res;
 
37
        }
 
38
 
 
39
        return MP_SUCCESS;
 
40
}
 
41
 
 
42
/*
 
43
 * Receive string of exact length from file descriptor 
 
44
 */
 
45
 
 
46
static Py_ssize_t
 
47
_conn_recvall(HANDLE h, char *buffer, size_t length)
 
48
{
 
49
        size_t remaining = length;
 
50
        Py_ssize_t temp;
 
51
        char *p = buffer;
 
52
 
 
53
        while (remaining > 0) {
 
54
                temp = READ(h, p, remaining);
 
55
                if (temp <= 0) {
 
56
                        if (temp == 0)
 
57
                                return remaining == length ? 
 
58
                                        MP_END_OF_FILE : MP_EARLY_END_OF_FILE;
 
59
                        else
 
60
                                return temp;
 
61
                }
 
62
                remaining -= temp;
 
63
                p += temp;
 
64
        }
 
65
 
 
66
        return MP_SUCCESS;
 
67
}
 
68
 
 
69
/*
 
70
 * Send a string prepended by the string length in network byte order
 
71
 */
 
72
 
 
73
static Py_ssize_t
 
74
conn_send_string(ConnectionObject *conn, char *string, size_t length)
 
75
{
 
76
        Py_ssize_t res;
 
77
        /* The "header" of the message is a 32 bit unsigned number (in
 
78
           network order) which specifies the length of the "body".  If
 
79
           the message is shorter than about 16kb then it is quicker to
 
80
           combine the "header" and the "body" of the message and send
 
81
           them at once. */
 
82
        if (length < (16*1024)) {
 
83
                char *message;
 
84
 
 
85
                message = PyMem_Malloc(length+4);
 
86
                if (message == NULL)
 
87
                        return MP_MEMORY_ERROR;
 
88
 
 
89
                *(UINT32*)message = htonl((UINT32)length);     
 
90
                memcpy(message+4, string, length);
 
91
                Py_BEGIN_ALLOW_THREADS
 
92
                res = _conn_sendall(conn->handle, message, length+4);
 
93
                Py_END_ALLOW_THREADS
 
94
                PyMem_Free(message);
 
95
        } else {
 
96
                UINT32 lenbuff;
 
97
 
 
98
                if (length > MAX_MESSAGE_LENGTH)
 
99
                        return MP_BAD_MESSAGE_LENGTH;
 
100
 
 
101
                lenbuff = htonl((UINT32)length);
 
102
                Py_BEGIN_ALLOW_THREADS
 
103
                res = _conn_sendall(conn->handle, (char*)&lenbuff, 4) || 
 
104
                        _conn_sendall(conn->handle, string, length);
 
105
                Py_END_ALLOW_THREADS
 
106
        }
 
107
        return res;
 
108
}
 
109
 
 
110
/*
 
111
 * Attempts to read into buffer, or failing that into *newbuffer
 
112
 *
 
113
 * Returns number of bytes read.
 
114
 */
 
115
 
 
116
static Py_ssize_t
 
117
conn_recv_string(ConnectionObject *conn, char *buffer, 
 
118
                 size_t buflength, char **newbuffer, size_t maxlength)
 
119
{
 
120
        int res;
 
121
        UINT32 ulength;
 
122
 
 
123
        *newbuffer = NULL;
 
124
 
 
125
        Py_BEGIN_ALLOW_THREADS
 
126
        res = _conn_recvall(conn->handle, (char*)&ulength, 4);
 
127
        Py_END_ALLOW_THREADS
 
128
        if (res < 0)
 
129
                return res;
 
130
 
 
131
        ulength = ntohl(ulength);
 
132
        if (ulength > maxlength)
 
133
                return MP_BAD_MESSAGE_LENGTH;
 
134
 
 
135
        if (ulength <= buflength) {
 
136
                Py_BEGIN_ALLOW_THREADS
 
137
                res = _conn_recvall(conn->handle, buffer, (size_t)ulength);
 
138
                Py_END_ALLOW_THREADS
 
139
                return res < 0 ? res : ulength;
 
140
        } else {
 
141
                *newbuffer = PyMem_Malloc((size_t)ulength);
 
142
                if (*newbuffer == NULL)
 
143
                        return MP_MEMORY_ERROR;
 
144
                Py_BEGIN_ALLOW_THREADS
 
145
                res = _conn_recvall(conn->handle, *newbuffer, (size_t)ulength);
 
146
                Py_END_ALLOW_THREADS
 
147
                return res < 0 ? (Py_ssize_t)res : (Py_ssize_t)ulength;
 
148
        }
 
149
}
 
150
 
 
151
/*
 
152
 * Check whether any data is available for reading -- neg timeout blocks
 
153
 */
 
154
 
 
155
static int
 
156
conn_poll(ConnectionObject *conn, double timeout, PyThreadState *_save)
 
157
{
 
158
        int res;
 
159
        fd_set rfds;
 
160
 
 
161
        /*
 
162
         * Verify the handle, issue 3321. Not required for windows.
 
163
         */ 
 
164
        #ifndef MS_WINDOWS
 
165
                if (((int)conn->handle) < 0 || ((int)conn->handle) >= FD_SETSIZE) {
 
166
                        Py_BLOCK_THREADS
 
167
                        PyErr_SetString(PyExc_IOError, "handle out of range in select()");
 
168
                        Py_UNBLOCK_THREADS
 
169
                        return MP_EXCEPTION_HAS_BEEN_SET;
 
170
                }
 
171
        #endif
 
172
 
 
173
        FD_ZERO(&rfds);
 
174
        FD_SET((SOCKET)conn->handle, &rfds);
 
175
 
 
176
        if (timeout < 0.0) {
 
177
                res = select((int)conn->handle+1, &rfds, NULL, NULL, NULL);
 
178
        } else {
 
179
                struct timeval tv;
 
180
                tv.tv_sec = (long)timeout;
 
181
                tv.tv_usec = (long)((timeout - tv.tv_sec) * 1e6 + 0.5);
 
182
                res = select((int)conn->handle+1, &rfds, NULL, NULL, &tv);
 
183
        }
 
184
 
 
185
        if (res < 0) {
 
186
                return MP_SOCKET_ERROR;
 
187
        } else if (FD_ISSET(conn->handle, &rfds)) {
 
188
                return TRUE;
 
189
        } else {
 
190
                assert(res == 0);
 
191
                return FALSE;
 
192
        }
 
193
}
 
194
 
 
195
/*
 
196
 * "connection.h" defines the Connection type using defs above
 
197
 */
 
198
 
 
199
#define CONNECTION_NAME "Connection"
 
200
#define CONNECTION_TYPE ConnectionType
 
201
 
 
202
#include "connection.h"