async def refresh_cb(kernel_id: str, data: bytes) -> None:
now = await redis_live.time()
await asyncio.shield(call_non_bursty(
conn_tracker_key,
apartial(redis_live.zadd, conn_tracker_key, now, conn_tracker_val),
max_bursts=64, max_idle=2000,
))
async def refresh_cb(kernel_id: str, data: bytes):
await asyncio.shield(
call_non_bursty(kernel_id,
apartial(registry.refresh_session, session_name,
access_key),
max_bursts=64,
max_idle=2000))
def create_app(default_cors_options: CORSOptions) -> Tuple[web.Application, Iterable[WebMiddleware]]:
app = web.Application()
app['api_versions'] = (1, 2, 3, 4)
app.on_startup.append(init)
app.on_shutdown.append(shutdown)
# middleware must be wrapped by web.middleware at the outermost level.
return app, [web.middleware(apartial(rlim_middleware, app))]
def get_loader(self, objtype_name, *args, **kwargs):
k = self._get_key(objtype_name, args, kwargs)
loader = self.cache.get(k)
if loader is None:
objtype_name, has_variant, variant_name = objtype_name.partition(
'.')
objtype = getattr(self.mod, objtype_name)
if has_variant:
batch_load_fn = getattr(objtype, 'batch_load_' + variant_name)
else:
batch_load_fn = objtype.batch_load
loader = DataLoader(apartial(batch_load_fn, *self.common_args,
*args, **kwargs),
max_batch_size=16)
self.cache[k] = loader
return loader
async def stream_proxy(defer, request: web.Request, params: Mapping[str, Any]) -> web.StreamResponse:
registry: AgentRegistry = request.app['registry']
session_name: str = request.match_info['session_name']
access_key: AccessKey = request['keypair']['access_key']
service: str = params['app']
local_config: LocalConfig = request.app['local_config']
myself = asyncio.current_task()
try:
kernel = await asyncio.shield(registry.get_session(session_name, access_key))
except (SessionNotFound, TooManySessionsMatched):
raise
stream_key = kernel['id']
stream_id = uuid.uuid4().hex
request.app['stream_proxy_handlers'][stream_key].add(myself)
defer(lambda: request.app['stream_proxy_handlers'][stream_key].discard(myself))
if kernel['kernel_host'] is None:
kernel_host = urlparse(kernel['agent_addr']).hostname
else:
kernel_host = kernel['kernel_host']
for sport in kernel['service_ports']:
if sport['name'] == service:
if params['port']:
# using one of the primary/secondary ports of the app
try:
hport_idx = sport['container_ports'].index(params['port'])
except ValueError:
raise InvalidAPIParameters(
f"Service {service} does not open the port number {params['port']}.")
host_port = sport['host_ports'][hport_idx]
else:
# using the default (primary) port of the app
if 'host_ports' not in sport:
host_port = sport['host_port'] # legacy kernels
else:
host_port = sport['host_ports'][0]
dest = (kernel_host, host_port)
break
else:
raise AppNotFound(f'{session_name}:{service}')
log.info('STREAM_WSPROXY (ak:{}, s:{}): tunneling {}:{} to {}',
access_key, session_name, service, sport['protocol'], '{}:{}'.format(*dest))
if sport['protocol'] == 'tcp':
proxy_cls = TCPProxy
elif sport['protocol'] == 'pty':
raise NotImplementedError
elif sport['protocol'] == 'http':
proxy_cls = TCPProxy
elif sport['protocol'] == 'preopen':
proxy_cls = TCPProxy
else:
raise InvalidAPIParameters(
f"Unsupported service protocol: {sport['protocol']}")
redis_live: aioredis.Redis = request.app['redis_live']
conn_tracker_key = f"session.{kernel['id']}.active_app_connections"
conn_tracker_val = f"{kernel['id']}:{service}:{stream_id}"
async def refresh_cb(kernel_id: str, data: bytes) -> None:
now = await redis_live.time()
await asyncio.shield(call_non_bursty(
conn_tracker_key,
apartial(redis_live.zadd, conn_tracker_key, now, conn_tracker_val),
max_bursts=64, max_idle=2000,
))
down_cb = apartial(refresh_cb, kernel['id'])
up_cb = apartial(refresh_cb, kernel['id'])
ping_cb = apartial(refresh_cb, kernel['id'])
active_session_ids: DefaultDict[str, int] = request.app['active_session_ids']
kernel_id = kernel['id']
async def add_conn_track() -> None:
async with request.app['conn_tracker_lock']:
active_session_ids[kernel_id] += 1
now = await redis_live.time()
await redis_live.zadd(conn_tracker_key, now, conn_tracker_val)
for idle_checker in request.app['idle_checkers']:
await idle_checker.update_app_streaming_status(
kernel_id,
AppStreamingStatus.HAS_ACTIVE_CONNECTIONS,
)
async def clear_conn_track() -> None:
async with request.app['conn_tracker_lock']:
active_session_ids[kernel_id] -= 1
if active_session_ids[kernel_id] <= 0:
del active_session_ids[kernel_id]
await redis_live.zrem(conn_tracker_key, conn_tracker_val)
remaining_count = await redis_live.zcount(conn_tracker_key)
if remaining_count == 0:
for idle_checker in request.app['idle_checkers']:
await idle_checker.update_app_streaming_status(
kernel_id,
AppStreamingStatus.NO_ACTIVE_CONNECTIONS,
)
try:
await asyncio.shield(add_conn_track())
await asyncio.shield(registry.increment_session_usage(session_name, access_key))
opts: MutableMapping[str, Union[None, str, List[str]]] = {}
if params['arguments'] is not None:
opts['arguments'] = json.loads(params['arguments'])
if params['envs'] is not None:
opts['envs'] = json.loads(params['envs'])
result = await asyncio.shield(
registry.start_service(session_name, access_key, service, opts))
if result['status'] == 'failed':
raise InternalServerError(
"Failed to launch the app service",
extra_data=result['error'])
# TODO: weakref to proxies for graceful shutdown?
ws = web.WebSocketResponse(
autoping=False,
max_msg_size=local_config['manager']['max-wsmsg-size'],
)
await ws.prepare(request)
proxy = proxy_cls(
ws, dest[0], dest[1],
downstream_callback=down_cb,
upstream_callback=up_cb,
ping_callback=ping_cb,
)
return await proxy.proxy()
except asyncio.CancelledError:
log.debug('stream_proxy({}, {}) cancelled', stream_key, service)
raise
finally:
await asyncio.shield(clear_conn_track())
async def stream_proxy(defer, request: web.Request,
params: Mapping[str, Any]) -> web.StreamResponse:
registry = request.app['registry']
session_name = request.match_info['session_name']
access_key = request['keypair']['access_key']
service = params['app']
stream_key = (session_name, access_key)
myself = asyncio.Task.current_task()
request.app['stream_proxy_handlers'][stream_key].add(myself)
defer(lambda: request.app['stream_proxy_handlers'][stream_key].discard(
myself))
try:
kernel = await asyncio.shield(
registry.get_session(session_name, access_key))
except SessionNotFound:
raise
if kernel.kernel_host is None:
kernel_host = urlparse(kernel.agent_addr).hostname
else:
kernel_host = kernel.kernel_host
for sport in kernel.service_ports:
if sport['name'] == service:
if params['port']:
# using one of the primary/secondary ports of the app
try:
hport_idx = sport['container_ports'].index(params['port'])
except ValueError:
raise InvalidAPIParameters(
f"Service {service} does not open the port number {params['port']}."
)
host_port = sport['host_ports'][hport_idx]
else:
# using the default (primary) port of the app
if 'host_ports' not in sport:
host_port = sport['host_port'] # legacy kernels
else:
host_port = sport['host_ports'][0]
dest = (kernel_host, host_port)
break
else:
raise AppNotFound(f'{session_name}:{service}')
log.info('STREAM_WSPROXY (ak:{}, s:{}): tunneling {}:{} to {}', access_key,
session_name, service, sport['protocol'], '{}:{}'.format(*dest))
if sport['protocol'] == 'tcp':
proxy_cls = TCPProxy
elif sport['protocol'] == 'pty':
raise NotImplementedError
elif sport['protocol'] == 'http':
proxy_cls = TCPProxy
elif sport['protocol'] == 'preopen':
proxy_cls = TCPProxy
else:
raise InvalidAPIParameters(
f"Unsupported service protocol: {sport['protocol']}")
# TODO: apply a (distributed) semaphore to limit concurrency per user.
await asyncio.shield(
registry.increment_session_usage(session_name, access_key))
async def refresh_cb(kernel_id: str, data: bytes):
await asyncio.shield(
call_non_bursty(kernel_id,
apartial(registry.refresh_session, session_name,
access_key),
max_bursts=64,
max_idle=2000))
down_cb = apartial(refresh_cb, kernel.id)
up_cb = apartial(refresh_cb, kernel.id)
ping_cb = apartial(refresh_cb, kernel.id)
try:
opts: MutableMapping[str, Union[None, str, List[str]]] = {}
if params['arguments'] is not None:
opts['arguments'] = json.loads(params['arguments'])
if params['envs'] is not None:
opts['envs'] = json.loads(params['envs'])
result = await asyncio.shield(
registry.start_service(session_name, access_key, service, opts))
if result['status'] == 'failed':
raise InternalServerError("Failed to launch the app service",
extra_data=result['error'])
# TODO: weakref to proxies for graceful shutdown?
ws = web.WebSocketResponse(autoping=False)
await ws.prepare(request)
proxy = proxy_cls(ws,
dest[0],
dest[1],
downstream_callback=down_cb,
upstream_callback=up_cb,
ping_callback=ping_cb)
return await proxy.proxy()
except asyncio.CancelledError:
log.debug('stream_proxy({}, {}) cancelled', stream_key, service)
raise
async def stream_proxy(request) -> web.StreamResponse:
registry = request.app['registry']
sess_id = request.match_info['sess_id']
access_key = request['keypair']['access_key']
service = request.query.get('app', None) # noqa
stream_key = (sess_id, access_key)
myself = asyncio.Task.current_task()
request.app['stream_proxy_handlers'][stream_key].add(myself)
try:
kernel = await asyncio.shield(registry.get_session(
sess_id, access_key))
except KernelNotFound:
raise
if kernel.kernel_host is None:
kernel_host = urlparse(kernel.agent_addr).hostname
else:
kernel_host = kernel.kernel_host
for sport in kernel.service_ports:
if sport['name'] == service:
dest = (kernel_host, sport['host_port'])
break
else:
raise AppNotFound(f'{sess_id}:{service}')
log.info('STREAM_WSPROXY: {0} ==[{1}:{2}]==> {3}', sess_id, service,
sport['protocol'], dest)
if sport['protocol'] == 'tcp':
proxy_cls = TCPProxy
elif sport['protocol'] == 'pty':
raise NotImplementedError
# proxy_cls = TermProxy
elif sport['protocol'] == 'http':
proxy_cls = TCPProxy
# proxy_cls = HTTPProxy
else:
raise InvalidAPIParameters(
f"Unsupported service protocol: {sport['protocol']}")
# TODO: apply a (distributed) semaphore to limit concurrency per user.
await asyncio.shield(registry.increment_session_usage(sess_id, access_key))
async def refresh_cb(kernel_id: str, data: bytes):
await asyncio.shield(
call_non_bursty(kernel_id,
apartial(registry.refresh_session, sess_id,
access_key),
max_bursts=64,
max_idle=2000))
down_cb = apartial(refresh_cb, kernel.id)
up_cb = apartial(refresh_cb, kernel.id)
ping_cb = apartial(refresh_cb, kernel.id)
try:
opts: Mapping[str, Any] = {}
result = await asyncio.shield(
registry.start_service(sess_id, access_key, service, opts))
if result['status'] == 'failed':
msg = f"Failed to launch the app service: {result['error']}"
raise InternalServerError(msg)
# TODO: weakref to proxies for graceful shutdown?
ws = web.WebSocketResponse(autoping=False)
await ws.prepare(request)
proxy = proxy_cls(ws,
dest[0],
dest[1],
downstream_callback=down_cb,
upstream_callback=up_cb,
ping_callback=ping_cb)
return await proxy.proxy()
except asyncio.CancelledError:
log.debug('stream_proxy({}, {}) cancelled', stream_key, service)
raise
finally:
request.app['stream_proxy_handlers'][stream_key].discard(myself)