def _sock_send(self, data):
target_size = len(data)
sent_size = 0
while sent_size < target_size:
try:
pack_size = self._sock.send(data[sent_size:])
except socket.error:
raise exceptions.TransportError(
'Failed to send request to server %s:%d.' %
(self._host, self._port))
if pack_size == 0:
raise exceptions.TransportError(
'Failed to send request to server %s:%d.' %
(self._host, self._port))
sent_size += pack_size
def _connect(self):
"""Connect to the server and send "Hello" message.
# Message consists of three components:
# 1) b'HBas\x00\x50'.
# 2) Big-endian uint32 indicating length of serialized ConnectionHeader.
# 3) Serialized ConnectionHeader.
Raises:
exceptions.TransportError: Failed to connect or send messages.
"""
try:
self._sock.connect((self._host, self._port))
except socket.error:
raise exceptions.TransportError(
'Failed to connect to server %s:%d.' %
(self._host, self._port))
header = pb.ConnectionHeader()
header.user_info.effective_user = '******'
header.service_name = self._service_name
header_bytes = header.SerializeToString()
message = b'HBas\x00\x50' + struct.pack(
'>I', len(header_bytes)) + header_bytes
self._sock.settimeout(60)
self._sock_send(message)
def _sock_recv(self, n):
buffer = io.BytesIO()
received = 0
while received < n:
try:
data = self._sock.recv(n - received)
except socket.error:
raise exceptions.TransportError(
'Failed to receive response to server %s:%d.' %
(self._host, self._port))
if not data:
raise exceptions.TransportError(
'Failed to receive response to server %s:%d.' %
(self._host, self._port))
received += len(data)
buffer.write(data)
return buffer.getvalue()
def _get_response(self, call_id):
with self._resp_lock:
if call_id not in self._resp_dict:
raise exceptions.TransportError(
'FATAL ERROR: Response for call [%d] does not exist.' %
call_id)
resp = self._resp_dict[call_id]
del self._resp_dict[call_id]
return resp
def _get_address(zkquorum, path, timeout=9, retries=3):
"""Get specific server address.
Args:
zkquorum (str):
path (str):
timeout (int):
retries (int):
Returns:
tuple: (hostname, port)
Raises:
exceptions.TransportError: Failed to connect.
exceptions.NoSuchZookeeperNodeError: The required node not found.
exceptions.ZookeeperProtocolError: Invalid response.
"""
zk_client = KazooClient(hosts=zkquorum)
try:
zk_client.start(timeout=timeout)
except KazooTimeoutError:
raise exceptions.TransportError(
'Failed to connect to zookeeper at %s.' % zkquorum)
response, znodestat = None, None
for _ in range(retries + 1):
try:
response, znodestat = zk_client.get(path)
except NoNodeError:
time.sleep(3.0)
continue
else:
break
if response is None:
raise exceptions.NoSuchZookeeperNodeError(
'ZooKeeper does not contain a %s node.' % path)
zk_client.stop()
# the message contains at least 5 bytes with the following structure:
# (1B)(4B)... => (b'\xff')(meta_size)...
if len(response) < 5:
raise exceptions.ZookeeperProtocolError(
'ZooKeeper returned too few response. Response size: %d.' %
len(response))
tag, meta_size = struct.unpack('>cI', response[:5])
if tag != b'\xff':
raise exceptions.ZookeeperProtocolError(
'ZooKeeper returned an invalid response. b\'\\xff\' expected, got %s'
% tag)
if meta_size <= 0 or meta_size > 65000:
raise exceptions.ZookeeperProtocolError(
'ZooKeeper returned an invalid meta size %d.' % meta_size)
# (meta_size B)(4B)... => (meta)(b'PBUF')...
magic = struct.unpack('>4s', response[meta_size + 5:meta_size + 9])[0]
if magic != b'PBUF':
raise exceptions.ZookeeperProtocolError(
'ZooKeeper returned an invalid response. b\'PBUF\' expected, got %s.'
% magic)
meta = pb.MetaRegionServer()
try:
meta.ParseFromString(response[meta_size + 9:])
except DecodeError:
raise exceptions.ZookeeperProtocolError(
'Failed to parse MetaRegionServer from response.')
return meta.server.host_name, meta.server.port
def _put_response(self, call_id, header, data, error):
with self._resp_lock:
if call_id in self._resp_dict:
raise exceptions.TransportError(
'FATAL ERROR: Response for call [%d] exists.' % call_id)
self._resp_dict[call_id] = (header, data, error)