blob: 3c7a3ca7db6f37440687adaded026cefcbd09a83 [file] [log] [blame]
Mark Slee89e2bb82007-03-01 00:20:36 +00001#
David Reissea2cba82009-03-30 21:35:00 +00002# 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 Slee89e2bb82007-03-01 00:20:36 +000019
David Reisse29995e2008-07-31 20:15:17 +000020import errno
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +090021import logging
Bryan Duxbury69720412012-01-03 17:32:30 +000022import os
Mark Sleecde2b612006-09-03 21:13:07 +000023import socket
David Reiss73af3b72010-08-30 21:57:07 +000024import sys
Mark Sleecde2b612006-09-03 21:13:07 +000025
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090026from .TTransport import TTransportBase, TTransportException, TServerTransportBase
Bryan Duxbury69720412012-01-03 17:32:30 +000027
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +090028logger = logging.getLogger(__name__)
29
Bryan Duxbury69720412012-01-03 17:32:30 +000030
David Reisse29995e2008-07-31 20:15:17 +000031class TSocketBase(TTransportBase):
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090032 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 ulidtko9b9567b2020-04-27 16:04:27 +030042 socket.AI_PASSIVE)
Mark Sleecde2b612006-09-03 21:13:07 +000043
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090044 def close(self):
45 if self.handle:
46 self.handle.close()
47 self.handle = None
David Reisse29995e2008-07-31 20:15:17 +000048
Bryan Duxbury69720412012-01-03 17:32:30 +000049
David Reisse29995e2008-07-31 20:15:17 +000050class TSocket(TSocketBase):
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090051 """Socket implementation of TTransport base."""
Mark Sleecde2b612006-09-03 21:13:07 +000052
Junf6511c92019-02-01 12:07:58 +080053 def __init__(self, host='localhost', port=9090, unix_socket=None,
54 socket_family=socket.AF_UNSPEC,
55 socket_keepalive=False):
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090056 """Initialize a TSocket
David Reissc16a8f62007-12-14 23:46:47 +000057
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090058 @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.
Junf6511c92019-02-01 12:07:58 +080063 @param socket_keepalive(bool) enable TCP keepalive, default off.
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090064 """
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
Junf6511c92019-02-01 12:07:58 +080071 self._socket_keepalive = socket_keepalive
David Reiss0c90f6f2008-02-06 22:18:40 +000072
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090073 def setHandle(self, h):
74 self.handle = h
Mark Sleec9676562006-09-05 17:34:52 +000075
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +090076 def isOpen(self):
Neil Williams01d53f42020-07-07 07:27:29 -070077 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 Agarwalf954f972007-02-06 01:26:12 +0000102
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900103 def setTimeout(self, ms):
104 if ms is None:
105 self._timeout = None
106 else:
107 self._timeout = ms / 1000.0
David Reiss0c90f6f2008-02-06 22:18:40 +0000108
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900109 if self.handle is not None:
110 self.handle.settimeout(self._timeout)
Mark Sleecde2b612006-09-03 21:13:07 +0000111
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +0900112 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 Sukegawa10308cb2016-02-03 01:57:03 +0900119 def open(self):
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +0900120 if self.handle:
James E. King III3131fe92019-07-02 14:21:05 -0400121 raise TTransportException(type=TTransportException.ALREADY_OPEN, message="already open")
Mark Slee22974602007-07-06 22:20:19 +0000122 try:
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +0900123 addrs = self._resolveAddr()
James E. King III3131fe92019-07-02 14:21:05 -0400124 except socket.gaierror as gai:
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +0900125 msg = 'failed to resolve sockaddr for ' + str(self._address)
126 logger.exception(msg)
James E. King III3131fe92019-07-02 14:21:05 -0400127 raise TTransportException(type=TTransportException.NOT_OPEN, message=msg, inner=gai)
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +0900128 for family, socktype, _, _, sockaddr in addrs:
129 handle = self._do_open(family, socktype)
Junf6511c92019-02-01 12:07:58 +0800130
131 # TCP_KEEPALIVE
132 if self._socket_keepalive:
133 handle.setsockopt(socket.IPPROTO_TCP, socket.SO_KEEPALIVE, 1)
134
Nobuaki Sukegawa1c8b5cb2016-02-14 22:04:38 +0900135 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 III3131fe92019-07-02 14:21:05 -0400146 raise TTransportException(type=TTransportException.NOT_OPEN, message=msg)
Mark Sleecde2b612006-09-03 21:13:07 +0000147
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900148 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 III3131fe92019-07-02 14:21:05 -0400161 elif e.args[0] == errno.ETIMEDOUT:
162 raise TTransportException(type=TTransportException.TIMED_OUT, message="read timeout", inner=e)
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900163 else:
James E. King III3131fe92019-07-02 14:21:05 -0400164 raise TTransportException(message="unexpected exception", inner=e)
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900165 if len(buff) == 0:
166 raise TTransportException(type=TTransportException.END_OF_FILE,
167 message='TSocket read 0 bytes')
168 return buff
Mark Sleecde2b612006-09-03 21:13:07 +0000169
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900170 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 III3131fe92019-07-02 14:21:05 -0400177 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 Sleecde2b612006-09-03 21:13:07 +0000186
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900187 def flush(self):
188 pass
Mark Sleec9676562006-09-05 17:34:52 +0000189
Bryan Duxbury69720412012-01-03 17:32:30 +0000190
David Reissd73255d2009-03-24 22:51:02 +0000191class TServerSocket(TSocketBase, TServerTransportBase):
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900192 """Socket implementation of TServerTransport base."""
Mark Sleec9676562006-09-05 17:34:52 +0000193
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900194 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
lshgdut7af79c82018-11-21 10:09:42 +0800200 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 Slee256bdc42007-11-27 08:42:19 +0000209
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900210 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 Slee22974602007-07-06 22:20:19 +0000216
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900217 # 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 Reisse29995e2008-07-31 20:15:17 +0000227
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900228 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])
lshgdut7af79c82018-11-21 10:09:42 +0800233 self.handle.listen(self._backlog)
Mark Sleec9676562006-09-05 17:34:52 +0000234
Nobuaki Sukegawa10308cb2016-02-03 01:57:03 +0900235 def accept(self):
236 client, addr = self.handle.accept()
237 result = TSocket()
238 result.setHandle(client)
239 return result