def shutdown_callback(self, request=None):
self.logger.info('Service shutting down')
self.demoted_callback()
# wailt for all channel loops to exit
gevent.sleep(Channel.MESSAGE_WAIT_PERIOD_IN_SECONDS)
self.shutdown = True
self.shutdown_event.set()
# self.demoted_callback()
ChannelFactory.get_instance().refresh_master_message_broker(ttl=0)
def test_register_framework():
print("Registering framework")
cm = ConfigManager.get_instance()
cf = ChannelFactory.get_instance()
framework_env = {
'URB_CONFIG_FILE' : cm.get_config_file(),
'URB_FRAMEWORK_ID' : '1',
'URB_FRAMEWORK_NAME':'test',
'URB_MASTER' : cf.get_message_broker_connection_url(),
}
adapter = K8SAdapter()
adapter.set_command(["/bin/sh", "-c", "env; sleep 100"])
max_tasks = 5
concurrent_tasks = 1
# specify unused parameters
kwargs = {'class': 'class',
'submit_options': 'submit_options'};
job_ids = adapter.register_framework(max_tasks, concurrent_tasks, framework_env, **kwargs)
print 'JOB IDs: ', job_ids
framework = {
'job_ids' : job_ids,
'name' : 'framework_name',
'id' : {'value' : 1},
'config' : {'mem' : 1024, 'disk' : 16384, 'cpus' : 1,
'ports' : '[(30000,31000)]',
'max_rejected_offers' : 1,
'max_tasks' : 5
}
}
FrameworkTracker.get_instance().add(1, framework)
assert job_ids != None
def test_register_framework():
print 'Creating service and client channels'
client_channel_name = 'urb.endpoint.1.notify'
channel_factory = ChannelFactory.get_instance()
client_channel = channel_factory.create_channel(client_channel_name)
client_channel.clear()
framework = {}
framework['name'] = 'TestFrameworkC++'
register_msg = {}
register_msg['framework'] = framework
payload = {'mesos.internal.RegisterFrameworkMessage' : register_msg}
msg = {'target' : 'RegisterFrameworkMessage',
'reply_to' : client_channel_name,
'source_id' : 'urb.endpoint.1',
'payload' : payload,
}
print 'Writing register framework message'
SERVICE_CHANNEL.write(json.dumps(msg))
msg2 = client_channel.read_blocking(timeout=5)
print 'Service response: ', msg2
assert msg2 != None
msg2 = json.loads(msg2[1])
# Allow errors temporarily until UGE test config is done
assert (msg2.get('payload') != None or msg2.get('error') != None)
def __init__(self, channel, initial_retry_interval, max_retry_count):
self.__manage = False
self.logger = LogManager.get_instance().get_logger(
self.__class__.__name__)
self.__thread_event = gevent.event.Event()
self.channel = channel
retry_channel_name = channel.name + '.retry'
cf = ChannelFactory.get_instance()
self.retry_channel = cf.create_channel(retry_channel_name)
self.initial_retry_interval = initial_retry_interval
self.max_retry_count = max_retry_count
def test_write_read():
channel_factory = ChannelFactory.get_instance()
channel = channel_factory.create_channel('test_channel')
channel.clear()
msg = 'test msg'
print 'Writing message: ', msg
channel.write(msg)
msg2 = channel.read()
if msg2:
msg2 = msg2[1]
print 'Read message: ', msg2
assert msg2 == msg
def send_shutdown_message(self):
try:
cf = ChannelFactory.get_instance()
channel = cf.create_channel(self.service_monitor_channel)
payload = { 'timestamp' : time.time() }
message = ServiceShutdownMessage(source_id='urb.service.controller', payload=payload)
self.logger.debug(
'Sending shutdown message to channel %s' % \
self.service_monitor_channel)
channel.write(message.to_json())
except Exception, ex:
self.logger.error(
'Could not send shutdown message to channel %s (error: %s)' % \
(self.service_monitor_channel, ex))
def send_response(self, request, response):
if response is None:
return
channel_name = "Unset"
try:
channel_name = self.get_response_channel_name(request)
if channel_name is None:
return
cf = ChannelFactory.get_instance()
channel = cf.create_channel(channel_name)
channel.write(response.to_json())
self.logger.debug('Wrote response to %s: %s' %
(channel_name, response.to_json()))
except Exception, ex:
self.logger.error(
'Could not send response %s to channel %s (error: %s)' % \
(response, channel_name, ex))
def elect(self):
""" Monitoring thread. """
self.logger.debug('MasterElector: Entering election loop')
cf = ChannelFactory.get_instance()
while True:
if not self.__run:
break
self.__thread_event.clear()
were_mb = self.is_master_message_broker()
try:
self.__master_broker = cf.determine_master_message_broker(
value=self.__server_id,
ttl=MasterElector.ELECTOR_TTL_IN_SECONDS)
except Exception, ex:
self.logger.error(
"MasterElector: Error determining master broker: %s" % ex)
self.__thread_event.wait(
MasterElector.ELECTOR_SLEEP_PERIOD_IN_SECONDS)
continue
self.logger.debug("The current master broker is: %s" %
self.__master_broker)
if self.is_master_message_broker():
# We are the master broker... lets make sure that our ttl is updated
self.logger.debug("MasterElector: we are master broker")
try:
cf.refresh_master_message_broker(
ttl=MasterElector.ELECTOR_TTL_IN_SECONDS)
except Exception, ex:
self.logger.error(
"MasterElector: Unable to refresh the master broker: %s"
% ex)
# We are now and if we weren't before trigger the callback
if not were_mb and self.elected_callback:
try:
self.elected_callback()
self.logger.debug(
"MasterElector: Elected callback complete.")
except Exception, ex:
self.logger.warn(
"MasterElector: Exception calling demote callback")
self.logger.exception(ex)
def test_write_callback():
channel_factory = ChannelFactory.get_instance()
channel = channel_factory.create_channel('test_channel')
channel.clear()
class WriteCallback:
def __init__(self):
self.written_msg = None
def write_callback(self, msg):
print 'Entering write callback with message: ', msg
self.written_msg = msg
msg = 'test msg'
write_callback = WriteCallback()
channel.register_write_callback(write_callback.write_callback)
channel.write(msg)
gevent.sleep(0.001)
assert write_callback.written_msg == msg
def test_read_callback():
channel_factory = ChannelFactory.get_instance()
channel = channel_factory.create_channel('test_channel')
channel.clear()
class ReadCallback:
def __init__(self):
self.read_msg = None
def read_callback(self, msg):
print 'Entering read callback with message: ', msg
self.read_msg = msg[1]
msg = 'test msg'
read_callback = ReadCallback()
channel.register_read_callback(read_callback.read_callback)
channel.write(msg)
channel.read()
gevent.sleep(0.001)
assert read_callback.read_msg == msg
def test_ping():
service_thread = create_service_thread(serve_time=3)
print 'Starting service thread'
service_thread.start()
gevent.sleep(1)
print 'Creating service and client channels'
channel_factory = ChannelFactory.get_instance()
service_channel = channel_factory.create_channel('urb.service.monitor')
client_channel_name = 'urb.client.response'
client_channel = channel_factory.create_channel(client_channel_name)
client_channel.clear()
msg = {'target': 'PingMessage', 'reply_to': client_channel_name}
print 'Writing ping message'
service_channel.write(json.dumps(msg))
msg2 = client_channel.read_blocking(timeout=5)
assert msg2 != None
msg2 = json.loads(msg2[1])
payload = msg2.get('payload')
assert payload.get('ack') != None
print 'Service response: ', payload.get('ack')
service_thread.urb_service.shutdown_callback()
service_thread.join()
def __init__(self, channel_name=None):
self.name = self.__class__.__name__
self.logger = LogManager.get_instance().get_logger(self.name)
cf = ChannelFactory.get_instance()
self.channel = cf.create_channel(channel_name)
self.channel.register_read_callback(self.handle)
def configure(self):
cm = ConfigManager.get_instance()
self.service_monitor_endpoint = cm.get_config_option('ExecutorRunner',
'service_monitor_endpoint')
mesos_master_endpoint = cm.get_config_option('ExecutorRunner',
'mesos_master_endpoint')
self.mesos_work_dir = cm.get_config_option('ExecutorRunner',
'mesos_work_dir') % { 'tmp' : os.environ.get('TMP','') }
# Create our working directory before we start logging
if not os.path.exists(self.mesos_work_dir):
os.makedirs(self.mesos_work_dir)
os.chdir(self.mesos_work_dir)
self.logger = LogManager.get_instance().get_logger(self.__class__.__name__)
self.logger.debug('Config file: %s' % cm.get_config_file())
cf = ChannelFactory.get_instance()
self.service_monitor_channel = cf.create_channel(self.service_monitor_endpoint)
self.mesos_channel = cf.create_channel(mesos_master_endpoint)
self.channel_name = cf.get_unique_channel_name()
self.logger.debug('Got my channel name: %s' % self.channel_name)
self.host = cf.get_message_broker_connection_host()
self.port = cf.get_message_broker_connection_port()
# The service will send messages to our notify channel
channel_id = MessagingUtility.get_endpoint_id(self.channel_name)
self.notify_channel_name = MessagingUtility.get_notify_channel_name(None,channel_id)
self.handler_list = [ExecutorHandler(self.notify_channel_name,self)]
# Get various id objects
self.framework_id = {
'value' : os.environ['URB_FRAMEWORK_ID']
}
if 'JOB_ID' in os.environ:
self.job_id = os.environ['JOB_ID']
self.task_id = os.environ.get('SGE_TASK_ID', '1')
slave_id = NamingUtility.create_slave_id(self.job_id, self.task_id, self.notify_channel_name)
else:
self.job_id = uuid.uuid1().hex
self.logger.error('Environment variable JOB_ID is not defined, autogenerating job id: ' % self.job_id)
self.task_id = "1"
slave_id = NamingUtility.create_slave_id(self.job_id, self.task_id , self.notify_channel_name)
self.slave_id = {
"value" : slave_id
}
self.logger.debug('slave id: %s' % self.slave_id)
self.slave_info = {
'hostname' : self.host,
'port' : self.port,
'id' : self.slave_id
}
self.dummy_framework_info = {
'name':'default',
'user':'******',
}
self.heartbeat_channel_info = {
'channel_id' : self.notify_channel_name,
'framework_id' : self.framework_id['value'],
'slave_id' : self.slave_id['value'],
'endpoint_type' : 'executor_runner',
}
def test_constructor():
channel_factory = ChannelFactory.get_instance()
assert isinstance(channel_factory, ChannelFactory)
def test_listen():
# Try following:
# 1) Service listens on service read channel
# 2) Client writes on client write (service read) channel,
# and enters write callback to listen on client read channel
# 3) Service receives message on service read channel,
# and enters read callback to write on service write (client read)
# channel
# 4) Client receives response message
channel_factory = ChannelFactory.get_instance()
service_read_channel = channel_factory.create_channel('urb.service.ping')
client_write_channel = channel_factory.create_channel('urb.service.ping')
service_write_channel = channel_factory.create_channel('urb.client.ping')
client_read_channel = channel_factory.create_channel('urb.client.ping')
client_read_channel.clear()
service_read_channel.clear()
class ClientWriteCallback:
def __init__(self, response_channel):
self.response_msg = None
self.response_channel = response_channel
def write_callback(self, msg):
gevent.sleep(0.001)
print 'Entering client write callback with message: ', msg
self.response_msg = self.response_channel.read_blocking(timeout=1)
if self.response_msg:
self.response_msg = self.response_msg[1]
print 'Exiting client write callback with response message: ', \
self.response_msg
cwc = ClientWriteCallback(client_read_channel)
client_write_channel.register_write_callback(cwc.write_callback)
class ServiceReadCallback:
def __init__(self, response_channel):
self.read_msg = None
self.response_channel = response_channel
self.response_msg = 'Ping Response'
def read_callback(self, msg):
print 'Entering service read callback with message: ', msg
self.response_channel.write(self.response_msg)
print 'Exiting service read callback with response message: ', \
self.response_msg
src = ServiceReadCallback(service_write_channel)
service_read_channel.register_read_callback(src.read_callback)
print 'Starting service listener'
service_read_channel.start_listener()
gevent.sleep(0.001)
msg = 'Ping'
print 'Writing client message'
client_write_channel.write(msg)
print 'Wrote client message, waiting for response'
for i in range(0, 1000):
gevent.sleep(0.001)
if cwc.response_msg is not None:
break
service_read_channel.stop_listener()
assert cwc.response_msg == src.response_msg
def create_service_channel(channel_name):
from urb.messaging.channel_factory import ChannelFactory
create_config_manager()
return ChannelFactory.get_instance().create_channel(channel_name)
def test_create_channel():
channel_factory = ChannelFactory.get_instance()
channel = channel_factory.create_channel('test_channel')
assert isinstance(channel, Channel)
def shutdown_callback(self, request=None):
self.logger.info('Shutting down')
self.shutdown = True
self.shutdown_event.set()
#self.demoted_callback()
ChannelFactory.get_instance().refresh_master_message_broker(ttl=0)