Fix segment collection end timestamp calculation

lib/membrane_mp4/muxer/cmaf/segment_helper.ex

@@ -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
 

lib/membrane_mp4/muxer/cmaf/track_samples_queue.ex

@@ -299,6 +299,33 @@ defmodule Membrane.MP4.Muxer.CMAF.TrackSamplesQueue do
   defp latest_collected_dts([]), do: nil
   defp latest_collected_dts([sample | _rest]), do: Ratio.to_float(sample.dts)
 
+  @doc """
+  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 collectable_end_timestamp(t()) :: integer()
+  def collectable_end_timestamp(%__MODULE__{
+        collectable?: false,
+        target_samples: target_samples,
+        excess_samples: excess_samples
+      }) do
+    sample = List.first(excess_samples) || List.first(target_samples)
+
+    if sample do
+      sample.dts + sample.metadata.duration
+    else
+      -1
+    end
+  end
+
+  def collectable_end_timestamp(%__MODULE__{collectable?: true, target_samples: target_samples}) do
+    sample = List.last(target_samples)
+
+    sample.dts + sample.metadata.duration
+  end
+
   @doc """
   Returns the most recenlty pushed sample.
   """