def test_gen_request(mocker):
mock_addr = 'mock_addr:1234'
operator = mocker.Mock()
master = mocker.Mock()
driver = MesosOperatorMasterDriver(master, operator)
driver._master = mock_addr
req = driver.gen_request()
parser = HttpParser(0)
assert len(req) == parser.execute(req, len(req))
assert parser.is_headers_complete()
assert parser.get_method() == 'POST'
assert parser.get_url() == '/api/v1/operator'
assert parser.is_partial_body()
body = parser.recv_body()
result = json.loads(body.decode('utf-8'))
assert result['type'] == 'SUBSCRIBE'
headers = {k.upper(): v for k, v in parser.get_headers().items()}
assert headers == {
'HOST': mock_addr,
'CONTENT-TYPE': 'application/json',
'ACCEPT': 'application/json',
'CONNECTION': 'close',
'CONTENT-LENGTH': str(len(body))
}
assert parser.is_message_complete()
def test_gen_request(mocker):
mock_addr = 'mock_addr:1234'
sched = mocker.Mock()
framework = {'failover_timeout': 0}
master = mocker.Mock()
driver = MesosSchedulerDriver(sched, framework, master)
driver._master = mock_addr
req = driver.gen_request()
parser = HttpParser(0)
assert len(req) == parser.execute(req, len(req))
assert parser.is_headers_complete()
assert parser.get_method() == 'POST'
assert parser.get_url() == '/api/v1/scheduler'
assert parser.is_partial_body()
body = parser.recv_body()
result = json.loads(body.decode('utf-8'))
assert result['type'] == 'SUBSCRIBE'
assert result['subscribe'] == {'framework_info': framework}
headers = {k.upper(): v for k, v in parser.get_headers().items()}
assert headers == {
'HOST': mock_addr,
'CONTENT-TYPE': 'application/json',
'ACCEPT': 'application/json',
'CONNECTION': 'close',
'CONTENT-LENGTH': str(len(body))
}
assert parser.is_message_complete()
def test_gen_request(mocker):
mock_addr = 'mock_addr:1234'
sched = mocker.Mock()
framework = {
'failover_timeout': 0
}
master = mocker.Mock()
driver = MesosSchedulerDriver(sched, framework, master)
driver._master = mock_addr
req = driver.gen_request()
parser = HttpParser(0)
assert len(req) == parser.execute(req, len(req))
assert parser.is_headers_complete()
assert parser.get_method() == 'POST'
assert parser.get_url() == '/api/v1/scheduler'
assert parser.is_partial_body()
body = parser.recv_body()
result = json.loads(body.decode('utf-8'))
assert result['type'] == 'SUBSCRIBE'
assert result['subscribe'] == {
'framework_info': framework
}
headers = {k.upper(): v for k, v in parser.get_headers().items()}
assert headers == {
'HOST': mock_addr,
'CONTENT-TYPE': 'application/json',
'ACCEPT': 'application/json',
'CONNECTION': 'close',
'CONTENT-LENGTH': str(len(body))
}
assert parser.is_message_complete()
def test_gen_request(mocker):
mock_addr = 'mock_addr:1234'
operator = mocker.Mock()
master = mocker.Mock()
driver = MesosOperatorMasterDriver(master, operator)
driver._master = mock_addr
req = driver.gen_request()
parser = HttpParser(0)
assert len(req) == parser.execute(req, len(req))
assert parser.is_headers_complete()
assert parser.get_method() == 'POST'
assert parser.get_url() == '/api/v1/operator'
assert parser.is_partial_body()
body = parser.recv_body()
result = json.loads(body.decode('utf-8'))
assert result['type'] == 'SUBSCRIBE'
headers = {k.upper(): v for k, v in parser.get_headers().items()}
assert headers == {
'HOST': mock_addr,
'CONTENT-TYPE': 'application/json',
'ACCEPT': 'application/json',
'CONNECTION': 'close',
'CONTENT-LENGTH': str(len(body))
}
assert parser.is_message_complete()
def test_gen_request(mocker):
agent_addr = 'mock_addr:12345'
framework_id = str(uuid.uuid4())
executor_id = str(uuid.uuid4())
env = {
'MESOS_LOCAL': 'true',
'MESOS_AGENT_ENDPOINT': agent_addr,
'MESOS_FRAMEWORK_ID': framework_id,
'MESOS_EXECUTOR_ID': executor_id,
}
mocker.patch('os.environ', env)
exc = mocker.Mock()
driver = MesosExecutorDriver(exc)
driver._master = agent_addr
assert driver.framework_id == dict(value=framework_id)
assert driver.executor_id == dict(value=executor_id)
assert -1e-5 < driver.grace_shutdown_period < 1e-5
assert not driver.checkpoint
assert driver.executor is exc
req = driver.gen_request()
parser = HttpParser(0)
assert len(req) == parser.execute(req, len(req))
assert parser.is_headers_complete()
assert parser.get_method() == 'POST'
assert parser.get_url() == '/api/v1/executor'
assert parser.is_partial_body()
body = parser.recv_body()
result = json.loads(body.decode('utf-8'))
assert result == {
'type': 'SUBSCRIBE',
'framework_id': {
'value': framework_id,
},
'executor_id': {
'value': executor_id,
},
'subscribe': {
'unacknowledged_tasks': [],
'unacknowledged_updates': [],
}
}
headers = {k.upper(): v for k, v in parser.get_headers().items()}
assert headers == {
'HOST': agent_addr,
'CONTENT-TYPE': 'application/json',
'ACCEPT': 'application/json',
'CONNECTION': 'close',
'CONTENT-LENGTH': str(len(body))
}
assert parser.is_message_complete()
def handle(cls, msg):
parser = HttpParser()
parser.execute(msg, len(msg))
url = parser.get_url()
if cls.path and cls.path in url:
cls.record(
"%s://%s%s" %
(str("https" if cls.ishttps else "http"), cls.host, url),
json.dumps(parser.get_headers()), parser.recv_body())
# Support extension
callback = None
if cls.ext_file:
ext = __import__(cls.ext_file[:-3]).EXT
callback = ext.do(url)
return cls.send(msg, callback)
class Connection(object):
def __init__(self, addr, callback):
host, port = addr.split(':', 2)
port = int(port)
self._addr = (host, port)
self._sock = socket.socket()
self._sock.setblocking(0)
try:
self._sock.connect(self._addr)
except socket.error as e:
if e.errno != errno.EAGAIN and e.errno != errno.EINPROGRESS:
raise
self._parser = HttpParser()
self._callback = callback
self._stream_id = None
self._request = callback.gen_request()
self._response = b''
@property
def addr(self):
return self._addr
@property
def stream_id(self):
return self._stream_id
def write(self):
try:
sent = self._sock.send(self._request)
self._request = self._request[sent:]
return True
except socket.error as e:
if e.errno == errno.EAGAIN:
return True
logger.exception('Failed to send to %s', self._addr)
return False
def read(self):
try:
buf = self._sock.recv(select.PIPE_BUF)
n_recv = len(buf)
if n_recv == 0:
logger.error('Remote %s closed', self.addr)
return False
n_parsed = self._parser.execute(buf, n_recv)
if n_parsed != n_recv:
raise RuntimeError('Failed to parse')
if self._stream_id is None and self._parser.is_headers_complete():
code = self._parser.get_status_code()
if code == TEMPORARY_REDIRECT:
headers = {
k.upper(): v
for k, v in list(self._parser.get_headers().items())
}
new_master = headers['LOCATION']
new_master = urlparse(new_master).netloc or new_master
logger.warning('Try to redirect to new master: %s',
new_master)
self._callback.change_master(new_master)
return False
elif code == SERVICE_UNAVAILABLE:
logger.warnig('Master is not available, retry.')
return False
elif code != OK:
msg = self._parser.recv_body()
if not self._parser.is_message_complete():
msg += ' ...'
raise RuntimeError('Failed with HTTP %s: %s' % (code, msg))
if not self._parser.is_chunked():
raise RuntimeError('Response is not chunked')
headers = {
k.upper(): v
for k, v in list(self._parser.get_headers().items())
}
self._stream_id = headers.get('MESOS-STREAM-ID', '')
self._callback.stream_id = self._stream_id
if self._parser.is_partial_body():
self._response += self._parser.recv_body()
while True:
m = LENGTH_PATTERN.match(self._response)
if not m:
break
captured = m.group(0)
length = int(captured.strip())
if len(self._response) < len(captured) + length:
break
data = self._response[len(captured):len(captured) + length]
self._response = self._response[len(captured) + length:]
try:
event = json.loads(data.decode('utf-8'))
except Exception:
logger.exception('Failed parse json %s', data)
raise
try:
self._callback.process_event(event)
except Exception:
logger.exception('Failed to process event')
raise
if self._parser.is_message_complete():
logger.debug('Event stream ended')
return False
return True
except socket.error as e:
if e.errno == errno.EAGAIN:
return True
logger.exception('Failed to recv from %s', self._addr)
return False
def want_write(self):
return bool(self._request)
def fileno(self):
return self._sock.fileno()
def close(self):
self._sock.close()
self._sock = None
self._parser = None
self._request = None
self._response = None
self._callback.on_close()
class Connection(object):
def __init__(self, addr, callback):
host, port = addr.split(':', 2)
port = int(port)
self._addr = (host, port)
self._sock = socket.socket()
self._sock.setblocking(0)
self.connected = False
try:
self._sock.connect(self._addr)
except socket.error as e:
if e.errno != errno.EAGAIN and e.errno != errno.EINPROGRESS:
raise
self._parser = HttpParser()
self._callback = callback
self._stream_id = None
self._request = callback.gen_request()
self._response = b''
@property
def addr(self):
return self._addr
@property
def stream_id(self):
return self._stream_id
def handle_connect_event(self):
err = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
if err == 0:
self.connected = True
return ""
return _strerror(err)
def write(self):
try:
sent = self._sock.send(self._request)
self._request = self._request[sent:]
return True
except socket.error as e:
if e.errno == errno.EAGAIN:
return True
logger.exception('Failed to send to %s', self._addr)
return False
def read(self):
try:
buf = self._sock.recv(PIPE_BUF)
n_recv = len(buf)
if n_recv == 0:
logger.error('Remote %s closed', self.addr)
return False
n_parsed = self._parser.execute(buf, n_recv)
if n_parsed != n_recv:
if hasattr(self._parser, 'errno'):
# using http_parser.pyparser
from http_parser.pyparser import INVALID_CHUNK
if self._parser.errno == INVALID_CHUNK:
# need more chunk data
return True
if self._parser.errno:
raise RuntimeError(
'Failed to parse, code:%s %s' % (
self._parser.errno,
self._parser.errstr,
)
)
else:
raise RuntimeError(
'Failed to parse, code:%s' % (
self._parser.get_errno(),
)
)
if self._stream_id is None and self._parser.is_headers_complete():
code = self._parser.get_status_code()
if code == TEMPORARY_REDIRECT:
headers = {
k.upper(): v
for k, v in list(self._parser.get_headers().items())
}
new_master = headers['LOCATION']
new_master = urlparse(new_master).netloc or new_master
logger.warning(
'Try to redirect to new master: %s', new_master
)
self._callback.change_master(new_master)
return False
elif code == SERVICE_UNAVAILABLE:
logger.warning('Master is not available, retry.')
return False
elif code != OK:
msg = self._parser.recv_body()
if not self._parser.is_message_complete():
msg += ' ...'
raise RuntimeError('Failed with HTTP %s: %s' % (code, msg))
if not self._parser.is_chunked():
raise RuntimeError('Response is not chunked')
headers = {
k.upper(): v
for k, v in list(self._parser.get_headers().items())
}
self._stream_id = headers.get('MESOS-STREAM-ID', '')
self._callback.stream_id = self._stream_id
if self._parser.is_partial_body():
self._response += self._parser.recv_body()
while True:
m = LENGTH_PATTERN.match(self._response)
if not m:
break
captured = m.group(0)
length = int(captured.strip())
if len(self._response) < len(captured) + length:
break
data = self._response[
len(captured):len(captured) + length]
self._response = self._response[
len(captured) + length:]
try:
event = json.loads(data.decode('utf-8'))
except Exception:
logger.exception('Failed parse json %s', data)
raise
try:
self._callback.process_event(event)
except Exception:
logger.exception('Failed to process event')
raise
if self._parser.is_message_complete():
logger.debug('Event stream ended')
return False
return True
except socket.error as e:
if e.errno == errno.EAGAIN:
return True
logger.exception('Failed to recv from %s', self._addr)
return False
def want_write(self):
return bool(self._request)
def fileno(self):
return self._sock.fileno()
def close(self):
self._sock.close()
self._sock = None
self.connected = False
self._parser = None
self._request = None
self._response = None
self._callback.on_close()