From 61c0a76fc0d0306bcddb6ca6b56600855cdcd4d4 Mon Sep 17 00:00:00 2001
From: Norihiro Kamae <norihiro@nagater.net>
Date: Sat, 11 Dec 2021 21:46:10 +0900
Subject: [PATCH] obs-ffmpeg: Split file by PTS instead of DTS

If B-frame is enabled in video encoder, video packets have different PTS
and DTS, however audio packets do not. That caused audio packets at the
splitting point goes to a different file and audio glitch appeared in an
NLE.
---
 plugins/obs-ffmpeg/obs-ffmpeg-mux.c | 57 ++++++++++++++++++++++++++++-
 plugins/obs-ffmpeg/obs-ffmpeg-mux.h |  1 +
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/plugins/obs-ffmpeg/obs-ffmpeg-mux.c b/plugins/obs-ffmpeg/obs-ffmpeg-mux.c
index 9bd0f5b0c..b2f918147 100644
--- a/plugins/obs-ffmpeg/obs-ffmpeg-mux.c
+++ b/plugins/obs-ffmpeg/obs-ffmpeg-mux.c
@@ -320,6 +320,11 @@ inline static void ts_offset_clear(struct ffmpeg_muxer *stream)
 	}
 }
 
+static inline int64_t packet_pts_usec(struct encoder_packet *packet)
+{
+	return packet->pts * 1000000 / packet->timebase_den;
+}
+
 inline static void ts_offset_update(struct ffmpeg_muxer *stream,
 				    struct encoder_packet *packet)
 {
@@ -731,6 +736,19 @@ static bool prepare_split_file(struct ffmpeg_muxer *stream,
 	return true;
 }
 
+static inline bool has_audio(struct ffmpeg_muxer *stream)
+{
+	return !!obs_output_get_audio_encoder(stream->output, 0);
+}
+
+static void push_back_packet(struct darray *packets,
+			     struct encoder_packet *packet)
+{
+	struct encoder_packet pkt;
+	obs_encoder_packet_ref(&pkt, packet);
+	darray_push_back(sizeof(pkt), packets, &pkt);
+}
+
 static void ffmpeg_mux_data(void *data, struct encoder_packet *packet)
 {
 	struct ffmpeg_muxer *stream = data;
@@ -744,9 +762,31 @@ static void ffmpeg_mux_data(void *data, struct encoder_packet *packet)
 		return;
 	}
 
-	if (stream->split_file && should_split(stream, packet)) {
-		if (!prepare_split_file(stream, packet))
+	if (stream->split_file && stream->mux_packets.num) {
+		int64_t pts_usec = packet_pts_usec(packet);
+		struct encoder_packet *first_pkt = stream->mux_packets.array;
+		int64_t first_pts_usec = packet_pts_usec(first_pkt);
+
+		if (pts_usec >= first_pts_usec) {
+			if (packet->type != OBS_ENCODER_AUDIO) {
+				push_back_packet(&stream->mux_packets.da,
+						 packet);
+				return;
+			}
+
+			if (!prepare_split_file(stream, first_pkt))
+				return;
+			stream->split_file_ready = true;
+		}
+	} else if (stream->split_file && should_split(stream, packet)) {
+		if (has_audio(stream)) {
+			push_back_packet(&stream->mux_packets.da, packet);
 			return;
+		} else {
+			if (!prepare_split_file(stream, packet))
+				return;
+			stream->split_file_ready = true;
+		}
 	}
 
 	if (!stream->sent_headers) {
@@ -766,6 +806,19 @@ static void ffmpeg_mux_data(void *data, struct encoder_packet *packet)
 		}
 	}
 
+	if (stream->split_file && stream->split_file_ready) {
+		for (size_t i = 0; i < stream->mux_packets.num; i++) {
+			struct encoder_packet *pkt =
+				&stream->mux_packets.array[i];
+			if (stream->reset_timestamps)
+				ts_offset_update(stream, pkt);
+			write_packet(stream, pkt);
+			obs_encoder_packet_release(pkt);
+		}
+		da_free(stream->mux_packets);
+		stream->split_file_ready = false;
+	}
+
 	if (stream->split_file && stream->reset_timestamps)
 		ts_offset_update(stream, packet);
 
diff --git a/plugins/obs-ffmpeg/obs-ffmpeg-mux.h b/plugins/obs-ffmpeg/obs-ffmpeg-mux.h
index 7d8a55add..1dcaa2c6e 100644
--- a/plugins/obs-ffmpeg/obs-ffmpeg-mux.h
+++ b/plugins/obs-ffmpeg/obs-ffmpeg-mux.h
@@ -42,6 +42,7 @@ struct ffmpeg_muxer {
 	bool found_audio[MAX_AUDIO_MIXES];
 	int64_t video_pts_offset;
 	int64_t audio_dts_offsets[MAX_AUDIO_MIXES];
+	bool split_file_ready;
 
 	/* these are accessed both by replay buffer and by HLS */
 	pthread_t mux_thread;