introduce websocket store, to kill open connections on session destroy

facefusion/apis/endpoints/metrics.py

@@ -5,6 +5,7 @@ from starlette.responses import JSONResponse, Response
 from starlette.status import HTTP_404_NOT_FOUND
 from starlette.websockets import WebSocket
 
+from facefusion.apis import websocket_store
 from facefusion.apis.api_helper import get_sec_websocket_protocol
 from facefusion.system import get_metrics_set
 
@@ -21,6 +22,7 @@ async def get_metrics(request : Request) -> Response:
 async def websocket_metrics(websocket : WebSocket) -> None:
 	subprotocol = get_sec_websocket_protocol(websocket.scope)
 	await websocket.accept(subprotocol = subprotocol)
+	websocket_store.set_websocket(websocket)
 
 	try:
 		while True:
@@ -30,3 +32,5 @@ async def websocket_metrics(websocket : WebSocket) -> None:
 
 	except Exception:
 		pass
+
+	websocket_store.delete_websocket(websocket)

facefusion/apis/endpoints/ping.py

@@ -1,5 +1,6 @@
 from starlette.websockets import WebSocket
 
+from facefusion.apis import websocket_store
 from facefusion.apis.api_helper import get_sec_websocket_protocol
 
 
@@ -7,6 +8,7 @@ async def websocket_ping(websocket : WebSocket) -> None:
 	subprotocol = get_sec_websocket_protocol(websocket.scope)
 
 	await websocket.accept(subprotocol = subprotocol)
+	websocket_store.set_websocket(websocket)
 
 	try:
 		while True:
@@ -14,3 +16,5 @@ async def websocket_ping(websocket : WebSocket) -> None:
 
 	except Exception:
 		pass
+
+	websocket_store.delete_websocket(websocket)

facefusion/apis/endpoints/session.py

@@ -1,11 +1,12 @@
 import secrets
 
+from starlette.concurrency import run_in_threadpool
 from starlette.requests import Request
 from starlette.responses import JSONResponse
 from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_401_UNAUTHORIZED, HTTP_404_NOT_FOUND, HTTP_429_TOO_MANY_REQUESTS
 
 from facefusion import content_store, face_store, inference_manager, process_manager, rtc_store, session_context, session_manager, state_manager, store_manager, translator, video_manager
-from facefusion.apis import asset_store
+from facefusion.apis import asset_store, websocket_store
 from facefusion.apis.session_helper import validate_api_key
 from facefusion.filesystem import is_directory, remove_directory
 from facefusion.jobs import job_manager
@@ -31,6 +32,7 @@ async def create_session(request : Request) -> JSONResponse:
 			video_manager.init()
 			process_manager.init()
 			rtc_store.init()
+			websocket_store.init()
 
 			store_manager.observe_api_session(session_id)
 
@@ -105,7 +107,7 @@ async def destroy_session(request : Request) -> JSONResponse:
 		}, status_code = HTTP_404_NOT_FOUND)
 
 	session_manager.clear_api_session(session_id)
-	store_manager.destroy(session_id)
+	await run_in_threadpool(store_manager.destroy, session_id)
 
 	return JSONResponse(
 	{

facefusion/apis/endpoints/stream.py

@@ -4,6 +4,7 @@ from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_404_NOT_FOUND,
 from starlette.websockets import WebSocket, WebSocketState
 
 from facefusion import rtc_store
+from facefusion.apis import websocket_store
 from facefusion.apis.api_helper import get_sec_websocket_protocol
 from facefusion.apis.stream_manager import destroy_stream, process_image, process_video
 
@@ -12,7 +13,9 @@ async def websocket_stream(websocket : WebSocket) -> None:
 	subprotocol = get_sec_websocket_protocol(websocket.scope)
 
 	await websocket.accept(subprotocol = subprotocol)
+	websocket_store.set_websocket(websocket)
 	await process_image(websocket)
+	websocket_store.delete_websocket(websocket)
 
 	if websocket.client_state == WebSocketState.CONNECTED:
 		await websocket.close()

facefusion/apis/websocket_store.py

@@ -0,0 +1,41 @@
+from anyio.from_thread import run
+from anyio.lowlevel import current_token
+
+from facefusion import store_creator
+from facefusion.session_manager import resolve_owner_id
+from facefusion.types import SessionId, SessionWebsocketSet, Store, Websocket
+
+WEBSOCKET_STORE : Store = store_creator.create_store({})
+
+
+def init() -> None:
+	owner_id = resolve_owner_id()
+	store_creator.init_content(WEBSOCKET_STORE, owner_id)
+
+
+def set_websocket(websocket : Websocket) -> None:
+	owner_id = resolve_owner_id()
+
+	if store_creator.has_content(WEBSOCKET_STORE, owner_id):
+		store_creator.get_content(WEBSOCKET_STORE, owner_id)[id(websocket)] =\
+		{
+			'websocket': websocket,
+			'event_loop_token': current_token()
+		}
+
+
+def delete_websocket(websocket : Websocket) -> None:
+	owner_id = resolve_owner_id()
+	session_websocket_set : SessionWebsocketSet = store_creator.get_content(WEBSOCKET_STORE, owner_id)
+
+	if session_websocket_set and id(websocket) in session_websocket_set:
+		del session_websocket_set[id(websocket)]
+
+
+def destroy(session_id : SessionId) -> None:
+	session_websocket_set : SessionWebsocketSet = store_creator.get_content(WEBSOCKET_STORE, session_id)
+	store_creator.delete_content(WEBSOCKET_STORE, session_id)
+
+	if session_websocket_set:
+		for session_websocket in session_websocket_set.values():
+			run(session_websocket.get('websocket').close, token = session_websocket.get('event_loop_token'))

facefusion/store_manager.py

@@ -5,12 +5,12 @@ from types import ModuleType
 from typing import List
 
 from facefusion import content_store, face_store, inference_manager, process_manager, rtc_store, session_manager, state_manager, video_manager
-from facefusion.apis import asset_store
+from facefusion.apis import asset_store, websocket_store
 from facefusion.types import SessionId
 
 
 def get_stores() -> List[ModuleType]:
-	return [ state_manager, asset_store, content_store, face_store, inference_manager, process_manager, rtc_store, video_manager ]
+	return [ state_manager, asset_store, content_store, face_store, inference_manager, process_manager, rtc_store, video_manager, websocket_store ]
 
 
 def observe_api_session(session_id : SessionId) -> None:

facefusion/types.py

@@ -156,6 +156,15 @@ Content : TypeAlias = Dict[str, Any]
 
 Token : TypeAlias = str
 SessionId : TypeAlias = str
+Websocket : TypeAlias = Any
+WebsocketId : TypeAlias = int
+EventLoopToken : TypeAlias = Any
+SessionWebsocket = TypedDict('SessionWebsocket',
+{
+	'websocket' : Websocket,
+	'event_loop_token' : EventLoopToken
+})
+SessionWebsocketSet : TypeAlias = Dict[WebsocketId, SessionWebsocket]
 ApiSession = TypedDict('ApiSession',
 {
 	'access_token' : Token,

tests/test_api_metrics.py

@@ -5,6 +5,7 @@ from pytest_mock import MockerFixture
 from starlette.testclient import TestClient
 
 from facefusion import metadata, session_manager, state_manager
+from facefusion.apis import websocket_store
 from facefusion.apis.core import create_api
 from .assert_helper import get_test_jobs_directory
 
@@ -247,3 +248,7 @@ def test_websocket_metrics(test_client : TestClient) -> None:
 		assert metrics_set.get('processor').get('cores').get('value') == 16
 		assert metrics_set.get('processor').get('frequency').get('value') == 3500
 		assert metrics_set.get('processor').get('utilization').get('value') == 25
+
+		websocket_store.destroy(session_manager.find_api_session_id(create_session_body.get('access_token')))
+
+		assert websocket.receive() == { 'type': 'websocket.close', 'code': 1000, 'reason': '' }

tests/test_api_ping.py

@@ -4,6 +4,7 @@ import pytest
 from starlette.testclient import TestClient
 
 from facefusion import metadata, session_manager, state_manager
+from facefusion.apis import websocket_store
 from facefusion.apis.core import create_api
 from .assert_helper import get_test_jobs_directory
 
@@ -38,3 +39,7 @@ def test_ping(test_client : TestClient) -> None:
 		'access_token.' + create_session_body.get('access_token')
 	]) as websocket:
 		assert websocket
+
+		websocket_store.destroy(session_manager.find_api_session_id(create_session_body.get('access_token')))
+
+		assert websocket.receive() == { 'type': 'websocket.close', 'code': 1000, 'reason': '' }

tests/test_api_session.py

@@ -8,7 +8,7 @@ import pytest
 from starlette.testclient import TestClient
 
 from facefusion import content_store, face_store, inference_manager, metadata, process_manager, rtc_store, session_context, session_manager, state_manager, store_creator, store_manager, video_manager
-from facefusion.apis import asset_store
+from facefusion.apis import asset_store, websocket_store
 from facefusion.apis.core import create_api
 from facefusion.download import conditional_download
 from facefusion.libraries import datachannel as datachannel_module
@@ -300,7 +300,8 @@ def test_destroy_session_content(test_client : TestClient) -> None:
 		(inference_manager, inference_manager.INFERENCE_POOL_STORE),
 		(process_manager, process_manager.PROCESS_STORE),
 		(rtc_store, rtc_store.RTC_STORE),
-		(video_manager, video_manager.VIDEO_POOL_STORE)
+		(video_manager, video_manager.VIDEO_POOL_STORE),
+		(websocket_store, websocket_store.WEBSOCKET_STORE)
 	]
 
 	session_manager.set_api_session(session_id,

tests/test_api_stream.py

@@ -91,7 +91,14 @@ def test_stream_image(test_client : TestClient) -> None:
 		websocket.send_bytes(source_content)
 		output_buffer = websocket.receive_bytes()
 
-	assert create_hash(output_buffer) == '0142782f'
+		assert create_hash(output_buffer) == '0142782f'
+
+		test_client.delete('/session', headers =
+		{
+			'Authorization': 'Bearer ' + access_token
+		})
+
+		assert websocket.receive() == { 'type': 'websocket.close', 'code': 1000, 'reason': '' }
 
 
 @pytest.mark.parametrize('video_codec', [ 'av1', 'vp8' ])

tests/test_store_manager.py

@@ -0,0 +1,48 @@
+from typing import Iterator
+
+import pytest
+
+from facefusion import state_manager, store_creator, video_manager
+from facefusion.download import conditional_download
+from facefusion.session_context import resolve_local_id, set_session_id
+from facefusion.session_manager import clear_api_session, clear_cli_session, create_api_session, fork_session, set_api_session
+from facefusion.store_manager import conditional_cli_destroy
+from .assert_helper import get_test_example_file, get_test_examples_directory
+
+
+@pytest.fixture(scope = 'module', autouse = True)
+def before_all() -> None:
+	conditional_download(get_test_examples_directory(),
+	[
+		'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4'
+	])
+
+	state_manager.init()
+
+
+@pytest.fixture(scope = 'function', autouse = True)
+def before_each() -> Iterator[None]:
+	local_id = resolve_local_id()
+
+	set_session_id(local_id)
+
+	yield
+
+	set_session_id(local_id)
+
+
+def test_conditional_cli_destroy() -> None:
+	set_api_session('session-a', create_api_session())
+	set_session_id('session-a')
+
+	fork_id = fork_session()
+	video_manager.init()
+	video_reader = video_manager.get_reader(get_test_example_file('target-240p.mp4'), 'read_video_frame')
+
+	clear_api_session('session-a')
+	conditional_cli_destroy(fork_id)
+
+	assert video_reader.get('process').poll() == -9
+	assert store_creator.has_content(video_manager.VIDEO_POOL_STORE, fork_id) is False
+
+	clear_cli_session(fork_id)

tests/test_websocket_store.py

@@ -0,0 +1,72 @@
+from typing import Iterator
+from unittest.mock import AsyncMock
+
+import anyio
+import pytest
+
+from facefusion import store_creator
+from facefusion.apis.websocket_store import WEBSOCKET_STORE, delete_websocket, destroy, init, set_websocket
+from facefusion.session_context import resolve_local_id, set_session_id
+from facefusion.types import SessionId, Websocket
+
+
+@pytest.fixture(scope = 'function', autouse = True)
+def before_each() -> Iterator[None]:
+	local_id = resolve_local_id()
+
+	set_session_id(local_id)
+	init()
+
+	yield
+
+	set_session_id(local_id)
+
+
+async def set_and_destroy_websocket(websocket : Websocket, session_id : SessionId) -> None:
+	set_websocket(websocket)
+	await anyio.to_thread.run_sync(destroy, session_id)
+
+
+async def set_and_delete_websocket(websocket : Websocket) -> None:
+	set_websocket(websocket)
+	delete_websocket(websocket)
+
+
+async def set_websocket_in_event_loop(websocket : Websocket) -> None:
+	set_websocket(websocket)
+
+
+def test_init() -> None:
+	set_session_id('session-a')
+	init()
+
+	assert store_creator.get_content(WEBSOCKET_STORE, 'session-a') == {}
+
+	destroy('session-a')
+
+
+def test_set_websocket() -> None:
+	websocket_mock = AsyncMock()
+
+	anyio.run(set_websocket_in_event_loop, websocket_mock)
+
+	assert store_creator.get_content(WEBSOCKET_STORE, resolve_local_id()).get(id(websocket_mock)).get('websocket') is websocket_mock
+
+
+def test_delete_websocket() -> None:
+	websocket_mock = AsyncMock()
+
+	anyio.run(set_and_delete_websocket, websocket_mock)
+
+	assert store_creator.get_content(WEBSOCKET_STORE, resolve_local_id()) == {}
+
+
+def test_destroy() -> None:
+	websocket_mock = AsyncMock()
+
+	set_session_id('session-a')
+	init()
+	anyio.run(set_and_destroy_websocket, websocket_mock, 'session-a')
+
+	websocket_mock.close.assert_awaited_once()
+	assert store_creator.has_content(WEBSOCKET_STORE, 'session-a') is False