Mark Slee | 89e2bb8 | 2007-03-01 00:20:36 +0000 | [diff] [blame] | 1 | # |
David Reiss | ea2cba8 | 2009-03-30 21:35:00 +0000 | [diff] [blame] | 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 | # |
Mark Slee | 89e2bb8 | 2007-03-01 00:20:36 +0000 | [diff] [blame] | 19 | |
David Reiss | e29995e | 2008-07-31 20:15:17 +0000 | [diff] [blame] | 20 | import errno |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 21 | import logging |
Bryan Duxbury | 6972041 | 2012-01-03 17:32:30 +0000 | [diff] [blame] | 22 | import os |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 23 | import socket |
David Reiss | 73af3b7 | 2010-08-30 21:57:07 +0000 | [diff] [blame] | 24 | import sys |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 25 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 26 | from .TTransport import TTransportBase, TTransportException, TServerTransportBase |
Bryan Duxbury | 6972041 | 2012-01-03 17:32:30 +0000 | [diff] [blame] | 27 | |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 28 | logger = logging.getLogger(__name__) |
| 29 | |
Bryan Duxbury | 6972041 | 2012-01-03 17:32:30 +0000 | [diff] [blame] | 30 | |
David Reiss | e29995e | 2008-07-31 20:15:17 +0000 | [diff] [blame] | 31 | class TSocketBase(TTransportBase): |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 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, |
max ulidtko | 9b9567b | 2020-04-27 16:04:27 +0300 | [diff] [blame] | 42 | socket.AI_PASSIVE) |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 43 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 44 | def close(self): |
| 45 | if self.handle: |
| 46 | self.handle.close() |
| 47 | self.handle = None |
David Reiss | e29995e | 2008-07-31 20:15:17 +0000 | [diff] [blame] | 48 | |
Bryan Duxbury | 6972041 | 2012-01-03 17:32:30 +0000 | [diff] [blame] | 49 | |
David Reiss | e29995e | 2008-07-31 20:15:17 +0000 | [diff] [blame] | 50 | class TSocket(TSocketBase): |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 51 | """Socket implementation of TTransport base.""" |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 52 | |
Jun | f6511c9 | 2019-02-01 12:07:58 +0800 | [diff] [blame] | 53 | def __init__(self, host='localhost', port=9090, unix_socket=None, |
| 54 | socket_family=socket.AF_UNSPEC, |
| 55 | socket_keepalive=False): |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 56 | """Initialize a TSocket |
David Reiss | c16a8f6 | 2007-12-14 23:46:47 +0000 | [diff] [blame] | 57 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 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. |
Jun | f6511c9 | 2019-02-01 12:07:58 +0800 | [diff] [blame] | 63 | @param socket_keepalive(bool) enable TCP keepalive, default off. |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 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 |
Jun | f6511c9 | 2019-02-01 12:07:58 +0800 | [diff] [blame] | 71 | self._socket_keepalive = socket_keepalive |
David Reiss | 0c90f6f | 2008-02-06 22:18:40 +0000 | [diff] [blame] | 72 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 73 | def setHandle(self, h): |
| 74 | self.handle = h |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 75 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 76 | def isOpen(self): |
Neil Williams | 01d53f4 | 2020-07-07 07:27:29 -0700 | [diff] [blame] | 77 | if self.handle is None: |
| 78 | return False |
| 79 | |
| 80 | # this lets us cheaply see if the other end of the socket is still |
| 81 | # connected. if disconnected, we'll get EOF back (expressed as zero |
| 82 | # bytes of data) otherwise we'll get one byte or an error indicating |
| 83 | # we'd have to block for data. |
| 84 | # |
| 85 | # note that we're not doing this with socket.MSG_DONTWAIT because 1) |
| 86 | # it's linux-specific and 2) gevent-patched sockets hide EAGAIN from us |
| 87 | # when timeout is non-zero. |
| 88 | original_timeout = self.handle.gettimeout() |
| 89 | try: |
| 90 | self.handle.settimeout(0) |
| 91 | try: |
| 92 | peeked_bytes = self.handle.recv(1, socket.MSG_PEEK) |
| 93 | except (socket.error, OSError) as exc: # on modern python this is just BlockingIOError |
| 94 | if exc.errno in (errno.EWOULDBLOCK, errno.EAGAIN): |
| 95 | return True |
| 96 | return False |
| 97 | finally: |
| 98 | self.handle.settimeout(original_timeout) |
| 99 | |
| 100 | # the length will be zero if we got EOF (indicating connection closed) |
| 101 | return len(peeked_bytes) == 1 |
Aditya Agarwal | f954f97 | 2007-02-06 01:26:12 +0000 | [diff] [blame] | 102 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 103 | def setTimeout(self, ms): |
| 104 | if ms is None: |
| 105 | self._timeout = None |
| 106 | else: |
| 107 | self._timeout = ms / 1000.0 |
David Reiss | 0c90f6f | 2008-02-06 22:18:40 +0000 | [diff] [blame] | 108 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 109 | if self.handle is not None: |
| 110 | self.handle.settimeout(self._timeout) |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 111 | |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 112 | def _do_open(self, family, socktype): |
| 113 | return socket.socket(family, socktype) |
| 114 | |
| 115 | @property |
| 116 | def _address(self): |
| 117 | return self._unix_socket if self._unix_socket else '%s:%d' % (self.host, self.port) |
| 118 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 119 | def open(self): |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 120 | if self.handle: |
James E. King III | 3131fe9 | 2019-07-02 14:21:05 -0400 | [diff] [blame] | 121 | raise TTransportException(type=TTransportException.ALREADY_OPEN, message="already open") |
Mark Slee | 2297460 | 2007-07-06 22:20:19 +0000 | [diff] [blame] | 122 | try: |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 123 | addrs = self._resolveAddr() |
James E. King III | 3131fe9 | 2019-07-02 14:21:05 -0400 | [diff] [blame] | 124 | except socket.gaierror as gai: |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 125 | msg = 'failed to resolve sockaddr for ' + str(self._address) |
| 126 | logger.exception(msg) |
James E. King III | 3131fe9 | 2019-07-02 14:21:05 -0400 | [diff] [blame] | 127 | raise TTransportException(type=TTransportException.NOT_OPEN, message=msg, inner=gai) |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 128 | for family, socktype, _, _, sockaddr in addrs: |
| 129 | handle = self._do_open(family, socktype) |
Jun | f6511c9 | 2019-02-01 12:07:58 +0800 | [diff] [blame] | 130 | |
| 131 | # TCP_KEEPALIVE |
| 132 | if self._socket_keepalive: |
| 133 | handle.setsockopt(socket.IPPROTO_TCP, socket.SO_KEEPALIVE, 1) |
| 134 | |
Nobuaki Sukegawa | 1c8b5cb | 2016-02-14 22:04:38 +0900 | [diff] [blame] | 135 | handle.settimeout(self._timeout) |
| 136 | try: |
| 137 | handle.connect(sockaddr) |
| 138 | self.handle = handle |
| 139 | return |
| 140 | except socket.error: |
| 141 | handle.close() |
| 142 | logger.info('Could not connect to %s', sockaddr, exc_info=True) |
| 143 | msg = 'Could not connect to any of %s' % list(map(lambda a: a[4], |
| 144 | addrs)) |
| 145 | logger.error(msg) |
James E. King III | 3131fe9 | 2019-07-02 14:21:05 -0400 | [diff] [blame] | 146 | raise TTransportException(type=TTransportException.NOT_OPEN, message=msg) |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 147 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 148 | def read(self, sz): |
| 149 | try: |
| 150 | buff = self.handle.recv(sz) |
| 151 | except socket.error as e: |
| 152 | if (e.args[0] == errno.ECONNRESET and |
| 153 | (sys.platform == 'darwin' or sys.platform.startswith('freebsd'))): |
| 154 | # freebsd and Mach don't follow POSIX semantic of recv |
| 155 | # and fail with ECONNRESET if peer performed shutdown. |
| 156 | # See corresponding comment and code in TSocket::read() |
| 157 | # in lib/cpp/src/transport/TSocket.cpp. |
| 158 | self.close() |
| 159 | # Trigger the check to raise the END_OF_FILE exception below. |
| 160 | buff = '' |
James E. King III | 3131fe9 | 2019-07-02 14:21:05 -0400 | [diff] [blame] | 161 | elif e.args[0] == errno.ETIMEDOUT: |
| 162 | raise TTransportException(type=TTransportException.TIMED_OUT, message="read timeout", inner=e) |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 163 | else: |
James E. King III | 3131fe9 | 2019-07-02 14:21:05 -0400 | [diff] [blame] | 164 | raise TTransportException(message="unexpected exception", inner=e) |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 165 | if len(buff) == 0: |
| 166 | raise TTransportException(type=TTransportException.END_OF_FILE, |
| 167 | message='TSocket read 0 bytes') |
| 168 | return buff |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 169 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 170 | def write(self, buff): |
| 171 | if not self.handle: |
| 172 | raise TTransportException(type=TTransportException.NOT_OPEN, |
| 173 | message='Transport not open') |
| 174 | sent = 0 |
| 175 | have = len(buff) |
| 176 | while sent < have: |
James E. King III | 3131fe9 | 2019-07-02 14:21:05 -0400 | [diff] [blame] | 177 | try: |
| 178 | plus = self.handle.send(buff) |
| 179 | if plus == 0: |
| 180 | raise TTransportException(type=TTransportException.END_OF_FILE, |
| 181 | message='TSocket sent 0 bytes') |
| 182 | sent += plus |
| 183 | buff = buff[plus:] |
| 184 | except socket.error as e: |
| 185 | raise TTransportException(message="unexpected exception", inner=e) |
Mark Slee | cde2b61 | 2006-09-03 21:13:07 +0000 | [diff] [blame] | 186 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 187 | def flush(self): |
| 188 | pass |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 189 | |
Bryan Duxbury | 6972041 | 2012-01-03 17:32:30 +0000 | [diff] [blame] | 190 | |
David Reiss | d73255d | 2009-03-24 22:51:02 +0000 | [diff] [blame] | 191 | class TServerSocket(TSocketBase, TServerTransportBase): |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 192 | """Socket implementation of TServerTransport base.""" |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 193 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 194 | def __init__(self, host=None, port=9090, unix_socket=None, socket_family=socket.AF_UNSPEC): |
| 195 | self.host = host |
| 196 | self.port = port |
| 197 | self._unix_socket = unix_socket |
| 198 | self._socket_family = socket_family |
| 199 | self.handle = None |
lshgdut | 7af79c8 | 2018-11-21 10:09:42 +0800 | [diff] [blame] | 200 | self._backlog = 128 |
| 201 | |
| 202 | def setBacklog(self, backlog=None): |
| 203 | if not self.handle: |
| 204 | self._backlog = backlog |
| 205 | else: |
| 206 | # We cann't update backlog when it is already listening, since the |
| 207 | # handle has been created. |
| 208 | logger.warn('You have to set backlog before listen.') |
Mark Slee | 256bdc4 | 2007-11-27 08:42:19 +0000 | [diff] [blame] | 209 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 210 | def listen(self): |
| 211 | res0 = self._resolveAddr() |
| 212 | socket_family = self._socket_family == socket.AF_UNSPEC and socket.AF_INET6 or self._socket_family |
| 213 | for res in res0: |
| 214 | if res[0] is socket_family or res is res0[-1]: |
| 215 | break |
Mark Slee | 2297460 | 2007-07-06 22:20:19 +0000 | [diff] [blame] | 216 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 217 | # We need remove the old unix socket if the file exists and |
| 218 | # nobody is listening on it. |
| 219 | if self._unix_socket: |
| 220 | tmp = socket.socket(res[0], res[1]) |
| 221 | try: |
| 222 | tmp.connect(res[4]) |
| 223 | except socket.error as err: |
| 224 | eno, message = err.args |
| 225 | if eno == errno.ECONNREFUSED: |
| 226 | os.unlink(res[4]) |
David Reiss | e29995e | 2008-07-31 20:15:17 +0000 | [diff] [blame] | 227 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 228 | self.handle = socket.socket(res[0], res[1]) |
| 229 | self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) |
| 230 | if hasattr(self.handle, 'settimeout'): |
| 231 | self.handle.settimeout(None) |
| 232 | self.handle.bind(res[4]) |
lshgdut | 7af79c8 | 2018-11-21 10:09:42 +0800 | [diff] [blame] | 233 | self.handle.listen(self._backlog) |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 234 | |
Nobuaki Sukegawa | 10308cb | 2016-02-03 01:57:03 +0900 | [diff] [blame] | 235 | def accept(self): |
| 236 | client, addr = self.handle.accept() |
| 237 | result = TSocket() |
| 238 | result.setHandle(client) |
| 239 | return result |