obs-ffmpeg: Add replay buffer output

This output allows buffering the encoded data, and then muxing the
current buffer to a file on the spot via a hotkey the user sets.

plugins/obs-ffmpeg/data/locale/en-US.ini

@@ -48,3 +48,6 @@ MediaFileFilter.AllMediaFiles="All Media Files"
 MediaFileFilter.VideoFiles="Video Files"
 MediaFileFilter.AudioFiles="Audio Files"
 MediaFileFilter.AllFiles="All Files"
+
+ReplayBuffer="Replay Buffer"
+ReplayBuffer.Save="Save Replay"

plugins/obs-ffmpeg/obs-ffmpeg-mux.c

@@ -16,9 +16,13 @@
 ******************************************************************************/
 
 #include <obs-module.h>
+#include <obs-hotkey.h>
 #include <obs-avc.h>
 #include <util/dstr.h>
 #include <util/pipe.h>
+#include <util/darray.h>
+#include <util/platform.h>
+#include <util/circlebuf.h>
 #include <util/threading.h>
 #include "ffmpeg-mux/ffmpeg-mux.h"
 
@@ -40,17 +44,55 @@ struct ffmpeg_muxer {
 	volatile bool     active;
 	volatile bool     stopping;
 	volatile bool     capturing;
+
+	/* replay buffer */
+	struct circlebuf  packets;
+	int64_t           cur_size;
+	int64_t           cur_time;
+	int64_t           max_size;
+	int64_t           max_time;
+	int64_t           save_ts;
+	int               keyframes;
+	obs_hotkey_id     hotkey;
+
+	DARRAY(struct encoder_packet) mux_packets;
+	pthread_t                     mux_thread;
+	bool                          mux_thread_joinable;
+	volatile bool                 muxing;
 };
 
-static const char *ffmpeg_mux_getname(void *unused)
+static const char *ffmpeg_mux_getname(void *type)
 {
-	UNUSED_PARAMETER(unused);
+	UNUSED_PARAMETER(type);
 	return obs_module_text("FFmpegMuxer");
 }
 
+static inline void replay_buffer_clear(struct ffmpeg_muxer *stream)
+{
+	while (stream->packets.size > 0) {
+		struct encoder_packet pkt;
+		circlebuf_pop_front(&stream->packets, &pkt, sizeof(pkt));
+		obs_encoder_packet_release(&pkt);
+	}
+
+	circlebuf_free(&stream->packets);
+	stream->cur_size = 0;
+	stream->cur_time = 0;
+	stream->max_size = 0;
+	stream->max_time = 0;
+	stream->save_ts = 0;
+	stream->keyframes = 0;
+}
+
 static void ffmpeg_mux_destroy(void *data)
 {
 	struct ffmpeg_muxer *stream = data;
+
+	replay_buffer_clear(stream);
+	if (stream->mux_thread_joinable)
+		pthread_join(stream->mux_thread, NULL);
+	da_free(stream->mux_packets);
+
 	os_process_pipe_destroy(stream->pipe);
 	dstr_free(&stream->path);
 	bfree(stream);
@@ -177,7 +219,8 @@ static void add_muxer_params(struct dstr *cmd, struct ffmpeg_muxer *stream)
 	dstr_free(&mux);
 }
 
-static void build_command_line(struct ffmpeg_muxer *stream, struct dstr *cmd)
+static void build_command_line(struct ffmpeg_muxer *stream, struct dstr *cmd,
+		const char *path)
 {
 	obs_encoder_t *vencoder = obs_output_get_video_encoder(stream->output);
 	obs_encoder_t *aencoders[MAX_AUDIO_MIXES];
@@ -196,7 +239,11 @@ static void build_command_line(struct ffmpeg_muxer *stream, struct dstr *cmd)
 	dstr_init_move_array(cmd, obs_module_file(FFMPEG_MUX));
 	dstr_insert_ch(cmd, 0, '\"');
 	dstr_cat(cmd, "\" \"");
+
+	dstr_copy(&stream->path, path);
+	dstr_replace(&stream->path, "\"", "\"\"");
 	dstr_cat_dstr(cmd, &stream->path);
+
 	dstr_catf(cmd, "\" %d %d ", vencoder ? 1 : 0, num_tracks);
 
 	if (vencoder)
@@ -213,11 +260,18 @@ static void build_command_line(struct ffmpeg_muxer *stream, struct dstr *cmd)
 	add_muxer_params(cmd, stream);
 }
 
+static inline void start_pipe(struct ffmpeg_muxer *stream, const char *path)
+{
+	struct dstr cmd;
+	build_command_line(stream, &cmd, path);
+	stream->pipe = os_process_pipe_create(cmd.array, "w");
+	dstr_free(&cmd);
+}
+
 static bool ffmpeg_mux_start(void *data)
 {
 	struct ffmpeg_muxer *stream = data;
 	obs_data_t *settings;
-	struct dstr cmd;
 	const char *path;
 
 	if (!obs_output_can_begin_data_capture(stream->output, 0))
@@ -227,14 +281,9 @@ static bool ffmpeg_mux_start(void *data)
 
 	settings = obs_output_get_settings(stream->output);
 	path = obs_data_get_string(settings, "path");
-	dstr_copy(&stream->path, path);
-	dstr_replace(&stream->path, "\"", "\"\"");
+	start_pipe(stream, path);
 	obs_data_release(settings);
 
-	build_command_line(stream, &cmd);
-	stream->pipe = os_process_pipe_create(cmd.array, "w");
-	dstr_free(&cmd);
-
 	if (!stream->pipe) {
 		warn("Failed to create process pipe");
 		return false;
@@ -424,3 +473,330 @@ struct obs_output_info ffmpeg_muxer = {
 	.encoded_packet = ffmpeg_mux_data,
 	.get_properties = ffmpeg_mux_properties
 };
+
+/* ------------------------------------------------------------------------ */
+
+static const char *replay_buffer_getname(void *type)
+{
+	UNUSED_PARAMETER(type);
+	return obs_module_text("ReplayBuffer");
+}
+
+static bool replay_buffer_hotkey(void *data, obs_hotkey_id id,
+		obs_hotkey_t *hotkey, bool pressed)
+{
+	struct ffmpeg_muxer *stream = data;
+	if (os_atomic_load_bool(&stream->active))
+		stream->save_ts = os_gettime_ns() / 1000LL;
+	return true;
+}
+
+static void *replay_buffer_create(obs_data_t *settings, obs_output_t *output)
+{
+	struct ffmpeg_muxer *stream = bzalloc(sizeof(*stream));
+	stream->output = output;
+
+	stream->hotkey = obs_hotkey_register_output(output,
+			"ReplayBuffer.Save",
+			obs_module_text("ReplayBuffer.Save"),
+			replay_buffer_hotkey, stream);
+
+	UNUSED_PARAMETER(settings);
+	return stream;
+}
+
+static void replay_buffer_destroy(void *data)
+{
+	struct ffmpeg_muxer *stream = data;
+	if (stream->hotkey)
+		obs_hotkey_unregister(stream->hotkey);
+	ffmpeg_mux_destroy(data);
+}
+
+static bool replay_buffer_start(void *data)
+{
+	struct ffmpeg_muxer *stream = data;
+
+	if (!obs_output_can_begin_data_capture(stream->output, 0))
+		return false;
+	if (!obs_output_initialize_encoders(stream->output, 0))
+		return false;
+
+	obs_data_t *s = obs_output_get_settings(stream->output);
+	stream->max_time = obs_data_get_int(s, "max_time_sec") * 1000000LL;
+	stream->max_size = obs_data_get_int(s, "max_size_mb") * (1024 * 1024);
+	obs_data_release(s);
+
+	os_atomic_set_bool(&stream->active, true);
+	os_atomic_set_bool(&stream->capturing, true);
+	obs_output_begin_data_capture(stream->output, 0);
+
+	return true;
+}
+
+static bool purge_front(struct ffmpeg_muxer *stream)
+{
+	struct encoder_packet pkt;
+	bool keyframe;
+
+	circlebuf_pop_front(&stream->packets, &pkt, sizeof(pkt));
+
+	keyframe = pkt.type == OBS_ENCODER_VIDEO && pkt.keyframe;
+
+	if (keyframe)
+		stream->keyframes--;
+
+	if (!stream->packets.size) {
+		stream->cur_size = 0;
+		stream->cur_time = 0;
+	} else {
+		struct encoder_packet first;
+		circlebuf_peek_front(&stream->packets, &first, sizeof(first));
+		stream->cur_time = first.dts_usec;
+		stream->cur_size -= (int64_t)pkt.size;
+	}
+
+	obs_encoder_packet_release(&pkt);
+	return keyframe;
+}
+
+static inline void purge(struct ffmpeg_muxer *stream)
+{
+	if (purge_front(stream)) {
+		struct encoder_packet pkt;
+
+		for (;;) {
+			circlebuf_peek_front(&stream->packets, &pkt,
+					sizeof(pkt));
+			if (pkt.type == OBS_ENCODER_VIDEO && pkt.keyframe)
+				return;
+
+			purge_front(stream);
+		}
+	}
+}
+
+static inline void replay_buffer_purge(struct ffmpeg_muxer *stream,
+		struct encoder_packet *pkt)
+{
+	if (stream->max_size) {
+		if (!stream->packets.size || stream->keyframes <= 2)
+			return;
+
+		while ((stream->cur_size + (int64_t)pkt->size) >
+				stream->max_size)
+			purge(stream);
+	}
+
+	if (!stream->packets.size || stream->keyframes <= 2)
+		return;
+
+	while ((pkt->dts_usec - stream->cur_time) > stream->max_time)
+		purge(stream);
+}
+
+static void insert_packet(struct darray *array, struct encoder_packet *packet,
+		int64_t video_offset, int64_t *audio_offsets,
+		int64_t video_dts_offset, int64_t *audio_dts_offsets)
+{
+	struct encoder_packet pkt;
+	DARRAY(struct encoder_packet) packets;
+	packets.da = *array;
+	size_t idx;
+
+	obs_encoder_packet_ref(&pkt, packet);
+
+	if (pkt.type == OBS_ENCODER_VIDEO) {
+		pkt.dts_usec -= video_offset;
+		pkt.dts -= video_dts_offset;
+		pkt.pts -= video_dts_offset;
+	} else {
+		pkt.dts_usec -= audio_offsets[pkt.track_idx];
+		pkt.dts -= audio_dts_offsets[pkt.track_idx];
+		pkt.pts -= audio_dts_offsets[pkt.track_idx];
+	}
+
+	for (idx = packets.num; idx > 0; idx--) {
+		struct encoder_packet *p = packets.array + (idx - 1);
+		if (p->dts_usec < pkt.dts_usec)
+			break;
+	}
+
+	da_insert(packets, idx, &pkt);
+	*array = packets.da;
+}
+
+static void *replay_buffer_mux_thread(void *data)
+{
+	struct ffmpeg_muxer *stream = data;
+
+	start_pipe(stream, stream->path.array);
+
+	if (!stream->pipe) {
+		warn("Failed to create process pipe");
+		goto error;
+	}
+
+	if (!send_headers(stream)) {
+		warn("Could not write headers for file '%s'",
+				stream->path.array);
+		goto error;
+	}
+
+	for (size_t i = 0; i < stream->mux_packets.num; i++) {
+		struct encoder_packet *pkt = &stream->mux_packets.array[i];
+		write_packet(stream, pkt);
+		obs_encoder_packet_release(pkt);
+	}
+
+	info("Wrote replay buffer to '%s'", stream->path.array);
+
+error:
+	os_process_pipe_destroy(stream->pipe);
+	stream->pipe = NULL;
+	da_free(stream->mux_packets);
+	os_atomic_set_bool(&stream->muxing, false);
+	return NULL;
+}
+
+static void replay_buffer_save(struct ffmpeg_muxer *stream)
+{
+	const size_t size = sizeof(struct encoder_packet);
+	size_t num_packets = stream->packets.size / size;
+
+	da_reserve(stream->mux_packets, num_packets);
+
+	/* ---------------------------- */
+	/* reorder packets */
+
+	bool found_video = false;
+	bool found_audio[MAX_AUDIO_MIXES] = {0};
+	int64_t video_offset = 0;
+	int64_t video_dts_offset = 0;
+	int64_t audio_offsets[MAX_AUDIO_MIXES] = {0};
+	int64_t audio_dts_offsets[MAX_AUDIO_MIXES] = {0};
+
+	for (size_t i = 0; i < num_packets; i++) {
+		struct encoder_packet *pkt;
+		pkt = circlebuf_data(&stream->packets, i * size);
+
+		if (pkt->type == OBS_ENCODER_VIDEO) {
+			if (!found_video) {
+				video_offset = pkt->dts_usec;
+				video_dts_offset = pkt->dts;
+				found_video = true;
+			}
+		} else {
+			if (!found_audio[pkt->track_idx]) {
+				found_audio[pkt->track_idx] = true;
+				audio_offsets[pkt->track_idx] = pkt->dts_usec;
+				audio_dts_offsets[pkt->track_idx] = pkt->dts;
+			}
+		}
+
+		insert_packet(&stream->mux_packets.da, pkt,
+				video_offset, audio_offsets,
+				video_dts_offset, audio_dts_offsets);
+	}
+
+	/* ---------------------------- */
+	/* generate filename */
+
+	obs_data_t *settings = obs_output_get_settings(stream->output);
+	const char *dir = obs_data_get_string(settings, "directory");
+	const char *fmt = obs_data_get_string(settings, "format");
+	const char *ext = obs_data_get_string(settings, "extension");
+	bool space = obs_data_get_bool(settings, "allow_spaces");
+
+	char *filename = os_generate_formatted_filename(ext, space, fmt);
+
+	dstr_copy(&stream->path, dir);
+	dstr_replace(&stream->path, "\\", "/");
+	if (dstr_end(&stream->path) != '/')
+		dstr_cat_ch(&stream->path, '/');
+	dstr_cat(&stream->path, filename);
+
+	bfree(filename);
+	obs_data_release(settings);
+
+	/* ---------------------------- */
+
+	os_atomic_set_bool(&stream->muxing, true);
+	stream->mux_thread_joinable = pthread_create(&stream->mux_thread, NULL,
+			replay_buffer_mux_thread, stream) == 0;
+}
+
+static void deactivate_replay_buffer(struct ffmpeg_muxer *stream)
+{
+	if (stopping(stream))
+		obs_output_end_data_capture(stream->output);
+
+	os_atomic_set_bool(&stream->active, false);
+	os_atomic_set_bool(&stream->sent_headers, false);
+	os_atomic_set_bool(&stream->stopping, false);
+	replay_buffer_clear(stream);
+}
+
+static void replay_buffer_data(void *data, struct encoder_packet *packet)
+{
+	struct ffmpeg_muxer *stream = data;
+	struct encoder_packet pkt;
+
+	if (!active(stream))
+		return;
+
+	if (stopping(stream)) {
+		if (packet->sys_dts_usec >= stream->stop_ts) {
+			deactivate_replay_buffer(stream);
+			return;
+		}
+	}
+
+	obs_encoder_packet_ref(&pkt, packet);
+	replay_buffer_purge(stream, &pkt);
+
+	if (!stream->packets.size)
+		stream->cur_time = pkt.dts_usec;
+	stream->cur_size += pkt.size;
+
+	circlebuf_push_back(&stream->packets, packet, sizeof(*packet));
+
+	if (packet->type == OBS_ENCODER_VIDEO && packet->keyframe)
+		stream->keyframes++;
+
+	if (stream->save_ts && packet->sys_dts_usec >= stream->save_ts) {
+		if (os_atomic_load_bool(&stream->muxing))
+			return;
+
+		if (stream->mux_thread_joinable) {
+			pthread_join(stream->mux_thread, NULL);
+			stream->mux_thread_joinable = false;
+		}
+
+		stream->save_ts = 0;
+		replay_buffer_save(stream);
+	}
+}
+
+static void replay_buffer_defaults(obs_data_t *s)
+{
+	obs_data_set_default_int(s, "max_time_sec", 15);
+	obs_data_set_default_int(s, "max_size_mb", 500);
+	obs_data_set_default_string(s, "format", "%CCYY-%MM-%DD %hh-%mm-%ss");
+	obs_data_set_default_string(s, "extension", "mp4");
+	obs_data_set_default_bool(s, "allow_spaces", true);
+}
+
+struct obs_output_info replay_buffer = {
+	.id             = "replay_buffer",
+	.flags          = OBS_OUTPUT_AV |
+	                  OBS_OUTPUT_ENCODED |
+	                  OBS_OUTPUT_MULTI_TRACK,
+	.get_name       = replay_buffer_getname,
+	.create         = replay_buffer_create,
+	.destroy        = replay_buffer_destroy,
+	.start          = replay_buffer_start,
+	.stop           = ffmpeg_mux_stop,
+	.encoded_packet = replay_buffer_data,
+	.get_defaults   = replay_buffer_defaults
+};

plugins/obs-ffmpeg/obs-ffmpeg.c

@@ -11,6 +11,7 @@ OBS_MODULE_USE_DEFAULT_LOCALE("obs-ffmpeg", "en-US")
 extern struct obs_source_info  ffmpeg_source;
 extern struct obs_output_info  ffmpeg_output;
 extern struct obs_output_info  ffmpeg_muxer;
+extern struct obs_output_info  replay_buffer;
 extern struct obs_encoder_info aac_encoder_info;
 extern struct obs_encoder_info nvenc_encoder_info;
 
@@ -145,6 +146,7 @@ bool obs_module_load(void)
 	obs_register_source(&ffmpeg_source);
 	obs_register_output(&ffmpeg_output);
 	obs_register_output(&ffmpeg_muxer);
+	obs_register_output(&replay_buffer);
 	obs_register_encoder(&aac_encoder_info);
 	if (nvenc_supported()) {
 		blog(LOG_INFO, "NVENC supported");