Add av1 support

Author: gBillal

lib/shinkai/pipeline.ex

@@ -15,19 +15,19 @@ defmodule Shinkai.Pipeline do
 
   @spec add_rtmp_client(String.t()) :: :ok
   def add_rtmp_client(source_id) do
-    Sink.RTMP.add_client({:via, Registry, {Source.Registry, {:rtmp_sink, source_id}}}, self())
+    Sink.RTMP.add_client(rtmp_name(source_id), self())
   end
 
   def add_webrtc_peer(source_id) do
-    Sink.WebRTC.add_new_peer(:"webrtc_sink_#{source_id}")
+    Sink.WebRTC.add_new_peer(webrtc_name(source_id))
   end
 
   def handle_webrtc_peer_answer(source_id, session_id, sdp_answer) do
-    Sink.WebRTC.handle_peer_answer(:"webrtc_sink_#{source_id}", session_id, sdp_answer)
+    Sink.WebRTC.handle_peer_answer(webrtc_name(source_id), session_id, sdp_answer)
   end
 
   def remove_webrtc_peer(source_id, session_id) do
-    Sink.WebRTC.remove_peer(:"webrtc_sink_#{source_id}", session_id)
+    Sink.WebRTC.remove_peer(webrtc_name(source_id), session_id)
   end
 
   def stop(source_id) do
@@ -42,7 +42,7 @@ defmodule Shinkai.Pipeline do
     children =
       [
         {Sink.Hls, [id: id] ++ hls_config},
-        {Sink.WebRTC, id: id, name: :"webrtc_sink_#{id}"}
+        {Sink.WebRTC, id: id}
       ] ++ rtmp_sink(rtmp_config[:enabled], id) ++ source(source)
 
     Supervisor.init(children, strategy: :one_for_all)
@@ -54,4 +54,7 @@ defmodule Shinkai.Pipeline do
 
   defp rtmp_sink(false, _id), do: []
   defp rtmp_sink(true, id), do: [{Sink.RTMP, [id: id]}]
+
+  defp webrtc_name(id), do: {:via, Registry, {Source.Registry, {:webrtc_sink, id}}}
+  defp rtmp_name(id), do: {:via, Registry, {Source.Registry, {:rtmp_sink, id}}}
 end

lib/shinkai/sink/webrtc.ex

@@ -12,28 +12,29 @@ defmodule Shinkai.Sink.WebRTC do
   alias Phoenix.PubSub
   alias RTSP.RTP.Encoder, as: RTPEncoder
 
-  @supported_codecs [:h264, :h265, :pcma]
+  @supported_codecs [:h264, :h265, :av1, :pcma]
   @video_clock_rate 90_000
 
   def start_link(opts) do
-    GenServer.start_link(__MODULE__, opts, name: opts[:name])
+    id = {:via, Registry, {Source.Registry, {:webrtc_sink, opts[:id]}}}
+    GenServer.start_link(__MODULE__, opts, name: id)
   end
 
-  @spec add_new_peer(server :: pid() | atom()) :: {:ok, String.t(), String.t()} | {:error, any()}
+  @spec add_new_peer(GenServer.server()) :: {:ok, String.t(), String.t()} | {:error, any()}
   def add_new_peer(server) do
     GenServer.call(server, :add_new_peer)
   end
 
   @spec handle_peer_answer(
-          server :: pid() | atom(),
+          GenServer.server(),
           session_id :: String.t(),
           sdp :: String.t()
         ) :: :ok | {:error, any()}
   def handle_peer_answer(server, session_id, sdp) do
     GenServer.call(server, {:handle_peer_answer, session_id, sdp})
   end
 
-  @spec remove_peer(server :: pid() | atom(), session_id :: String.t()) :: :ok
+  @spec remove_peer(GenServer.server(), session_id :: String.t()) :: :ok
   def remove_peer(server, session_id) do
     GenServer.cast(server, {:remove_peer, session_id})
   end
@@ -165,7 +166,7 @@ defmodule Shinkai.Sink.WebRTC do
 
     track_id = track_ctx.id
 
-    Registry.dispatch(Shinkai.Registry, {:webrtc, source_id}, fn peers ->
+    Registry.dispatch(Sink.Registry, {:webrtc, source_id}, fn peers ->
       for {_pid, {pc, _session_id}} <- peers do
         Enum.each(packets, fn rtp_packet ->
           :ok = ExWebRTC.PeerConnection.send_rtp(pc, track_id, rtp_packet)
@@ -196,6 +197,7 @@ defmodule Shinkai.Sink.WebRTC do
 
   defp mime_type(:h264), do: "video/H264"
   defp mime_type(:h265), do: "video/H265"
+  defp mime_type(:av1), do: "video/AV1"
   defp mime_type(:pcma), do: "audio/PCMA"
 
   defp sdp_fmtp_line(:h264, pt) do
@@ -207,10 +209,10 @@ defmodule Shinkai.Sink.WebRTC do
     }
   end
 
-  defp sdp_fmtp_line(:h265, _pt), do: nil
   defp sdp_fmtp_line(_codec, _pt), do: nil
 
   defp payloader_mod(:h264), do: RTPEncoder.H264
   defp payloader_mod(:h265), do: RTPEncoder.H265
+  defp payloader_mod(:av1), do: RTPEncoder.AV1
   defp payloader_mod(:pcma), do: RTPEncoder.G711
 end

lib/shinkai/sink/webrtc/peer_manager.ex

@@ -65,7 +65,7 @@ defmodule Shinkai.Sink.WebRTC.PeerManager do
   def handle_call({:remove_peer, session_id}, _from, state) do
     Logger.info("Removing WebRTC peer with session ID: #{session_id}")
 
-    case Registry.match(Shinkai.Registry, {:webrtc, state.source_id}, {:_, session_id}) do
+    case Registry.match(Sink.Registry, {:webrtc, state.source_id}, {:_, session_id}) do
       [{_pid, {pc, _session_id}}] -> unregister(state.source_id, pc)
       [] -> :ok
     end
@@ -138,8 +138,9 @@ defmodule Shinkai.Sink.WebRTC.PeerManager do
   end
 
   def handle_info({:ex_webrtc, pid, {:connection_state_change, :connected}}, state) do
+    Logger.info("New WebRTC peer connected")
     {session_id, pc} = Enum.find(state.sessions, fn {_session_id, p} -> p == pid end)
-    Registry.register(Shinkai.Registry, {:webrtc, state.source_id}, {pc, session_id})
+    Registry.register(Sink.Registry, {:webrtc, state.source_id}, {pc, session_id})
     {:noreply, %{state | sessions: Map.delete(state.sessions, session_id)}}
   end
 
@@ -157,7 +158,7 @@ defmodule Shinkai.Sink.WebRTC.PeerManager do
     {:noreply, state}
   end
 
-  def handle_info({:ex_webrtc, _pid, msg}, state) do
+  def handle_info({:ex_webrtc, _pid, _msg}, state) do
     # Logger.info("Unhandled ExWebRTC message: #{inspect(msg)}")
     {:noreply, state}
   end
@@ -177,6 +178,6 @@ defmodule Shinkai.Sink.WebRTC.PeerManager do
 
   defp unregister(source_id, pc) do
     PeerConnection.stop(pc)
-    Registry.unregister_match(Shinkai.Registry, {:webrtc, source_id}, {pc, :_})
+    Registry.unregister_match(Sink.Registry, {:webrtc, source_id}, {pc, :_})
   end
 end