cctv/server/ingest/branch_thumbnail.go

package ingest

import (
	"bytes"
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-gst/go-gst/gst"
	"github.com/go-gst/go-gst/gst/app"
)

type thumbnailReq struct {
	Result chan<- thumbnailRes
}
type thumbnailRes struct {
	Reader *bytes.Reader
	Err    error
}

// addThumbnailBranch populates the pipeline with:
//
//	[vTee] ---> identity[drop D, limit fps] ---> queue ---> [dynamic decoder]
//	 ---> videoconvert ---> aspectratiocrop ---> videoscale ---> capsfilter
//	 ---> jpegenc ---> appsink
//
// Essentially, this branch:
//   - Removes any deltas from the stream.
//   - Limits FPS (we only have keyframes, so this is roughly possible before decoding).
//   - Decode, convert, crop to aspect ratio, scale and encode to JPEG
//   - Then we can query the appsink for thumbnails
//
// width and height are for the dimensions of the generated thumbnails.
// minInterval has no effect when zero or negative, otherwise specifies the minimum interval between thumbnails.
// timeoutSecs is used to deactivate this branch after a period of inactivity (to save CPU), this is done by blocking after the queue.
func (p *CameraPipeline) addThumbnailBranch(width, height int, minInterval, branchTimeout time.Duration) error {
	if p.thumbnailSink != nil {
		return fmt.Errorf("thumbnail branch already added")
	}
	if p.vTee == nil {
		return fmt.Errorf("video tee not initialized")
	}
	vPad := p.vTee.GetRequestPad("src_%u")
	if vPad == nil {
		return fmt.Errorf("failed to get request pad from video tee")
	}

	preDecodeChain, err := p.addStaticChainProps(
		vPad, nil,
		elementFactory{name: "identity", props: map[string]any{
			"drop-buffer-flags": gst.BufferFlagDeltaUnit,
		}},
		elementFactory{name: "queue", props: map[string]any{
			"max-size-buffers": uint(1),
			"leaky":            2, // downstream
		}},
	)
	if err != nil {
		return fmt.Errorf("failed to add pre-decode chain: %w", err)
	}
	preDecodeSrc := preDecodeChain[len(preDecodeChain)-1].GetStaticPad("src")
	if preDecodeSrc == nil {
		return fmt.Errorf("failed to get src pad of pre-decode chain")
	}
	identitySrc := preDecodeChain[0].GetStaticPad("src")
	if identitySrc == nil {
		return fmt.Errorf("failed to get src pad of identity element in pre-decode chain")
	}

	transformChain, err := p.addStaticChainProps(
		nil, nil,
		elementFactory{name: "videoconvert", props: nil},
		elementFactory{name: "aspectratiocrop", props: map[string]any{
			"aspect-ratio": gst.Fraction(width, height),
		}},
		elementFactory{name: "videoscale", props: nil},
		elementFactory{name: "capsfilter", props: map[string]any{
			"caps": gst.NewCapsFromString(fmt.Sprintf("video/x-raw,width=%d,height=%d", width, height)),
		}},
		elementFactory{name: "jpegenc", props: map[string]any{
			"quality": 85,
		}},
		elementFactory{name: "appsink", props: map[string]any{
			"max-buffers": uint(1),
			"drop":        true,
			"sync":        false,
			"async":       false,
		}},
	)
	if err != nil {
		return fmt.Errorf("failed to add transform chain: %w", err)
	}

	if minInterval > 0 {
		minInterval := uint64(minInterval.Nanoseconds())
		lastKeptTS, haveLastTS := uint64(0), false
		identitySrc.AddProbe(gst.PadProbeTypeBuffer, func(_ *gst.Pad, info *gst.PadProbeInfo) gst.PadProbeReturn {
			buf := info.GetBuffer()
			if buf == nil {
				return gst.PadProbeOK
			}

			ts := buf.PresentationTimestamp()
			if ts == gst.ClockTimeNone {
				ts = buf.DecodingTimestamp()
			}
			if ts == gst.ClockTimeNone || ts < 0 {
				return gst.PadProbeOK
			}

			curr := uint64(ts)
			if haveLastTS {
				resetOccurred := curr < lastKeptTS
				if !resetOccurred && curr-lastKeptTS < minInterval {
					return gst.PadProbeDrop
				}
			}

			haveLastTS = true
			lastKeptTS = curr
			return gst.PadProbeOK
		})
	}

	// Dynamically add decoder (decodebin has an internal queue that is unnecessary for our use case,
	// so it just adds unwanted latency).
	transformSink := transformChain[0].GetStaticPad("sink")
	if transformSink == nil {
		return fmt.Errorf("failed to get sink pad for transform chain")
	}
	once := sync.Once{}
	preDecodeSrc.AddProbe(gst.PadProbeTypeEventDownstream, func(_ *gst.Pad, info *gst.PadProbeInfo) gst.PadProbeReturn {
		event := info.GetEvent()
		if event != nil && event.Type() == gst.EventTypeCaps {
			caps := event.ParseCaps()
			if caps == nil || caps.GetSize() == 0 {
				p.log.Error("Thumbnail pad probe received caps event with no or empty caps")
				return gst.PadProbeOK
			}
			if caps.GetSize() > 1 {
				// Note: should be impossible
				p.log.Warn("Thumbnail pad probe received caps event with multiple structures, using the first one", "caps", caps.String())
			}
			structure := caps.GetStructureAt(0)
			if structure == nil {
				p.log.Error("Thumbnail pad probe received caps event with nil structure")
				return gst.PadProbeOK
			}

			res := gst.PadProbeOK
			once.Do(func() {
				res = gst.PadProbeRemove
				var (
					chain []*gst.Element
					err   error
				)
				switch structure.Name() {
				case "video/x-h264":
					chain, err = p.addStaticChain(preDecodeSrc, transformSink, "avdec_h264")
				case "video/x-h265":
					chain, err = p.addStaticChain(preDecodeSrc, transformSink, "avdec_h265")
				default:
					p.log.Error("Thumbnail pad probe received caps event with unsupported encoding", "caps", caps.String())
					return
				}
				if err != nil {
					p.log.Error("Failed to add decoder chain in thumbnail pad probe", "caps", caps.String(), "error", err)
					return
				}
				for _, element := range chain {
					element.SyncStateWithParent()
				}
			})
			return res
		} else {
			return gst.PadProbeOK
		}
	})

	p.thumbnailBlockPad = preDecodeSrc
	p.thumbnailSink = app.SinkFromElement(transformChain[len(transformChain)-1])
	p.thumbnailTimeout = branchTimeout
	p.thumbnailReq = make(chan thumbnailReq)
	return nil
}

func (p *CameraPipeline) addThumbnailBlockProbe() {
	p.thumbnailBlockProbe = p.thumbnailBlockPad.AddProbe(gst.PadProbeTypeBlock, func(pad *gst.Pad, info *gst.PadProbeInfo) gst.PadProbeReturn {
		p.thumbnailBlocked.Store(true)
		return gst.PadProbeOK
	})
}

func (p *CameraPipeline) handleThumbnailReq(active bool, lastThumbnail []byte) (bool, []byte, error) {
	var sample *gst.Sample
	if active {
		sample = p.thumbnailSink.TryPullSample(0)
	} else {
		p.log.Info("Activating thumbnail branch")
		// Install one-shot probe to drop the first buffer (which is a stale frame that was blocked by the blocking probe
		// when the branch was first deactivated).
		if p.thumbnailBlocked.Swap(false) {
			dropProbeId := atomic.Pointer[uint64]{}
			once := sync.Once{}
			probeId := p.thumbnailBlockPad.AddProbe(gst.PadProbeTypeBuffer, func(_ *gst.Pad, _ *gst.PadProbeInfo) gst.PadProbeReturn {
				if probeId := dropProbeId.Load(); probeId != nil {
					p.thumbnailBlockPad.RemoveProbe(*probeId)
				}
				res := gst.PadProbeRemove
				once.Do(func() {
					res = gst.PadProbeDrop
				})
				return res
			})
			dropProbeId.Store(&probeId)
		}
		p.thumbnailBlockPad.RemoveProbe(p.thumbnailBlockProbe)
		active = true
		sample = p.thumbnailSink.PullSample()
	}
	if sample == nil {
		if lastThumbnail == nil {
			return active, lastThumbnail, fmt.Errorf("no thumbnail available")
		} else {
			p.log.Debug("Resolving last thumbnail")
			return active, lastThumbnail, nil
		}
	}
	p.log.Debug("Resolving new thumbnail")

	data, err := gstBufferToBytes(sample.GetBuffer())
	if err != nil {
		return active, lastThumbnail, fmt.Errorf("failed to convert thumbnail buffer to bytes: %w", err)
	}
	return active, data, nil
}

func (p *CameraPipeline) thumbnailManager(ctx context.Context) {
	active := true
	var lastThumbnail []byte

	defer p.log.Debug("Thumbnail manager exiting")
ManagerLoop:
	for {
		var timeout <-chan time.Time
		if active {
			timeout = time.After(p.thumbnailTimeout)
		}
		select {
		case <-ctx.Done():
			return
		case req := <-p.thumbnailReq:
			var (
				thumbnail []byte
				err       error
			)
			active, thumbnail, err = p.handleThumbnailReq(active, lastThumbnail)
			if err != nil {
				req.Result <- thumbnailRes{Err: err}
				continue ManagerLoop
			}

			lastThumbnail = thumbnail
			select {
			case req.Result <- thumbnailRes{Reader: bytes.NewReader(thumbnail)}:
			case <-ctx.Done():
				return
			}
		case <-timeout:
			p.log.Info("Deactivating thumbnail branch due to inactivity")
			// Install block probe to stop flow into the thumbnail branch
			p.addThumbnailBlockProbe()
			// Flush buffer to avoid sending stale thumbnail when reactivating
			for {
				sample := p.thumbnailSink.TryPullSample(0)
				if sample == nil {
					break
				}
			}
			active = false
		}
	}
}

// GetThumbnail either retrieves a JPEG thumbnail from the pipeline, or returns the previous if no new thumbnail is available.
func (p *CameraPipeline) GetThumbnail(ctx context.Context) (*bytes.Reader, error) {
	if p.thumbnailSink == nil {
		return nil, fmt.Errorf("thumbnail branch not initialized")
	}

	resp := make(chan thumbnailRes, 1)
	select {
	case p.thumbnailReq <- thumbnailReq{Result: resp}:
	case <-ctx.Done():
		return nil, fmt.Errorf("context cancelled while requesting thumbnail")
	}

	select {
	case thumbnail := <-resp:
		return thumbnail.Reader, thumbnail.Err
	case <-ctx.Done():
		return nil, fmt.Errorf("context cancelled while waiting for thumbnail")
	}
}