Improving the reconnection logic to be able to recreate the peer connection in some cases.

This commit is contained in:
Filipi Fuchter
2025-03-13 17:07:32 -03:00
parent 30432639b4
commit 8f2dadf5a0
4 changed files with 96 additions and 34 deletions

View File

@@ -14,7 +14,7 @@ from typing import Any, Awaitable, Callable, Optional
import cv2
import numpy as np
from aiortc import VideoStreamTrack
from aiortc.mediastreams import AudioStreamTrack, VideoFrame
from aiortc.mediastreams import AudioStreamTrack, MediaStreamError, VideoFrame
from av import AudioFrame, AudioResampler
from loguru import logger
from pydantic import BaseModel
@@ -191,6 +191,9 @@ class SmallWebRTCClient:
logger.warning("Timeout: No video frame received within the specified time.")
self._webrtcConnection.ask_to_renegotiate()
frame = None
except MediaStreamError:
logger.warning("Received an unexpected media stream error while reading the audio.")
frame = None
if frame is None or not isinstance(frame, VideoFrame):
# If no valid frame, sleep for a bit
@@ -237,6 +240,9 @@ class SmallWebRTCClient:
if self._webrtcConnection.is_connected():
logger.warning("Timeout: No audio frame received within the specified time.")
frame = None
except MediaStreamError:
logger.warning("Received an unexpected media stream error while reading the audio.")
frame = None
if frame is None or not isinstance(frame, AudioFrame):
# If we don't read any audio let's sleep for a little bit (i.e. busy wait).
@@ -265,11 +271,11 @@ class SmallWebRTCClient:
yield audio_frame
async def write_raw_audio_frames(self, data: bytes):
if self._can_send():
if self._can_send() and self._audio_output_track:
await self._audio_output_track.add_audio_bytes(data)
async def write_frame_to_camera(self, frame: OutputImageRawFrame):
if self._can_send():
if self._can_send() and self._video_output_track:
self._video_output_track.add_video_frame(frame)
async def setup(self, _params: TransportParams, frame):
@@ -279,23 +285,13 @@ class SmallWebRTCClient:
self._params = _params
async def connect(self):
if self._audio_output_track or self._video_output_track:
if self._webrtcConnection.is_connected() or self._webrtcConnection.is_connecting():
# already initialized
return
logger.info(f"Connecting to Small WebRTC")
await self._webrtcConnection.connect()
logger.info(f"Connecting to Small WebRTC")
if self._params.audio_out_enabled:
self._audio_output_track = RawAudioTrack(sample_rate=self._out_sample_rate)
self._webrtcConnection.replace_audio_track(self._audio_output_track)
if self._params.camera_out_enabled:
self._video_output_track = RawVideoTrack(
width=self._params.camera_out_width, height=self._params.camera_out_height
)
self._webrtcConnection.replace_video_track(self._video_output_track)
async def disconnect(self):
if self.is_connected and not self.is_closing:
@@ -311,12 +307,30 @@ class SmallWebRTCClient:
async def _handle_client_connected(self):
self._audio_input_track = self._webrtcConnection.audio_input_track()
self._video_input_track = self._webrtcConnection.video_input_track()
if self._params.audio_out_enabled:
self._audio_output_track = RawAudioTrack(sample_rate=self._out_sample_rate)
self._webrtcConnection.replace_audio_track(self._audio_output_track)
if self._params.camera_out_enabled:
self._video_output_track = RawVideoTrack(
width=self._params.camera_out_width, height=self._params.camera_out_height
)
self._webrtcConnection.replace_video_track(self._video_output_track)
await self._callbacks.on_client_connected(self._webrtcConnection)
async def _handle_client_disconnected(self):
self._audio_input_track = None
self._video_input_track = None
self._audio_output_track = None
self._video_output_track = None
await self._callbacks.on_client_disconnected(self._webrtcConnection)
async def _handle_client_closed(self):
self._audio_input_track = None
self._video_input_track = None
self._audio_output_track = None
self._video_output_track = None
await self._callbacks.on_client_closed(self._webrtcConnection)
async def _handle_app_message(self, message: Any):

View File

@@ -1,7 +1,7 @@
import asyncio
import json
import uuid
import time
import uuid
from enum import Enum
from typing import Any, Optional
@@ -20,6 +20,11 @@ class SignallingMessage(Enum):
class SmallWebRTCConnection(EventEmitter):
def __init__(self):
super().__init__()
self._is_connecting = False
self._initialize()
def _initialize(self):
logger.info("Initializing new peer connection")
self.answer: Optional[RTCSessionDescription] = None
self.pc = RTCPeerConnection()
self.pc_id = "PeerConnection(%s)" % uuid.uuid4()
@@ -87,10 +92,21 @@ class SmallWebRTCConnection(EventEmitter):
await self._create_answer(sdp, type)
async def connect(self):
self._is_connecting = True
await self.pc.setLocalDescription(self.answer)
async def renegotiate(self, sdp: str, type: str):
async def renegotiate(self, sdp: str, type: str, restart_pc: bool = False):
logger.info(f"Renegotiating {self.pc_id}")
if restart_pc:
await self.emit("disconnected", self)
logger.info("Closing old peer connection")
# removing the listeners to prevent the bot from closing
self.pc.remove_all_listeners()
await self.close()
# we are initializing a new peer connection in this case.
self._initialize()
await self._create_answer(sdp, type)
await self.pc.setLocalDescription(self.answer)
@@ -136,6 +152,7 @@ class SmallWebRTCConnection(EventEmitter):
async def close(self):
if self.pc:
await self.pc.close()
self._is_connecting = False
def get_answer(self):
if not self.answer:
@@ -166,6 +183,9 @@ class SmallWebRTCConnection(EventEmitter):
# Checks if the last received ping was within the last 3 seconds.
return (time.time() - self._last_received_time) < 3
def is_connecting(self):
return self._is_connecting
def audio_input_track(self):
# Transceivers always appear in creation-order for both peers
# For now we are only considering that we are going to have 02 transceivers,