def __init__(self, environments): self.environments = environments self._connections = {} self.server = EnvironmentsTCPServer() self.server.set_connection_message_callback(self.on_connection_message) self.server.listen(self.server.PORT)
class EnvironmentsController(object): ENVIRONMENT_CLIENT_CLS = EnvironmentClient def __init__(self, environments): self.environments = environments self._connections = {} self.server = EnvironmentsTCPServer() self.server.set_connection_message_callback(self.on_connection_message) self.server.listen(self.server.PORT) def get_environment(self, env_name, on_stdout, on_stderr): executable_path = self.get_executable_path(env_name) return self.start_env(executable_path, on_stdout, on_stderr) def get_executable_path(self, env_name): return self.environments[env_name] def start_env(self, executable, on_stdout, on_stderr): environment_id = uuid.uuid4().hex args = [ executable, str(self.server.PORT), environment_id ] stream = Subprocess.STREAM env = { 'PYTHONUNBUFFERED': '0', 'PYTHONIOENCODING': 'utf-8' } try: sub_process = Subprocess(args=args, executable=executable, stdout=stream, stderr=stream, env=env) except Exception as e: logger.error(e) raise def decode_data(func): def _decode(data): return func(data.decode('utf-8')) return _decode _on_stdout = decode_data(partial(on_stdout, environment_id)) _on_stderr = decode_data(partial(on_stderr, environment_id)) sub_process.stdout.read_until_close(lambda a: a, streaming_callback=_on_stdout) sub_process.stderr.read_until_close(lambda a: a, streaming_callback=_on_stderr) self._connections[environment_id] = Future() return self._connections[environment_id] def on_connection_message(self, data, stream): if data.get('status') != 'connected': raise CheckioEnvironmentError("Wrong connection message {}".format(str(data))) environment_id = data['environment_id'] environment_client = self.ENVIRONMENT_CLIENT_CLS(stream, environment_id) environment_client.set_on_stop_callback(self.on_environment_stopped) logger.debug("EnvironmentsController:: connected {}".format(environment_id)) self._connections[environment_id].set_result(environment_client) def on_environment_stopped(self, environment_id): del self._connections[environment_id] @gen.coroutine def stop_all_environments(self): stop_environments = [] for environment in self._connections.values(): stop_environments.append(environment.stop()) return (yield stop_environments) def is_valid_env(self, env_name): return env_name in self.environments