]>
Commit | Line | Data |
---|---|---|
f67539c2 TL |
1 | # |
2 | # Licensed to the Apache Software Foundation (ASF) under one | |
3 | # or more contributor license agreements. See the NOTICE file | |
4 | # distributed with this work for additional information | |
5 | # regarding copyright ownership. The ASF licenses this file | |
6 | # to you under the Apache License, Version 2.0 (the | |
7 | # "License"); you may not use this file except in compliance | |
8 | # with the License. You may obtain a copy of the License at | |
9 | # | |
10 | # http://www.apache.org/licenses/LICENSE-2.0 | |
11 | # | |
12 | # Unless required by applicable law or agreed to in writing, | |
13 | # software distributed under the License is distributed on an | |
14 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
15 | # KIND, either express or implied. See the License for the | |
16 | # specific language governing permissions and limitations | |
17 | # under the License. | |
18 | # | |
19 | ||
20 | import errno | |
21 | import logging | |
22 | import os | |
23 | import socket | |
24 | import sys | |
25 | ||
26 | from .TTransport import TTransportBase, TTransportException, TServerTransportBase | |
27 | ||
28 | logger = logging.getLogger(__name__) | |
29 | ||
30 | ||
31 | class TSocketBase(TTransportBase): | |
32 | def _resolveAddr(self): | |
33 | if self._unix_socket is not None: | |
34 | return [(socket.AF_UNIX, socket.SOCK_STREAM, None, None, | |
35 | self._unix_socket)] | |
36 | else: | |
37 | return socket.getaddrinfo(self.host, | |
38 | self.port, | |
39 | self._socket_family, | |
40 | socket.SOCK_STREAM, | |
41 | 0, | |
42 | socket.AI_PASSIVE | socket.AI_ADDRCONFIG) | |
43 | ||
44 | def close(self): | |
45 | if self.handle: | |
46 | self.handle.close() | |
47 | self.handle = None | |
48 | ||
49 | ||
50 | class TSocket(TSocketBase): | |
51 | """Socket implementation of TTransport base.""" | |
52 | ||
53 | def __init__(self, host='localhost', port=9090, unix_socket=None, | |
54 | socket_family=socket.AF_UNSPEC, | |
55 | socket_keepalive=False): | |
56 | """Initialize a TSocket | |
57 | ||
58 | @param host(str) The host to connect to. | |
59 | @param port(int) The (TCP) port to connect to. | |
60 | @param unix_socket(str) The filename of a unix socket to connect to. | |
61 | (host and port will be ignored.) | |
62 | @param socket_family(int) The socket family to use with this socket. | |
63 | @param socket_keepalive(bool) enable TCP keepalive, default off. | |
64 | """ | |
65 | self.host = host | |
66 | self.port = port | |
67 | self.handle = None | |
68 | self._unix_socket = unix_socket | |
69 | self._timeout = None | |
70 | self._socket_family = socket_family | |
71 | self._socket_keepalive = socket_keepalive | |
72 | ||
73 | def setHandle(self, h): | |
74 | self.handle = h | |
75 | ||
76 | def isOpen(self): | |
77 | return self.handle is not None | |
78 | ||
79 | def setTimeout(self, ms): | |
80 | if ms is None: | |
81 | self._timeout = None | |
82 | else: | |
83 | self._timeout = ms / 1000.0 | |
84 | ||
85 | if self.handle is not None: | |
86 | self.handle.settimeout(self._timeout) | |
87 | ||
88 | def _do_open(self, family, socktype): | |
89 | return socket.socket(family, socktype) | |
90 | ||
91 | @property | |
92 | def _address(self): | |
93 | return self._unix_socket if self._unix_socket else '%s:%d' % (self.host, self.port) | |
94 | ||
95 | def open(self): | |
96 | if self.handle: | |
97 | raise TTransportException(type=TTransportException.ALREADY_OPEN, message="already open") | |
98 | try: | |
99 | addrs = self._resolveAddr() | |
100 | except socket.gaierror as gai: | |
101 | msg = 'failed to resolve sockaddr for ' + str(self._address) | |
102 | logger.exception(msg) | |
103 | raise TTransportException(type=TTransportException.NOT_OPEN, message=msg, inner=gai) | |
104 | for family, socktype, _, _, sockaddr in addrs: | |
105 | handle = self._do_open(family, socktype) | |
106 | ||
107 | # TCP_KEEPALIVE | |
108 | if self._socket_keepalive: | |
109 | handle.setsockopt(socket.IPPROTO_TCP, socket.SO_KEEPALIVE, 1) | |
110 | ||
111 | handle.settimeout(self._timeout) | |
112 | try: | |
113 | handle.connect(sockaddr) | |
114 | self.handle = handle | |
115 | return | |
116 | except socket.error: | |
117 | handle.close() | |
118 | logger.info('Could not connect to %s', sockaddr, exc_info=True) | |
119 | msg = 'Could not connect to any of %s' % list(map(lambda a: a[4], | |
120 | addrs)) | |
121 | logger.error(msg) | |
122 | raise TTransportException(type=TTransportException.NOT_OPEN, message=msg) | |
123 | ||
124 | def read(self, sz): | |
125 | try: | |
126 | buff = self.handle.recv(sz) | |
127 | except socket.error as e: | |
128 | if (e.args[0] == errno.ECONNRESET and | |
129 | (sys.platform == 'darwin' or sys.platform.startswith('freebsd'))): | |
130 | # freebsd and Mach don't follow POSIX semantic of recv | |
131 | # and fail with ECONNRESET if peer performed shutdown. | |
132 | # See corresponding comment and code in TSocket::read() | |
133 | # in lib/cpp/src/transport/TSocket.cpp. | |
134 | self.close() | |
135 | # Trigger the check to raise the END_OF_FILE exception below. | |
136 | buff = '' | |
137 | elif e.args[0] == errno.ETIMEDOUT: | |
138 | raise TTransportException(type=TTransportException.TIMED_OUT, message="read timeout", inner=e) | |
139 | else: | |
140 | raise TTransportException(message="unexpected exception", inner=e) | |
141 | if len(buff) == 0: | |
142 | raise TTransportException(type=TTransportException.END_OF_FILE, | |
143 | message='TSocket read 0 bytes') | |
144 | return buff | |
145 | ||
146 | def write(self, buff): | |
147 | if not self.handle: | |
148 | raise TTransportException(type=TTransportException.NOT_OPEN, | |
149 | message='Transport not open') | |
150 | sent = 0 | |
151 | have = len(buff) | |
152 | while sent < have: | |
153 | try: | |
154 | plus = self.handle.send(buff) | |
155 | if plus == 0: | |
156 | raise TTransportException(type=TTransportException.END_OF_FILE, | |
157 | message='TSocket sent 0 bytes') | |
158 | sent += plus | |
159 | buff = buff[plus:] | |
160 | except socket.error as e: | |
161 | raise TTransportException(message="unexpected exception", inner=e) | |
162 | ||
163 | def flush(self): | |
164 | pass | |
165 | ||
166 | ||
167 | class TServerSocket(TSocketBase, TServerTransportBase): | |
168 | """Socket implementation of TServerTransport base.""" | |
169 | ||
170 | def __init__(self, host=None, port=9090, unix_socket=None, socket_family=socket.AF_UNSPEC): | |
171 | self.host = host | |
172 | self.port = port | |
173 | self._unix_socket = unix_socket | |
174 | self._socket_family = socket_family | |
175 | self.handle = None | |
176 | self._backlog = 128 | |
177 | ||
178 | def setBacklog(self, backlog=None): | |
179 | if not self.handle: | |
180 | self._backlog = backlog | |
181 | else: | |
182 | # We cann't update backlog when it is already listening, since the | |
183 | # handle has been created. | |
184 | logger.warn('You have to set backlog before listen.') | |
185 | ||
186 | def listen(self): | |
187 | res0 = self._resolveAddr() | |
188 | socket_family = self._socket_family == socket.AF_UNSPEC and socket.AF_INET6 or self._socket_family | |
189 | for res in res0: | |
190 | if res[0] is socket_family or res is res0[-1]: | |
191 | break | |
192 | ||
193 | # We need remove the old unix socket if the file exists and | |
194 | # nobody is listening on it. | |
195 | if self._unix_socket: | |
196 | tmp = socket.socket(res[0], res[1]) | |
197 | try: | |
198 | tmp.connect(res[4]) | |
199 | except socket.error as err: | |
200 | eno, message = err.args | |
201 | if eno == errno.ECONNREFUSED: | |
202 | os.unlink(res[4]) | |
203 | ||
204 | self.handle = socket.socket(res[0], res[1]) | |
205 | self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | |
206 | if hasattr(self.handle, 'settimeout'): | |
207 | self.handle.settimeout(None) | |
208 | self.handle.bind(res[4]) | |
209 | self.handle.listen(self._backlog) | |
210 | ||
211 | def accept(self): | |
212 | client, addr = self.handle.accept() | |
213 | result = TSocket() | |
214 | result.setHandle(client) | |
215 | return result |