Fix segment collection end timestamp calculation

README.md

@@ -12,7 +12,7 @@ The package can be installed by adding `membrane_mp4_plugin` to your list of dep
 ```elixir
 defp deps do
   [
-    {:membrane_mp4_plugin, "~> 0.32.0"}
+    {:membrane_mp4_plugin, "~> 0.33.0"}
   ]
 end
 ```

lib/membrane_mp4/muxer/cmaf.ex

@@ -335,7 +335,7 @@ defmodule Membrane.MP4.Muxer.CMAF do
 
     {sample, state} =
       state
-      |> maybe_init_segment_base_timestamp(pad, sample)
+      |> maybe_init_segment_timestamps(pad, sample)
       |> process_buffer_awaiting_duration(pad, sample)
 
     state = SegmentHelper.update_awaiting_stream_format(state, pad)
@@ -438,9 +438,12 @@ defmodule Membrane.MP4.Muxer.CMAF do
     track_data = %{
       id: track_id,
       track: nil,
-      # base timestamp of the current segment, initialized with DTS of the first buffer
+      # decoding timestamp of the current segment, initialized with DTS of the first buffer
       # and then incremented by duration of every produced segment
-      segment_base_timestamp: nil,
+      segment_decoding_timestamp: nil,
+      # presentation timestamp of the current segment, initialized with PTS of the first buffer
+      # and then incremented by duration of every produced segment
+      segment_presentation_timestamp: nil,
       end_timestamp: 0,
       buffer_awaiting_duration: nil,
       chunks_duration: Membrane.Time.seconds(0)
@@ -587,7 +590,7 @@ defmodule Membrane.MP4.Muxer.CMAF do
           sequence_number: state.seq_nums[output_pad],
           timescale: timescale,
           base_timestamp:
-            track_data.segment_base_timestamp
+            track_data.segment_presentation_timestamp
             |> Helper.timescalify(timescale)
             |> Ratio.trunc(),
           unscaled_duration: duration,
@@ -603,7 +606,12 @@ defmodule Membrane.MP4.Muxer.CMAF do
       state =
         tracks_data
         |> Enum.reduce(state, fn %{unscaled_duration: duration, pad: pad}, state ->
-          update_in(state, [:pad_to_track_data, pad, :segment_base_timestamp], &(&1 + duration))
+          state
+          |> update_in([:pad_to_track_data, pad, :segment_decoding_timestamp], &(&1 + duration))
+          |> update_in(
+            [:pad_to_track_data, pad, :segment_presentation_timestamp],
+            &(&1 + duration)
+          )
         end)
         |> update_in([:seq_nums, output_pad], &(&1 + 1))
 
@@ -712,10 +720,15 @@ defmodule Membrane.MP4.Muxer.CMAF do
     end
   end
 
-  defp maybe_init_segment_base_timestamp(state, pad, sample) do
+  defp maybe_init_segment_timestamps(state, pad, sample) do
     case state do
-      %{pad_to_track_data: %{^pad => %{segment_base_timestamp: nil}}} ->
-        put_in(state, [:pad_to_track_data, pad, :segment_base_timestamp], sample.dts)
+      %{pad_to_track_data: %{^pad => %{segment_decoding_timestamp: nil}}} ->
+        update_in(state, [:pad_to_track_data, pad], fn data ->
+          Map.merge(data, %{
+            segment_decoding_timestamp: sample.dts,
+            segment_presentation_timestamp: sample.pts
+          })
+        end)
 
       _else ->
         state

lib/membrane_mp4/muxer/cmaf/segment_helper.ex

@@ -120,7 +120,7 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
   end
 
   defp push_video_segment(state, queue, pad, sample) do
-    base_timestamp = max_segment_base_timestamp(state)
+    base_timestamp = max_segment_decoding_timestamp(state)
 
     queue =
       if state.finish_current_segment? do
@@ -144,7 +144,7 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
   end
 
   defp push_audio_segment(state, queue, pad, sample) do
-    base_timestamp = max_segment_base_timestamp(state)
+    base_timestamp = max_segment_decoding_timestamp(state)
 
     {video_pad, video_queue} =
       Enum.find(state.sample_queues, {nil, nil}, fn {_pad, queue} ->
@@ -190,7 +190,7 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
     total_collected_durations =
       Map.fetch!(state.pad_to_track_data, pad).chunks_duration + collected_duration
 
-    base_timestamp = state.pad_to_track_data[pad].segment_base_timestamp
+    base_timestamp = state.pad_to_track_data[pad].segment_decoding_timestamp
 
     queue =
       cond do
@@ -262,7 +262,7 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
       if video_queue do
         SamplesQueue.force_push(queue, sample)
       else
-        base_timestamp = max_segment_base_timestamp(state)
+        base_timestamp = max_segment_decoding_timestamp(state)
 
         SamplesQueue.plain_push_until_target(queue, sample, base_timestamp)
       end
@@ -288,7 +288,7 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
   defp update_queue_for(pad, queue, state), do: put_in(state, [:sample_queues, pad], queue)
 
   defp collect_samples_for_video_track(pad, queue, state) do
-    end_timestamp = SamplesQueue.last_collected_dts(queue)
+    end_timestamp = SamplesQueue.collectable_end_timestamp(queue)
     state = update_queue_for(pad, queue, state)
 
     if tracks_ready_for_collection?(state, end_timestamp) do
@@ -311,7 +311,7 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
   end
 
   defp collect_samples_for_audio_track(pad, queue, state) do
-    end_timestamp = SamplesQueue.last_collected_dts(queue)
+    end_timestamp = SamplesQueue.collectable_end_timestamp(queue)
     state = update_queue_for(pad, queue, state)
 
     if tracks_ready_for_collection?(state, end_timestamp) do
@@ -338,7 +338,7 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
 
   defp tracks_ready_for_collection?(state, end_timestamp) do
     Enum.all?(state.sample_queues, fn {_pad, queue} ->
-      SamplesQueue.last_collected_dts(queue) >= end_timestamp
+      SamplesQueue.collectable_end_timestamp(queue) >= end_timestamp
     end)
   end
 
@@ -393,11 +393,11 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelper do
     maybe_return_segment(segment, reset_chunks_duration(state))
   end
 
-  defp max_segment_base_timestamp(state) do
+  defp max_segment_decoding_timestamp(state) do
     state.pad_to_track_data
-    |> Enum.reject(fn {_key, track_data} -> is_nil(track_data.segment_base_timestamp) end)
+    |> Enum.reject(fn {_key, track_data} -> is_nil(track_data.segment_decoding_timestamp) end)
     |> Enum.map(fn {_key, track_data} ->
-      Ratio.to_float(track_data.segment_base_timestamp)
+      Ratio.to_float(track_data.segment_decoding_timestamp)
     end)
     |> Enum.max()
   end

lib/membrane_mp4/muxer/cmaf/track_samples_queue.ex

@@ -280,24 +280,31 @@ defmodule Membrane.MP4.Muxer.CMAF.TrackSamplesQueue do
   end
 
   @doc """
-  Returns dts of the latest sample that is eligible for collection.
+  Returns the end timestamp for latest sample that is eligible for collection.
 
   In case of collectable state it is the last sample that has been put to queue, otherwise
   it is the last sample that will be in return from 'collect/1'.
   """
-  @spec last_collected_dts(t()) :: integer()
-  def last_collected_dts(%__MODULE__{
+  @spec collectable_end_timestamp(t()) :: integer()
+  def collectable_end_timestamp(%__MODULE__{
         collectable?: false,
         target_samples: target_samples,
         excess_samples: excess_samples
-      }),
-      do: latest_collected_dts(excess_samples) || latest_collected_dts(target_samples) || -1
+      }) do
+    sample = List.first(excess_samples) || List.first(target_samples)
+
+    if sample do
+      sample.dts + sample.metadata.duration
+    else
+      -1
+    end
+  end
 
-  def last_collected_dts(%__MODULE__{collectable?: true, target_samples: target_samples}),
-    do: latest_collected_dts(List.last(target_samples, []) |> List.wrap()) || -1
+  def collectable_end_timestamp(%__MODULE__{collectable?: true, target_samples: target_samples}) do
+    sample = List.last(target_samples)
 
-  defp latest_collected_dts([]), do: nil
-  defp latest_collected_dts([sample | _rest]), do: Ratio.to_float(sample.dts)
+    sample.dts + sample.metadata.duration
+  end
 
   @doc """
   Returns the most recenlty pushed sample.

mix.exs

@@ -1,7 +1,7 @@
 defmodule Membrane.MP4.Plugin.MixProject do
   use Mix.Project
 
-  @version "0.32.0"
+  @version "0.33.0"
   @github_url "https://github.com/membraneframework/membrane_mp4_plugin"
 
   def project do

test/membrane_mp4/muxer/cmaf/segment_helper_test.exs

@@ -16,8 +16,18 @@ defmodule Membrane.MP4.Muxer.CMAF.SegmentHelperTest do
       input_to_output_pad: %{audio: :output, video: :output},
       input_groups: %{output: [:audio, :video]},
       pad_to_track_data: %{
-        audio: %{segment_base_timestamp: 0, chunks_duration: 0, buffer_awaiting_duration: nil},
-        video: %{segment_base_timestamp: 0, chunks_duration: 0, buffer_awaiting_duration: nil}
+        audio: %{
+          segment_decoding_timestamp: 0,
+          segment_presentation_timestamp: 0,
+          chunks_duration: 0,
+          buffer_awaiting_duration: nil
+        },
+        video: %{
+          segment_decoding_timestamp: 0,
+          segment_presentation_timestamp: 0,
+          chunks_duration: 0,
+          buffer_awaiting_duration: nil
+        }
       },
       sample_queues: %{
         audio: %Queue{track_with_keyframes?: false, duration_range: chunk_duration_range},