From 2019d5a9f42022ac55df937ff8fb67772c79ae7f Mon Sep 17 00:00:00 2001
From: David Colburn <xero73@gmail.com>
Date: Mon, 13 May 2024 17:15:29 -0400
Subject: [PATCH] log keyframes when stream starts

---
 pkg/config/base.go            |  1 +
 pkg/pipeline/builder/video.go | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/pkg/config/base.go b/pkg/config/base.go
index 8d20853a..fa118b37 100644
--- a/pkg/config/base.go
+++ b/pkg/config/base.go
@@ -52,6 +52,7 @@ type BaseConfig struct {
 type DebugConfig struct {
 	EnableProfiling bool             `yaml:"enable_profiling"` // create dot file and pprof on internal error
 	PathPrefix      string           `yaml:"path_prefix"`      // filepath prefix for uploads
+	LogKeyFrames    bool             `yaml:"log_keyframes"`    // log first 15s of keyframes when streaming
 	StorageConfig   `yaml:",inline"` // upload config (S3, Azure, GCP, or AliOSS)
 }
 
diff --git a/pkg/pipeline/builder/video.go b/pkg/pipeline/builder/video.go
index 7707ea95..e2582d6d 100644
--- a/pkg/pipeline/builder/video.go
+++ b/pkg/pipeline/builder/video.go
@@ -582,6 +582,28 @@ func (b *VideoBin) addEncoder() error {
 			return err
 		}
 
+		if b.conf.GetStreamConfig() != nil && b.conf.Debug.LogKeyFrames {
+			var firstKeyFrame *time.Duration
+			var keyframes []time.Duration
+			x264Enc.GetStaticPad("src").AddProbe(gst.PadProbeTypeBuffer, func(pad *gst.Pad, info *gst.PadProbeInfo) gst.PadProbeReturn {
+				buffer := info.GetBuffer()
+				if !buffer.HasFlags(gst.BufferFlagDeltaUnit) {
+					clockTime := buffer.PresentationTimestamp().AsDuration()
+					if firstKeyFrame == nil {
+						firstKeyFrame = clockTime
+					}
+					pts := *clockTime - *firstKeyFrame
+					keyframes = append(keyframes, pts)
+					logger.Debugw("keyframe generated", "pts", pts)
+					if pts > time.Second*15 {
+						return gst.PadProbeRemove
+					}
+				}
+
+				return gst.PadProbeOK
+			})
+		}
+
 		return nil
 
 	case types.MimeTypeVP9: