1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
|
# -*- coding: utf-8 -*-
from __future__ import absolute_import, division
import errno
import os
import socket
import struct
import sys
from . import TTransportException
class TSocket(object):
"""Socket implementation for client side."""
def __init__(self, host=None, port=None, unix_socket=None,
sock=None, socket_family=socket.AF_INET,
socket_timeout=3000, connect_timeout=None):
"""Initialize a TSocket
TSocket can be initialized in 3 ways:
* host + port. can configure to use AF_INET/AF_INET6
* unix_socket
* socket. should pass already opened socket here.
@param host(str) The host to connect to.
@param port(int) The (TCP) port to connect to.
@param unix_socket(str) The filename of a unix socket to connect to.
@param sock(socket) Initialize with opened socket directly.
If this param used, the host, port and unix_socket params will
be ignored.
@param socket_family(str) socket.AF_INET or socket.AF_INET6. only
take effect when using host/port
@param socket_timeout socket timeout in ms
@param connect_timeout connect timeout in ms, only used in
connection, will be set to socket_timeout if not set.
"""
if sock:
self.sock = sock
elif unix_socket:
self.unix_socket = unix_socket
self.host = None
self.port = None
self.sock = None
else:
self.unix_socket = None
self.host = host
self.port = port
self.sock = None
self.socket_family = socket_family
self.socket_timeout = socket_timeout / 1000 if socket_timeout else None
self.connect_timeout = connect_timeout / 1000 if connect_timeout \
else self.socket_timeout
def _init_sock(self):
if self.unix_socket:
_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
else:
_sock = socket.socket(self.socket_family, socket.SOCK_STREAM)
_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
# socket options
linger = struct.pack('ii', 0, 0)
_sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger)
_sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
self.sock = _sock
def set_handle(self, sock):
self.sock = sock
def set_timeout(self, ms):
"""Backward compat api, will bind the timeout to both connect_timeout
and socket_timeout.
"""
self.socket_timeout = ms / 1000 if (ms and ms > 0) else None
self.connect_timeout = self.socket_timeout
if self.sock is not None:
self.sock.settimeout(self.socket_timeout)
def is_open(self):
return bool(self.sock)
def open(self):
self._init_sock()
addr = self.unix_socket or (self.host, self.port)
try:
if self.connect_timeout:
self.sock.settimeout(self.connect_timeout)
self.sock.connect(addr)
if self.socket_timeout:
self.sock.settimeout(self.socket_timeout)
except (socket.error, OSError):
raise TTransportException(
type=TTransportException.NOT_OPEN,
message="Could not connect to %s" % str(addr))
def read(self, sz):
try:
buff = self.sock.recv(sz)
except socket.error as e:
if (e.args[0] == errno.ECONNRESET and
(sys.platform == 'darwin' or
sys.platform.startswith('freebsd'))):
# freebsd and Mach don't follow POSIX semantic of recv
# and fail with ECONNRESET if peer performed shutdown.
# See corresponding comment and code in TSocket::read()
# in lib/cpp/src/transport/TSocket.cpp.
self.close()
# Trigger the check to raise the END_OF_FILE exception below.
buff = ''
else:
raise
if len(buff) == 0:
raise TTransportException(type=TTransportException.END_OF_FILE,
message='TSocket read 0 bytes')
return buff
def write(self, buff):
self.sock.sendall(buff)
def flush(self):
pass
def close(self):
if not self.sock:
return
try:
self.sock.shutdown(socket.SHUT_RDWR)
self.sock.close()
self.sock = None
except (socket.error, OSError):
pass
class TServerSocket(object):
"""Socket implementation for server side."""
def __init__(self, host=None, port=None, unix_socket=None,
socket_family=socket.AF_INET, client_timeout=3000,
backlog=128):
"""Initialize a TServerSocket
TSocket can be initialized in 2 ways:
* host + port. can configure to use AF_INET/AF_INET6
* unix_socket
@param host(str) The host to connect to
@param port(int) The (TCP) port to connect to
@param unix_socket(str) The filename of a unix socket to connect to
@param socket_family(str) socket.AF_INET or socket.AF_INET6. only
take effect when using host/port
@param client_timeout client socket timeout
@param backlog backlog for server socket
"""
if unix_socket:
self.unix_socket = unix_socket
self.host = None
self.port = None
else:
self.unix_socket = None
self.host = host
self.port = port
self.socket_family = socket_family
self.client_timeout = client_timeout / 1000 if client_timeout else None
self.backlog = backlog
def _init_sock(self):
if self.unix_socket:
# try remove the sock file it already exists
_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
try:
_sock.connect(self.unix_socket)
except (socket.error, OSError) as err:
if err.args[0] == errno.ECONNREFUSED:
os.unlink(self.unix_socket)
else:
_sock = socket.socket(self.socket_family, socket.SOCK_STREAM)
_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
if hasattr(socket, "SO_REUSEPORT"):
_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
_sock.settimeout(None)
self.sock = _sock
def listen(self):
self._init_sock()
addr = self.unix_socket or (self.host, self.port)
self.sock.bind(addr)
self.sock.listen(self.backlog)
def accept(self):
client, _ = self.sock.accept()
if self.client_timeout:
client.settimeout(self.client_timeout)
return TSocket(sock=client)
def close(self):
if not self.sock:
return
try:
self.sock.shutdown(socket.SHUT_RDWR)
self.sock.close()
except (socket.error, OSError):
pass
|
