fork and join sessions (#1239)

* introduce fork and join session, fix inference manager race condition

* destroy stream on DELETE

tests/test_api_session.py

@@ -1,3 +1,4 @@
+import ctypes
 import os
 import tempfile
 from datetime import timedelta
@@ -7,11 +8,12 @@ from unittest.mock import patch
 import pytest
 from starlette.testclient import TestClient
 
-from facefusion import metadata, process_manager, session_manager, state_manager
+from facefusion import metadata, process_manager, rtc, rtc_store, session_manager, state_manager
 from facefusion.apis import asset_store
 from facefusion.apis.core import create_api
 from facefusion.download import conditional_download
-from facefusion.types import Session
+from facefusion.libraries import datachannel as datachannel_module
+from facefusion.types import RtcPeer, Session
 from .assert_helper import get_test_example_file, get_test_examples_directory
 
 
@@ -19,6 +21,8 @@ from .assert_helper import get_test_example_file, get_test_examples_directory
 def before_all() -> None:
 	state_manager.init()
 
+	datachannel_module.pre_check()
+
 	process_manager.start()
 	conditional_download(get_test_examples_directory(),
 	[
@@ -238,6 +242,21 @@ def test_destroy_session(test_client : TestClient) -> None:
 	assert session_manager.find_session_id(access_token) == session_id
 	assert delete_session_response.status_code == 404
 
+	peer_connection = rtc.create_peer_connection()
+	rtc_peer : RtcPeer =\
+	{
+		'peer_connection': peer_connection,
+		'video':
+		{
+			'sender_track': rtc.add_video_track(peer_connection, 'sendonly', 'vp8', 96),
+			'receiver_track': 0,
+			'codec': 'vp8'
+		},
+		'sender_bitrate': ctypes.c_uint(0),
+		'receiver_bitrate': ctypes.c_uint(0)
+	}
+	rtc_store.set_peer(session_id, rtc_peer)
+
 	delete_session_response = test_client.delete('/session', headers =
 	{
 		'Authorization': 'Bearer ' + access_token
@@ -245,6 +264,7 @@ def test_destroy_session(test_client : TestClient) -> None:
 
 	assert session_manager.find_session_id(access_token) is None
 	assert asset_store.get_assets(session_id) is None
+	assert rtc_store.has_peer(session_id) is False
 	assert delete_session_response.status_code == 200
 
 	for asset_path in asset_paths:

tests/test_session_manager.py

@@ -1,7 +1,35 @@
 import secrets
-from datetime import timedelta
+from datetime import datetime, timedelta
+from typing import Iterator
 
-from facefusion.session_manager import clear_session, create_session, get_session, set_session, validate_session
+import pytest
+
+from facefusion.session_context import get_session_id, resolve_local_id, set_session_id
+from facefusion.session_manager import clear_session, create_session, fork_session, get_session, join_session, resolve_owner_id, set_session, validate_session
+
+
+@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_fork_session() -> None:
+	local_id = resolve_local_id()
+	session_id = fork_session()
+
+	assert get_session_id() == session_id
+	assert resolve_owner_id() == local_id
+	assert get_session(session_id).get('owner_id') == local_id
+	assert get_session(session_id).get('access_token') is None
+	assert get_session(session_id).get('expires_at') is None
+
+	join_session()
 
 
 def test_get_and_set_session() -> None:
@@ -13,6 +41,25 @@ def test_get_and_set_session() -> None:
 	assert get_session(session_id) == session
 
 
+def test_resolve_owner_id() -> None:
+	local_id = resolve_local_id()
+
+	assert resolve_owner_id() == local_id
+
+	set_session('session-a', create_session())
+	set_session('session-a1', { 'owner_id': 'session-a', 'created_at': datetime.now() })
+	set_session_id('session-a')
+
+	assert resolve_owner_id() == 'session-a'
+
+	set_session_id('session-a1')
+
+	assert resolve_owner_id() == 'session-a'
+
+	clear_session('session-a1')
+	clear_session('session-a')
+
+
 def test_validate_session() -> None:
 	session = create_session()
 	session_id = secrets.token_urlsafe(16)
@@ -43,3 +90,12 @@ def test_clear_session() -> None:
 	clear_session(session_id)
 
 	assert validate_session(session_id) is None
+
+
+def test_join_session() -> None:
+	local_id = resolve_local_id()
+	session_id = fork_session()
+	join_session()
+
+	assert get_session_id() == local_id
+	assert get_session(session_id) is None

tests/test_state_manager.py

@@ -1,10 +1,11 @@
+from datetime import datetime
 from typing import Iterator
 
 import pytest
 
-from facefusion import store_creator
+from facefusion import session_manager, store_creator
 from facefusion.session_context import resolve_local_id, set_session_id
-from facefusion.state_manager import STATE_SET, clear, get_item, get_state, init, init_item, set_item, set_state
+from facefusion.state_manager import STATE_SET, clear, clone_state, get_item, get_state, init, init_item, set_item, set_state
 
 
 @pytest.fixture(scope = 'function', autouse = True)
@@ -14,6 +15,8 @@ def before_each() -> Iterator[None]:
 	set_session_id(local_id)
 	clear()
 	store_creator.delete_content(STATE_SET, 'session-a')
+	store_creator.delete_content(STATE_SET, 'session-a1')
+	session_manager.clear_session('session-a1')
 
 	yield
 
@@ -48,6 +51,20 @@ def test_set_state() -> None:
 	assert get_state() == { 'video_memory_strategy': 'strict' }
 
 
+def test_clone_state() -> None:
+	init_item('video_memory_strategy', 'tolerant')
+	session_manager.set_session('session-a1', { 'owner_id': resolve_local_id(), 'created_at': datetime.now() })
+	set_session_id('session-a1')
+	clone_state()
+	set_item('video_memory_strategy', 'strict')
+
+	assert get_state() == { 'video_memory_strategy': 'strict' }
+
+	set_session_id(resolve_local_id())
+
+	assert get_state() == { 'video_memory_strategy': 'tolerant' }
+
+
 def test_clear() -> None:
 	init_item('video_memory_strategy', 'tolerant')
 	clear()