Fix memory leak in the sink (#31)

Author: dmorn

c_src/membrane_rtmp_plugin/sink/rtmp_sink.c

@@ -134,24 +134,24 @@ UNIFEX_TERM write_video_frame(UnifexEnv *env, State *state,
       state->output_ctx->streams[state->video_stream_index]->time_base;
   AVPacket *packet = av_packet_alloc();
 
+  uint8_t* data = (uint8_t *)av_malloc(frame->size);
+  memcpy(data, frame->data, frame->size);
+  av_packet_from_data(packet, data, frame->size);
+
   UNIFEX_TERM write_frame_result;
 
   if (is_key_frame) {
     packet->flags |= AV_PKT_FLAG_KEY;
   }
 
   packet->stream_index = state->video_stream_index;
-  packet->size = frame->size;
-  packet->data = (uint8_t *)av_malloc(frame->size);
 
   if (!packet->data) {
     write_frame_result =
         unifex_raise(env, "Failed allocating video frame data");
     goto end;
   }
 
-  memcpy(packet->data, frame->data, frame->size);
-
   int64_t dts_scaled =
       av_rescale_q(dts, MEMBRANE_TIME_BASE, video_stream_time_base);
   int64_t pts_scaled =
@@ -162,15 +162,14 @@ UNIFEX_TERM write_video_frame(UnifexEnv *env, State *state,
   packet->duration = dts_scaled - state->current_video_dts;
   state->current_video_dts = dts_scaled;
 
-  if (av_interleaved_write_frame(state->output_ctx, packet)) {
+  if (av_write_frame(state->output_ctx, packet)) {
     write_frame_result =
         write_video_frame_result_error(env, "Failed writing video frame");
     goto end;
   }
   write_frame_result = write_video_frame_result_ok(env, state);
 
 end:
-  av_packet_unref(packet);
   av_packet_free(&packet);
   return write_frame_result;
 }
@@ -187,18 +186,20 @@ UNIFEX_TERM write_audio_frame(UnifexEnv *env, State *state,
       state->output_ctx->streams[state->audio_stream_index]->time_base;
   AVPacket *packet = av_packet_alloc();
 
+  uint8_t* data = (uint8_t *)av_malloc(frame->size);
+  memcpy(data, frame->data, frame->size);
+  av_packet_from_data(packet, data, frame->size);
+
   UNIFEX_TERM write_frame_result;
 
   packet->stream_index = state->audio_stream_index;
-  packet->size = frame->size;
 
-  packet->data = (uint8_t *)av_malloc(frame->size);
   if (!packet->data) {
     write_frame_result =
         unifex_raise(env, "Failed allocating audio frame data.");
     goto end;
   }
-  memcpy(packet->data, frame->data, frame->size);
+
   int64_t pts_scaled =
       av_rescale_q(pts, MEMBRANE_TIME_BASE, audio_stream_time_base);
   // Packet DTS is set to PTS since AAC buffers do not contain DTS
@@ -208,15 +209,14 @@ UNIFEX_TERM write_audio_frame(UnifexEnv *env, State *state,
   packet->duration = pts_scaled - state->current_audio_pts;
   state->current_audio_pts = pts_scaled;
 
-  if (av_interleaved_write_frame(state->output_ctx, packet)) {
+  if (av_write_frame(state->output_ctx, packet)) {
     write_frame_result =
         write_audio_frame_result_error(env, "Failed writing audio frame");
     goto end;
   }
   write_frame_result = write_audio_frame_result_ok(env, state);
 
 end:
-  av_packet_unref(packet);
   av_packet_free(&packet);
   return write_frame_result;
 }

lib/membrane_rtmp_plugin/rtmp/sink/sink.ex

@@ -12,16 +12,20 @@ defmodule Membrane.RTMP.Sink do
   require Membrane.Logger
 
   alias __MODULE__.Native
-  alias Membrane.{AAC, MP4}
+  alias Membrane.{AAC, Buffer, MP4}
 
   @supported_protocols ["rtmp://", "rtmps://"]
   @connection_attempt_interval 500
   @default_state %{
     attempts: 0,
     native: nil,
-    buffered_frame: nil,
-    ready: false,
-    current_timestamps: %{}
+    # Keys here are the pad names.
+    frame_buffer: %{audio: nil, video: nil},
+    ready?: false,
+    # Activated when one of the source inputs gets closed. Interleaving is
+    # disabled, frame buffer is flushed and from that point buffers on the
+    # remaining pad are simply forwarded to the output.
+    forward_mode?: false
   }
 
   def_input_pad :audio,
@@ -130,52 +134,52 @@ defmodule Membrane.RTMP.Sink do
 
   @impl true
   def handle_write(pad, buffer, _ctx, %{ready: false} = state) do
-    state = Map.put(state, :buffered_frame, {pad, buffer})
-    {:ok, state}
+    {:ok, fill_frame_buffer(state, pad, buffer)}
   end
 
-  @impl true
-  def handle_write(
-        pad,
-        buffer,
-        _ctx,
-        %{ready: true, buffered_frame: {frame_pad, frame}} = state
-      ) do
-    state = write_frame(state, frame_pad, frame) |> write_frame(pad, buffer)
-    {{:ok, demand: get_demand(state)}, Map.put(state, :buffered_frame, nil)}
+  def handle_write(pad, buffer, _ctx, %{forward_mode?: true} = state) do
+    {{:ok, demand: pad}, write_frame(state, pad, buffer)}
   end
 
-  @impl true
   def handle_write(pad, buffer, _ctx, state) do
-    state = write_frame(state, pad, buffer)
-    {{:ok, demand: get_demand(state)}, state}
+    state
+    |> fill_frame_buffer(pad, buffer)
+    |> write_frame_interleaved()
   end
 
   @impl true
-  def handle_end_of_stream(pad, ctx, state) do
-    if ctx.pads |> Map.values() |> Enum.all?(& &1.end_of_stream?) do
+  def handle_end_of_stream(pad, _ctx, state) do
+    if state.forward_mode? do
       Native.finalize_stream(state.native)
       {:ok, state}
     else
-      state = Map.update!(state, :current_timestamps, &Map.delete(&1, pad))
-      {{:ok, demand: get_demand(state)}, state}
+      # The interleave logic does not work if either one of the inputs does not
+      # produce buffers. From this point on we act as a "forward" filter.
+      other_pad =
+        case pad do
+          :audio -> :video
+          :video -> :audio
+        end
+
+      state = flush_frame_buffer(state)
+      {{:ok, demand: other_pad}, %{state | forward_mode?: true}}
     end
   end
 
   @impl true
   def handle_other(:try_connect, _ctx, %{attempts: attempts, max_attempts: max_attempts} = state)
       when max_attempts != :infinity and attempts >= max_attempts do
-    raise "Failed to connect to '#{state.rtmp_url}' #{attempts} times, aborting"
+    raise "failed to connect to '#{state.rtmp_url}' #{attempts} times, aborting"
   end
 
-  def handle_other(:try_connect, ctx, state) do
+  def handle_other(:try_connect, _ctx, state) do
     state = %{state | attempts: state.attempts + 1}
 
     case Native.try_connect(state.native) do
       :ok ->
         Membrane.Logger.debug("Correctly initialized connection with: #{state.rtmp_url}")
-        demands = ctx.pads |> Map.keys() |> Enum.map(&{:demand, &1})
-        {{:ok, [{:playback_change, :resume} | demands]}, state}
+
+        {{:ok, [{:playback_change, :resume} | build_demand(state)]}, state}
 
       {:error, error} when error in [:econnrefused, :etimedout] ->
         Process.send_after(self(), :try_connect, @connection_attempt_interval)
@@ -187,22 +191,70 @@ defmodule Membrane.RTMP.Sink do
         {:ok, state}
 
       {:error, reason} ->
-        raise "Failed to connect to '#{state.rtmp_url}': #{reason}"
+        raise "failed to connect to '#{state.rtmp_url}': #{reason}"
+    end
+  end
+
+  defp build_demand(%{frame_buffer: frame_buffer}) do
+    frame_buffer
+    |> Enum.filter(fn {_pad, buffer} -> buffer == nil end)
+    |> Enum.map(fn {pad, _} -> {:demand, pad} end)
+  end
+
+  defp fill_frame_buffer(state, pad, buffer) do
+    if get_in(state, [:frame_buffer, pad]) == nil do
+      put_in(state, [:frame_buffer, pad], buffer)
+    else
+      raise "attempted to overwrite frame buffer on pad #{inspect(pad)}"
     end
   end
 
+  defp write_frame_interleaved(state = %{frame_buffer: %{audio: audio, video: video}})
+       when audio == nil or video == nil do
+    # We still have to wait for the other frame.
+    {:ok, state}
+  end
+
+  defp write_frame_interleaved(%{frame_buffer: frame_buffer} = state) do
+    {pad, buffer} =
+      Enum.min_by(frame_buffer, fn {_, buffer} ->
+        buffer
+        |> Buffer.get_dts_or_pts()
+        |> Ratio.ceil()
+      end)
+
+    state =
+      state
+      |> write_frame(pad, buffer)
+      |> put_in([:frame_buffer, pad], nil)
+
+    {{:ok, build_demand(state)}, state}
+  end
+
+  defp flush_frame_buffer(%{frame_buffer: frame_buffer} = state) do
+    pads_with_buffer =
+      frame_buffer
+      |> Enum.filter(fn {_pad, buffer} -> buffer != nil end)
+      |> Enum.sort(fn {_, left}, {_, right} ->
+        Buffer.get_dts_or_pts(left) <= Buffer.get_dts_or_pts(right)
+      end)
+
+    Enum.reduce(pads_with_buffer, state, fn {pad, buffer}, state ->
+      state
+      |> write_frame(pad, buffer)
+      |> put_in([:frame_buffer, pad], nil)
+    end)
+  end
+
   defp write_frame(state, :audio, buffer) do
     buffer_pts = buffer.pts |> Ratio.ceil()
 
     case Native.write_audio_frame(state.native, buffer.payload, buffer_pts) do
       {:ok, native} ->
         Map.put(state, :native, native)
-        |> Map.update!(:current_timestamps, fn curr_tmps ->
-          Map.put(curr_tmps, :audio, buffer_pts)
-        end)
 
       {:error, reason} ->
-        raise("Writing audio frame failed with reason: #{reason}")
+        raise "writing audio frame failed with reason: #{inspect(reason)}"
     end
   end
 
@@ -216,19 +268,9 @@ defmodule Membrane.RTMP.Sink do
          ) do
       {:ok, native} ->
         Map.put(state, :native, native)
-        |> Map.update!(:current_timestamps, fn curr_tmps ->
-          Map.put(curr_tmps, :video, buffer.dts)
-        end)
 
       {:error, reason} ->
-        raise("Writing video frame failed with reason: #{reason}")
+        raise "writing video frame failed with reason: #{inspect(reason)}"
     end
   end
-
-  defp get_demand(state) do
-    {pad, _timestamp} =
-      state.current_timestamps |> Enum.min_by(fn {_pad, timestamp} -> timestamp end)
-
-    {pad, 1}
-  end
 end

mix.exs

@@ -43,9 +43,11 @@ defmodule Membrane.RTMP.Mixfile do
       {:membrane_h264_ffmpeg_plugin, "~> 0.22.0"},
       {:membrane_aac_plugin, "~> 0.12.1"},
       {:membrane_mp4_plugin, "~> 0.16.0"},
+      {:membrane_file_plugin, "~> 0.12.0"},
       # testing
       {:membrane_hackney_plugin, "~> 0.8.0", only: :test},
       {:ffmpex, "~> 0.7", only: :test},
+      {:membrane_stream_plugin, "~> 0.1", only: :test},
       # development
       {:ex_doc, "~> 0.28", only: :dev, runtime: false},
       {:dialyxir, "~> 1.1", only: :dev, runtime: false},

mix.lock

@@ -22,7 +22,7 @@
   "makeup_elixir": {:hex, :makeup_elixir, "0.16.0", "f8c570a0d33f8039513fbccaf7108c5d750f47d8defd44088371191b76492b0b", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "28b2cbdc13960a46ae9a8858c4bebdec3c9a6d7b4b9e7f4ed1502f8159f338e7"},
   "makeup_erlang": {:hex, :makeup_erlang, "0.1.1", "3fcb7f09eb9d98dc4d208f49cc955a34218fc41ff6b84df7c75b3e6e533cc65f", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "174d0809e98a4ef0b3309256cbf97101c6ec01c4ab0b23e926a9e17df2077cbb"},
   "membrane_aac_format": {:hex, :membrane_aac_format, "0.7.0", "93ea75b57d5cbee5db59f0250baba65a8949d32ee4cfa38224e93ea306c30048", [:mix], [{:bimap, "~> 1.1", [hex: :bimap, repo: "hexpm", optional: false]}], "hexpm", "3233d2a5a1b264039a198abcef3b7f1c64f5bde9704ec01afbf1bab43a56d415"},
-  "membrane_aac_plugin": {:hex, :membrane_aac_plugin, "0.12.1", "3cca913e084e331015513a6bf5fc83e40c92fbc97843d2f28ac4c4ac406b2cdb", [:mix], [{:bunch, "~> 1.0", [hex: :bunch, repo: "hexpm", optional: false]}, {:crc, "~> 0.10.2", [hex: :crc, repo: "hexpm", optional: false]}, {:membrane_aac_format, "~> 0.7.0", [hex: :membrane_aac_format, repo: "hexpm", optional: false]}, {:membrane_core, "~> 0.10.0", [hex: :membrane_core, repo: "hexpm", optional: false]}], "hexpm", "f6e097402e60e450b970a6ef4f8161161e0e3ba16de371ed15166237d4bb5226"},
+  "membrane_aac_plugin": {:hex, :membrane_aac_plugin, "0.12.2", "f688b8001e010377f33d4fd9ca9983f77a31b193f742c116923b51ffc07f1c0c", [:mix], [{:bunch, "~> 1.0", [hex: :bunch, repo: "hexpm", optional: false]}, {:crc, "~> 0.10.2", [hex: :crc, repo: "hexpm", optional: false]}, {:membrane_aac_format, "~> 0.7.0", [hex: :membrane_aac_format, repo: "hexpm", optional: false]}, {:membrane_core, "~> 0.10.0", [hex: :membrane_core, repo: "hexpm", optional: false]}], "hexpm", "cde8f7c5afdca51b25fa656778f57631074a4195e27ee1940ebf336e3740bede"},
   "membrane_cmaf_format": {:hex, :membrane_cmaf_format, "0.6.1", "89d130b5f8786d4285d395697b0f7763a2c82a02de1658cdeb4f8e37e6a6c85c", [:mix], [], "hexpm", "e916e3c8216f3bf999b069ffda94da48c9bdbe3181ce7155a458d1ccf1a97b3d"},
   "membrane_common_c": {:hex, :membrane_common_c, "0.13.0", "c314623f93209eb2fa092379954c686f6e50ac89baa48360f836d24f4d53f5ee", [:mix], [{:membrane_core, "~> 0.10.0", [hex: :membrane_core, repo: "hexpm", optional: false]}, {:shmex, "~> 0.5.0", [hex: :shmex, repo: "hexpm", optional: false]}, {:unifex, "~> 1.0", [hex: :unifex, repo: "hexpm", optional: false]}], "hexpm", "90181fbbe481ccd0a4a76daf0300f8ad1b5b0bf0ebd8b42c133904f8839663ca"},
   "membrane_core": {:hex, :membrane_core, "0.10.2", "d2d17039f6df746e4a3c47da32f51867fbafe528272cdd9b226d16b1032bc337", [:mix], [{:bunch, "~> 1.3", [hex: :bunch, repo: "hexpm", optional: false]}, {:qex, "~> 0.3", [hex: :qex, repo: "hexpm", optional: false]}, {:ratio, "~> 2.0", [hex: :ratio, repo: "hexpm", optional: false]}, {:telemetry, "~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "6a4f290f919ada66c772807d64d5830be2962b7c13a2f2bc9ace416a1cd19ee1"},
@@ -34,6 +34,7 @@
   "membrane_mp4_plugin": {:hex, :membrane_mp4_plugin, "0.16.2", "e893c0c107b0ec9f00918e8c8927d0cebb15f6cb6f91d08bef5c9a7c333f701d", [:mix], [{:membrane_aac_format, "~> 0.7.0", [hex: :membrane_aac_format, repo: "hexpm", optional: false]}, {:membrane_cmaf_format, "~> 0.6.0", [hex: :membrane_cmaf_format, repo: "hexpm", optional: false]}, {:membrane_core, "~> 0.10.0", [hex: :membrane_core, repo: "hexpm", optional: false]}, {:membrane_file_plugin, "~> 0.12.0", [hex: :membrane_file_plugin, repo: "hexpm", optional: false]}, {:membrane_h264_format, "~> 0.3", [hex: :membrane_h264_format, repo: "hexpm", optional: false]}, {:membrane_mp4_format, "~> 0.7.0", [hex: :membrane_mp4_format, repo: "hexpm", optional: false]}, {:membrane_opus_format, "~> 0.3.0", [hex: :membrane_opus_format, repo: "hexpm", optional: false]}], "hexpm", "3c3726e551eb14f8368d93849b79a765d7a08c66d29c33750cbc3e1696ea7b90"},
   "membrane_opus_format": {:hex, :membrane_opus_format, "0.3.0", "3804d9916058b7cfa2baa0131a644d8186198d64f52d592ae09e0942513cb4c2", [:mix], [], "hexpm", "8fc89c97be50de23ded15f2050fe603dcce732566fe6fdd15a2de01cb6b81afe"},
   "membrane_raw_video_format": {:hex, :membrane_raw_video_format, "0.2.0", "cda8eb207cf65c93690a19001aba3edbb2ba5d22abc8068a1f6a785ba871e8cf", [:mix], [], "hexpm", "6b716fc24f60834323637c95aaaa0f99be23fcc6a84a21af70195ef50185b634"},
+  "membrane_stream_plugin": {:hex, :membrane_stream_plugin, "0.1.0", "0a21e36c4523fae2c514d2661e49a0630f548dcdfe5d18bbf07a8c15e6d1b036", [:mix], [{:membrane_core, "~> 0.10.0", [hex: :membrane_core, repo: "hexpm", optional: false]}], "hexpm", "303658ac9e7915b77ed0a127df2813bdc3d57a5d287bbc0328506c21021138c4"},
   "metrics": {:hex, :metrics, "1.0.1", "25f094dea2cda98213cecc3aeff09e940299d950904393b2a29d191c346a8486", [:rebar3], [], "hexpm", "69b09adddc4f74a40716ae54d140f93beb0fb8978d8636eaded0c31b6f099f16"},
   "mimerl": {:hex, :mimerl, "1.2.0", "67e2d3f571088d5cfd3e550c383094b47159f3eee8ffa08e64106cdf5e981be3", [:rebar3], [], "hexpm", "f278585650aa581986264638ebf698f8bb19df297f66ad91b18910dfc6e19323"},
   "mockery": {:hex, :mockery, "2.3.1", "a02fd60b10ac9ed37a7a2ecf6786c1f1dd5c75d2b079a60594b089fba32dc087", [:mix], [], "hexpm", "1d0971d88ebf084e962da3f2cfee16f0ea8e04ff73a7710428500d4500b947fa"},

test/fixtures/audio.msr

@@ -17,6 +17,30 @@ defmodule Membrane.RTMP.Sink.Test do
     %{flv_output_file: flv_output_file}
   end
 
+  @tag :tmp_dir
+  test "Checks if audio and video are interleaved correctly", %{tmp_dir: tmp_dir} do
+    output_file = Path.join(tmp_dir, "rtmp_sink_interleave_test.flv")
+
+    rtmp_server = Task.async(fn -> start_rtmp_server(output_file) end)
+    {:ok, sink_pipeline_pid} = start_interleaving_sink_pipeline(@rtmp_server_url)
+
+    # There's an RC - there's no way to ensure RTMP server starts to listen before pipeline is started
+    # so it may retry a few times before succeeding
+    assert_pipeline_playback_changed(sink_pipeline_pid, :prepared, :playing, 5000)
+
+    assert_start_of_stream(sink_pipeline_pid, :rtmp_sink, :video, 5_000)
+    assert_start_of_stream(sink_pipeline_pid, :rtmp_sink, :audio, 5_000)
+
+    assert_end_of_stream(sink_pipeline_pid, :rtmp_sink, :video, 5_000)
+    assert_end_of_stream(sink_pipeline_pid, :rtmp_sink, :audio, 5_000)
+
+    Membrane.Testing.Pipeline.terminate(sink_pipeline_pid, blocking?: true)
+    # RTMP server should terminate when the connection is closed
+    assert :ok = Task.await(rtmp_server)
+
+    assert File.exists?(output_file)
+  end
+
   @tag :tmp_dir
   test "Check if audio and video streams are correctly received by RTMP server instance", %{
     flv_output_file: flv_output_file
@@ -43,6 +67,36 @@ defmodule Membrane.RTMP.Sink.Test do
     assert File.stat!(flv_output_file).size == File.stat!(@reference_flv_path).size
   end
 
+  defp start_interleaving_sink_pipeline(rtmp_url) do
+    import Membrane.ParentSpec
+
+    options = [
+      children: [
+        rtmp_sink: %Membrane.RTMP.Sink{rtmp_url: rtmp_url, max_attempts: 5}
+      ],
+      links: [
+        link(:video_source, %Membrane.File.Source{location: "test/fixtures/video.msr"})
+        |> to(:video_deserializer, Membrane.Stream.Deserializer)
+        |> to(:video_parser, %Membrane.H264.FFmpeg.Parser{
+          alignment: :au,
+          attach_nalus?: true,
+          skip_until_parameters?: true
+        })
+        |> to(:video_payloader, Membrane.MP4.Payloader.H264)
+        |> via_in(:video)
+        |> to(:rtmp_sink),
+        link(:audio_source, %Membrane.File.Source{location: "test/fixtures/audio.msr"})
+        |> to(:audio_deserializer, Membrane.Stream.Deserializer)
+        |> to(:audio_parser, Membrane.AAC.Parser)
+        |> via_in(:audio)
+        |> to(:rtmp_sink)
+      ],
+      test_process: self()
+    ]
+
+    Pipeline.start_link(options)
+  end
+
   defp start_sink_pipeline(rtmp_url) do
     import Membrane.ParentSpec