Don't accept streams until server is ready to read them

.gitignore

@@ -177,3 +177,4 @@ $RECYCLE.BIN/
 *.lnk
 
 # End of https://www.gitignore.io/api/c,vim,linux,macos,elixir,windows,visualstudiocode
+received.flv

examples/source_with_standalone_server.exs

@@ -45,25 +45,34 @@ defmodule Pipeline do
   end
 end
 
+Logger.configure(level: :error)
+
 # The client will connect on `rtmp://localhost:1935/app/stream_key`
 port = 1935
 app = "app"
 stream_key = "stream_key"
 
+# example lambda function that upon launching will send a message back to parent server.
+lambda = fn server, app, stream_key ->
+  send(server, {:lambda, "hello from the other side #{app} #{stream_key}"})
+end
+
 # Run the standalone server
 {:ok, server} =
   Membrane.RTMP.Server.start_link(
     handler: %Membrane.RTMP.Source.ClientHandler{controlling_process: self()},
     port: port,
-    use_ssl?: false
+    use_ssl?: false,
+    lambda: lambda
   )
 
-# Subscribe to receive client reference that connected to the
-# server with given app id and stream key
-:ok = Membrane.RTMP.Server.subscribe(server, app, stream_key)
+# Subscribe any app stream key to be informed when new client connects,
+# so you can await its ref and pull data from it.
+:ok = Membrane.RTMP.Server.subscribe_any(server)
+
+# Wait for the client reference for given app stream_key
+{:ok, client_ref} = Membrane.RTMP.Server.await_client_ref(app, stream_key)
 
-# Wait for the client reference
-{:ok, client_ref} = Membrane.RTMP.Server.await_subscription(app, stream_key)
 # Start the pipeline and provide it with the client_ref
 {:ok, _supervisor, pipeline} =
   Membrane.Pipeline.start_link(Pipeline, client_ref: client_ref)

lib/membrane_rtmp_plugin/rtmp/source/message_handler.ex

@@ -48,6 +48,9 @@ defmodule Membrane.RTMP.MessageHandler do
       header_sent?: false,
       events: [],
       receiver_pid: nil,
+      #
+      publish_msg: nil,
+      publish_header: nil,
       # how many times the Source tries to get control of the socket
       socket_retries: 3,
       # epoch required for performing a handshake with the pipeline
@@ -72,6 +75,17 @@ defmodule Membrane.RTMP.MessageHandler do
     end)
   end
 
+  @spec send_publish_success(map()) :: {map(), list()}
+  def send_publish_success(state) do
+    Responses.publish_success(state.publish_msg.stream_key)
+    |> send_rtmp_payload(state.socket,
+      chunk_stream_id: 3,
+      stream_id: state.publish_header.stream_id
+    )
+
+    {%{state | events: []}, [{:published, state.publish_msg} | state.events]}
+  end
+
   # Expected flow of messages:
   # 1.  [in] c0_c1 handshake -> [out] s0_s1_s2 handshake
   # 2.  [in] c2 handshake -> [out] empty
@@ -141,11 +155,10 @@ defmodule Membrane.RTMP.MessageHandler do
     %Messages.UserControl{event_type: @stream_begin_type, data: <<0, 0, 0, 1>>}
     |> send_rtmp_payload(state.socket, chunk_stream_id: 2)
 
-    Responses.publish_success(publish_msg.stream_key)
-    |> send_rtmp_payload(state.socket, chunk_stream_id: 3, stream_id: header.stream_id)
+    # at this point pause the unfinished handshake until pipeline demands data from this client
+    # (this mechanism prevents accepting streams with no listeners)
 
-    state = %{state | events: [{:published, publish_msg} | state.events]}
-    {:cont, state}
+    {:halt, %{state | publish_msg: publish_msg, publish_header: header}}
   end
 
   # A message containing stream metadata

lib/membrane_rtmp_plugin/rtmp/source/source.ex

@@ -85,7 +85,8 @@ defmodule Membrane.RTMP.Source do
       Membrane.RTMP.Server.start_link(
         handler: %SourceClientHandler{controlling_process: self()},
         port: port,
-        use_ssl?: use_ssl?
+        use_ssl?: use_ssl?,
+        lambda: nil
       )
 
     state = %{state | app: app, stream_key: stream_key, server: server_pid}
@@ -157,7 +158,7 @@ defmodule Membrane.RTMP.Source do
 
   @impl true
   def handle_info({:client_connected, app, stream_key}, _ctx, %{mode: :builtin_server} = state) do
-    :ok = Membrane.RTMP.Server.subscribe(state.server, state.app, state.stream_key)
+    :ok = Membrane.RTMP.Server.subscribe_any(state.server)
     state = %{state | app: app, stream_key: stream_key}
     {[], state}
   end

lib/membrane_rtmp_plugin/rtmp_server.ex

@@ -1,6 +1,12 @@
 defmodule Membrane.RTMP.Server do
   @moduledoc """
   A simple RTMP server, which handles each new incoming connection.
+
+  When new client connects to the server, it goes into :client_waiting_queue and its RTMP handshake will remanin unfinished.
+  Only when pipeline tries to pull data from client, its handshake will be finished, and client will be registered.
+
+  Also when new client connects, optional, annonymous function defined by user is triggered.
+  The lambda function is given PID of parent server, app and stream key.
   """
   use GenServer
 
@@ -15,7 +21,8 @@ defmodule Membrane.RTMP.Server do
           handler: ClientHandlerBehaviour.t(),
           port: :inet.port_number(),
           use_ssl?: boolean(),
-          name: atom() | nil
+          name: atom() | nil,
+          lambda: function() | nil
         ]
 
   @type server_identifier :: pid() | atom()
@@ -31,27 +38,26 @@ defmodule Membrane.RTMP.Server do
   end
 
   @doc """
-  Subscribes for the given app and stream key.
-  When a client connects (or has already connected) to the server with given app and stream key,
-  the subscriber will be informed.
+  Subscribes for any app/stream_key.
+  When a new client connects, subscriber will be informed, if currently is awaiting client_ref for a given app/stream_key
   """
-  @spec subscribe(server_identifier(), String.t(), String.t()) :: :ok
-  def subscribe(server_identifier, app, stream_key) do
-    GenServer.cast(server_identifier, {:subscribe, app, stream_key, self()})
+  @spec subscribe_any(server_identifier()) :: :ok
+  def subscribe_any(server_identifier) do
+    GenServer.cast(server_identifier, {:subscribe_any, self()})
     :ok
   end
 
   @doc """
-  Awaits for the client reference of the connection to which the user has previously subscribed.
+  Awaits for the client reference of the connection to specified app/stream_key
 
   Note: this function call is blocking!
-  Note: first you need to call `#{__MODULE__}.subscribe/3` to subscribe
-  for a given `app` and `stream_key`.
+  Note: first you need to call `#{__MODULE__}.subscribe_any/1`.
   """
-  @spec await_subscription(String.t(), String.t(), non_neg_integer()) :: {:ok, pid()} | :error
-  def await_subscription(app, stream_key, timeout \\ 5_000) do
+  @spec await_client_ref(String.t(), String.t(), non_neg_integer()) :: {:ok, pid()} | :error
+  def await_client_ref(app, stream_key, timeout \\ 5_000) do
     receive do
-      {:client_ref, client_ref, ^app, ^stream_key} -> {:ok, client_ref}
+      {:client_ref, client_ref, ^app, ^stream_key} ->
+        {:ok, client_ref}
     after
       timeout -> :error
     end
@@ -74,12 +80,14 @@ defmodule Membrane.RTMP.Server do
 
     {:ok,
      %{
-       subscriptions: %{},
+       subscriptions_any: [],
        client_reference_mapping: %{},
+       client_waiting_queue: %{},
        listener: pid,
        port: nil,
        to_reply: [],
-       use_ssl?: server_options.use_ssl?
+       use_ssl?: server_options.use_ssl?,
+       lambda: server_options.lambda
      }}
   end
 
@@ -93,16 +101,37 @@ defmodule Membrane.RTMP.Server do
   end
 
   @impl true
-  def handle_cast({:subscribe, app, stream_key, subscriber_pid}, state) do
-    state = put_in(state, [:subscriptions, {app, stream_key}], subscriber_pid)
-    maybe_send_subscription(app, stream_key, state)
+  def handle_cast({:subscribe_any, subscriber_pid}, state) do
+    subs = [state.subscriptions_any ++ subscriber_pid]
+    state = %{state | subscriptions_any: subs}
+
+    # try to send all client_refs from :client_waiting_queue to this subscriber, maybe is already awaiting one of /app/stream_keys
+    state.client_waiting_queue
+    |> Enum.each(fn {{app, stream_key}, client_ref} ->
+      send(subscriber_pid, {:client_ref, client_ref, app, stream_key})
+    end)
+
     {:noreply, state}
   end
 
   @impl true
   def handle_info({:register_client, app, stream_key, client_reference_pid}, state) do
     state = put_in(state, [:client_reference_mapping, {app, stream_key}], client_reference_pid)
-    maybe_send_subscription(app, stream_key, state)
+    client_waiting_queue = state.client_waiting_queue |> Map.delete({app, stream_key})
+    {:noreply, %{state | client_waiting_queue: client_waiting_queue}}
+  end
+
+  @impl true
+  def handle_info({:register_client_in_queue, app, stream_key, client_ref}, state) do
+    if is_function(state.lambda) do
+      state.lambda.(self(), app, stream_key)
+    end
+
+    state = put_in(state, [:client_waiting_queue, {app, stream_key}], client_ref)
+    # send client ref_to anyone possibly awaiting it
+    state.subscriptions_any
+    |> Enum.each(fn subscriber -> send(subscriber, {:client_ref, client_ref, app, stream_key}) end)
+
     {:noreply, state}
   end
 
@@ -112,13 +141,9 @@ defmodule Membrane.RTMP.Server do
     {:noreply, %{state | port: port, to_reply: []}}
   end
 
-  defp maybe_send_subscription(app, stream_key, state) do
-    if state.subscriptions[{app, stream_key}] != nil and
-         state.client_reference_mapping[{app, stream_key}] != nil do
-      send(
-        state.subscriptions[{app, stream_key}],
-        {:client_ref, state.client_reference_mapping[{app, stream_key}], app, stream_key}
-      )
-    end
+  @impl true
+  def handle_info({:lambda, message}, state) do
+    IO.inspect(message, label: "new message from lambda")
+    {:noreply, state}
   end
 end

lib/membrane_rtmp_plugin/rtmp_server/client_handler.ex

@@ -38,7 +38,8 @@ defmodule Membrane.RTMP.Server.ClientHandler do
        stream_key: nil,
        server: opts.server,
        buffers_demanded: 0,
-       published?: false
+       published?: false,
+       client_register_attempt_made?: false
      }}
   end
 
@@ -77,7 +78,7 @@ defmodule Membrane.RTMP.Server.ClientHandler do
 
   @impl true
   def handle_info({:demand_data, how_many_buffers_demanded}, state) do
-    state = %{state | buffers_demanded: how_many_buffers_demanded}
+    state = %{finish_handshake(state) | buffers_demanded: how_many_buffers_demanded}
     request_data(state)
     {:noreply, state}
   end
@@ -96,6 +97,17 @@ defmodule Membrane.RTMP.Server.ClientHandler do
     {message_handler_state, events} =
       MessageHandler.handle_client_messages(messages, state.message_handler_state)
 
+    state =
+      if message_handler_state.publish_msg != nil and not state.client_register_attempt_made? do
+        %{publish_msg: %Membrane.RTMP.Messages.Publish{stream_key: stream_key}} =
+          message_handler_state
+
+        send(state.server, {:register_client_in_queue, state.app, stream_key, self()})
+        %{state | client_register_attempt_made?: true}
+      else
+        state
+      end
+
     state = Enum.reduce(events, state, &handle_event/2)
 
     request_data(state)
@@ -159,4 +171,14 @@ defmodule Membrane.RTMP.Server.ClientHandler do
       end
     end
   end
+
+  defp finish_handshake(state) when not state.published? do
+    {message_handler_state, events} =
+      MessageHandler.send_publish_success(state.message_handler_state)
+
+    state = Enum.reduce(events, state, &handle_event/2)
+    %{state | message_handler_state: message_handler_state}
+  end
+
+  defp finish_handshake(state), do: state
 end

test/membrane_rtmp_plugin/rtmp_source_bin_test.exs

@@ -200,7 +200,7 @@ defmodule Membrane.RTMP.SourceBin.IntegrationTest do
   defp start_pipeline_with_builtin_rtmp_server(app, stream_key, use_ssl? \\ false) do
     options = [
       module: Membrane.RTMP.Source.WithBuiltinServerTestPipeline,
-      custom_args: %{app: app, stream_key: stream_key, port: @default_port, use_ssl?: use_ssl?},
+      custom_args: %{app: app, stream_key: stream_key, port: @default_port, use_ssl?: use_ssl?, lambda: nil},
       test_process: self()
     ]
 
@@ -221,7 +221,8 @@ defmodule Membrane.RTMP.SourceBin.IntegrationTest do
           controlling_process: self()
         },
         port: port,
-        use_ssl?: use_ssl?
+        use_ssl?: use_ssl?,
+        lambda: nil
       )
 
     {:ok, assigned_port} = Membrane.RTMP.Server.get_port(server_pid)
@@ -233,10 +234,10 @@ defmodule Membrane.RTMP.SourceBin.IntegrationTest do
         {:client_connected, ^app, ^stream_key} -> :ok
       end
 
-    :ok = Membrane.RTMP.Server.subscribe(server_pid, app, stream_key)
+    :ok = Membrane.RTMP.Server.subscribe_any(server_pid)
 
     {:ok, client_reference} =
-      Membrane.RTMP.Server.await_subscription(app, stream_key)
+      Membrane.RTMP.Server.await_client_ref(app, stream_key)
 
     options = [
       module: Membrane.RTMP.Source.WithExternalServerTestPipeline,