webrtc/athena: allow webrtc connection with only data channel (#38177)

* add param that requests keyframe

* set keyframe request to false on cleanup

* timing and errors

* video_enabled option and ability to add video later

* remove loggerd changes from this PR

* video enabled default to on, clean identifier

* request keyframe on start regardless of video enabled or not

* move error handling additions to different PR

system/athena/athenad.py

@@ -580,7 +580,7 @@ def getNetworks():
 
 
 @dispatcher.add_method
-def startStream(sdp: str) -> dict:
+def startStream(sdp: str, video_enabled: bool | None = None) -> dict:
   from openpilot.system.webrtc.webrtcd import StreamRequestBody
   bridge_services_in = []
 
@@ -591,17 +591,20 @@ def startStream(sdp: str) -> dict:
       if CP.notCar:
         bridge_services_in.append("testJoystick")
 
-  body = StreamRequestBody(sdp, "wideRoad", bridge_services_in, ["carState"])
+  t_start = time.monotonic()
+  body = StreamRequestBody(sdp=sdp, initCamera="wideRoad", bridge_services_in=bridge_services_in, bridge_services_out=["carState"], video_enabled=video_enabled)
   try:
-    resp = WEBRTCD_SESS.post(f"http://localhost:{WEBRTCD_PORT}/stream",
-                       json=asdict(body), timeout=10)
+    resp = WEBRTCD_SESS.post(f"http://localhost:{WEBRTCD_PORT}/stream", json=asdict(body), timeout=10)
+    t_end = time.monotonic()
     if not resp.ok:
       try:
         error_body = resp.json()
         raise Exception(error_body.get("message", f"webrtcd returned {resp.status_code}"))
       except ValueError:
         resp.raise_for_status()
-    return resp.json()
+    ret = resp.json()
+    ret["time"] = (t_end - t_start) * 1000
+    return ret
   except requests.ConnectTimeout as e:
     raise Exception("webrtc took too long to respond. is the comma body on?") from e
   except requests.ConnectionError as e:

system/webrtc/device/video.py

@@ -29,7 +29,7 @@ class LiveStreamVideoStreamTrack(TiciVideoStreamTrack):
     "road": "livestreamRoadEncodeData",
   }
 
-  def __init__(self, camera_type: str):
+  def __init__(self, camera_type: str, video_enabled: bool = True):
     dt = DT_DMON if camera_type == "driver" else DT_MDL
     super().__init__(camera_type, dt)
 
@@ -39,6 +39,7 @@ class LiveStreamVideoStreamTrack(TiciVideoStreamTrack):
     self.timing_sei_enabled = False
     self.params = Params()
     self._seen_keyframe = False
+    self.video_enabled = video_enabled
 
   def stop(self) -> None:
     super().stop()
@@ -50,6 +51,11 @@ class LiveStreamVideoStreamTrack(TiciVideoStreamTrack):
   def switch_camera(self, camera_type: str) -> None:
     self._sock = self._make_sock(camera_type)
 
+  def enable(self, enabled: bool):
+    self.video_enabled = enabled
+    if not enabled:
+      self._seen_keyframe = False
+
   def _build_frame_data(self, msg) -> bytes:
     encode_data = getattr(msg, msg.which())
     if not self.timing_sei_enabled:
@@ -68,6 +74,12 @@ class LiveStreamVideoStreamTrack(TiciVideoStreamTrack):
     while True:
       if self.readyState != "live":
         raise MediaStreamError
+
+      # while video is disabled, pause here without returning
+      if not self.video_enabled:
+        await asyncio.sleep(0.005)
+        continue
+
       msg = messaging.recv_one_or_none(self._sock)
       if msg is not None:
         if not self._seen_keyframe and (getattr(msg, msg.which()).idx.flags & V4L2_BUF_FLAG_KEYFRAME):

system/webrtc/tests/test_stream_session.py

@@ -32,12 +32,11 @@ class TestStreamSession:
     expected_json = json.dumps(expected_dict).encode()
 
     channel = mocker.Mock(spec=RTCDataChannel)
-    mocked_submaster = messaging.SubMaster(["customReservedRawData0"])
+    proxy = CerealOutgoingMessageProxy(["customReservedRawData0"])
     def mocked_update(t):
-      mocked_submaster.update_msgs(0, [test_msg])
+      proxy.sm.update_msgs(0, [test_msg])
 
     mocker.patch.object(messaging.SubMaster, "update", side_effect=mocked_update)
-    proxy = CerealOutgoingMessageProxy(mocked_submaster)
     proxy.add_channel(channel)
 
     proxy.update()

system/webrtc/webrtcd.py

@@ -49,7 +49,7 @@ def _primary_host_addresses(use_ipv4: bool, use_ipv6: bool) -> list[str]:
   primary = _default_route_ip()
   if primary not in addresses:
     return addresses
-  return [a for a in addresses if a == primary]
+  return [primary, ]
 aioice.ice.get_host_addresses = _primary_host_addresses
 
 
@@ -80,14 +80,19 @@ class AsyncTaskRunner:
 
 
 class CerealOutgoingMessageProxy(AsyncTaskRunner):
-  def __init__(self, sm: messaging.SubMaster):
+  def __init__(self, services: list[str], enabled: bool = True):
     super().__init__()
-    self.sm = sm
+    self.services = list(services)
+    self.sm = messaging.SubMaster(self.services)
     self.channels: list[RTCDataChannel] = []
+    self._enabled = enabled
 
   def add_channel(self, channel: 'RTCDataChannel'):
     self.channels.append(channel)
 
+  def enable(self, enable: bool):
+    self._enabled = enable
+
   def to_json(self, msg_content: Any):
     if isinstance(msg_content, capnp._DynamicStructReader):
       msg_dict = msg_content.to_dict()
@@ -117,6 +122,9 @@ class CerealOutgoingMessageProxy(AsyncTaskRunner):
     from aiortc.exceptions import InvalidStateError
 
     while True:
+      if not self._enabled:
+        await asyncio.sleep(0.01)
+        continue
       try:
         self.update()
       except InvalidStateError:
@@ -165,21 +173,27 @@ class LivestreamBitrateController(AsyncTaskRunner):
   down_samples = 5 # 1s
   param_name = "LivestreamEncoderBitrate"
 
-  def __init__(self, peer_connection: Any):
+  def __init__(self, peer_connection: Any, params: Params, enabled: bool = True):
     super().__init__()
     self.pc = peer_connection
-    self.params = Params()
+    self.params = params
 
-    self.level = 1
+    self.level = 2
     self._publish(self.bitrates[self.level])
     self.prev_lost, self.prev_sent = None, None
     self.counter = 0
     self.up_samples = 5 # 1s
     self._auto = True
+    self._enabled = enabled
+
+  def enable(self, enable: bool):
+    self._enabled = enable
 
   async def run(self):
     while True:
       await asyncio.sleep(self.sample_interval)
+      if not self._enabled:
+        continue
       if not self._auto:
         continue
 
@@ -231,20 +245,28 @@ class LivestreamBitrateController(AsyncTaskRunner):
 class StreamSession:
   shared_pub_master = DynamicPubMaster([])
 
-  def __init__(self, sdp: str, init_camera: str, incoming_services: list[str], outgoing_services: list[str], debug_mode: bool = False):
+  def __init__(
+    self,
+    sdp: str,
+    init_camera: str,
+    incoming_services: list[str],
+    outgoing_services: list[str],
+    enabled: bool | None = None,
+    debug_mode: bool = False
+  ):
     from aiortc.mediastreams import VideoStreamTrack
     from openpilot.system.webrtc.device.video import LiveStreamVideoStreamTrack
     from teleoprtc import WebRTCAnswerBuilder
 
+    self.identifier = str(uuid.uuid4())
+    self.params = Params()
     builder = WebRTCAnswerBuilder(sdp)
 
-    self.video_track = LiveStreamVideoStreamTrack(init_camera) if not debug_mode else VideoStreamTrack()
+    self.enabled = enabled if enabled else True # default to enabled
+    self.video_track = LiveStreamVideoStreamTrack(init_camera, self.enabled) if not debug_mode else VideoStreamTrack()
     builder.add_video_stream(init_camera, self.video_track)
-
     self.stream = builder.stream()
-    self.identifier = str(uuid.uuid4())
 
-    self.params = Params()
     self.incoming_bridge: CerealIncomingMessageProxy | None = None
     self.incoming_bridge_services = incoming_services
     self.outgoing_bridge: CerealOutgoingMessageProxy | None = None
@@ -252,16 +274,16 @@ class StreamSession:
     if len(incoming_services) > 0:
       self.incoming_bridge = CerealIncomingMessageProxy(self.shared_pub_master)
     if len(outgoing_services) > 0:
-      self.outgoing_bridge = CerealOutgoingMessageProxy(messaging.SubMaster(outgoing_services))
-    self.bitrate_controller = LivestreamBitrateController(self.stream.peer_connection)
+      self.outgoing_bridge = CerealOutgoingMessageProxy(outgoing_services, self.enabled)
+    self.bitrate_controller = LivestreamBitrateController(self.stream.peer_connection, self.params, self.enabled)
 
     self.run_task: asyncio.Task | None = None
     self._cleanup_lock = asyncio.Lock()
     self._cleanup_done = False
     self.logger = logging.getLogger("webrtcd")
     self.logger.info(
-      "New stream session (%s), init camera %s, incoming services %s, outgoing services %s",
-      self.identifier, init_camera, incoming_services, outgoing_services,
+      "New stream session (%s), init camera %s, video enabled %s, incoming services %s, outgoing services %s",
+      self.identifier, init_camera, enabled, incoming_services, outgoing_services,
     )
 
   def start(self):
@@ -279,7 +301,6 @@ class StreamSession:
     return await self.stream.start()
 
   def message_handler(self, message: bytes):
-    assert self.incoming_bridge is not None
     try:
       payload = json.loads(message) if isinstance(message, (bytes, str)) else None
       if isinstance(payload, dict):
@@ -290,6 +311,13 @@ class StreamSession:
             self.video_track.switch_camera(payload["data"]["camera"])
           case "livestreamSettings":
             self.bitrate_controller.set_quality(payload["data"]["quality"])
+          case "livestreamVideoEnable":
+            enabled = payload["data"]["enabled"]
+            self.video_track.enable(enabled)
+            self.outgoing_bridge.enable(enabled)
+            self.bitrate_controller.enable(enabled)
+            if not enabled:
+              self.params.put("LivestreamRequestKeyframe", True)
           case "clockSync":
             pong = json.dumps({"type": "clockSync", "data": {
               "action": "pong", "browserSendTime": payload["data"]["browserSendTime"], "deviceTime": time.time() * 1000, # noqa: TID251
@@ -320,9 +348,7 @@ class StreamSession:
       self.bitrate_controller.start()
 
       self.logger.info("Stream session (%s) connected", self.identifier)
-
       await self.stream.wait_for_disconnection()
-
       self.logger.info("Stream session (%s) ended", self.identifier)
     except Exception:
       self.logger.exception("Stream session failure")
@@ -335,12 +361,12 @@ class StreamSession:
         return
       self._cleanup_done = True
       self.params.put("LivestreamRequestKeyframe", False)
-      if self.bitrate_controller is not None:
-        await self.bitrate_controller.stop()
+      await self.bitrate_controller.stop()
       if self.outgoing_bridge is not None:
         await self.outgoing_bridge.stop()
       if self.video_track is not None:
         self.video_track.stop()
+        self.video_track = None
       await self.stream.stop()
 
 
@@ -348,6 +374,7 @@ class StreamSession:
 class StreamRequestBody:
   sdp: str
   initCamera: str
+  video_enabled: bool = True
   bridge_services_in: list[str] = field(default_factory=list)
   bridge_services_out: list[str] = field(default_factory=list)
 
@@ -369,7 +396,8 @@ async def get_stream(request: 'web.Request'):
       await s.stop()
       stream_dict.pop(sid, None)
 
-    session = StreamSession(body.sdp, body.initCamera, body.bridge_services_in, body.bridge_services_out, debug_mode)
+    session = StreamSession(body.sdp, body.initCamera, body.bridge_services_in, body.bridge_services_out, body.video_enabled, debug_mode)
+    stream_dict[session.identifier] = session
     try:
       answer = await session.get_answer()
     except ValueError as e:
@@ -380,8 +408,9 @@ async def get_stream(request: 'web.Request'):
       ) from e
     except Exception:
       await session.stop()
+      stream_dict.pop(session.identifier, None)
+      logging.getLogger("webrtcd").exception("Failed to create stream answer")
       raise
-    stream_dict[session.identifier] = session
     session.start()
 
     def remove_finished_session(_: asyncio.Task) -> None: