silver_pose/v2/internal/source/stream.go

// Package source owns the FFprobe/FFmpeg stream lifecycle for one camera.
package source

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Status string

const (
	Connecting Status = "CONNECTING"
	Connected  Status = "CONNECTED"
	Retrying   Status = "RETRYING"
	Stopped    Status = "STOPPED"
)

// Config is intentionally separate from the persisted configuration. SourceURL
// exists only in process memory after config resolves its named environment variable.
type Config struct {
	SourceURL   string
	FFmpegPath  string
	FFprobePath string
	Transport   string
	Timeout     time.Duration
	LowLatency  bool
}

type Metadata struct {
	Width  int
	Height int
	FPS    float64
}

type Frame struct {
	BGR       []byte
	Width     int
	Height    int
	Timestamp time.Duration
	Sequence  uint64
}

type Update struct {
	Status  Status
	Message string
	Attempt int
}

// Dependencies make lifecycle behavior testable without an FFmpeg installation
// or a real source URL. Nil functions use the safe default subprocess adapter.
type Dependencies struct {
	Probe  func(context.Context, Config) (Metadata, error)
	Decode func(context.Context, Config, Metadata, func([]byte)) error
	Wait   func(context.Context, time.Duration) error
}

type Stream struct {
	frames   chan Frame
	statuses chan Update
	done     chan struct{}
	once     sync.Once
}

// Start begins one reconnecting stream worker. Consumers receive at most one
// pending frame: a slow inference/UI path always observes the most recent frame.
func Start(ctx context.Context, config Config, dependencies Dependencies) *Stream {
	stream := &Stream{
		frames:   make(chan Frame, 1),
		statuses: make(chan Update, 16),
		done:     make(chan struct{}),
	}
	if dependencies.Probe == nil {
		dependencies.Probe = probe
	}
	if dependencies.Decode == nil {
		dependencies.Decode = decode
	}
	if dependencies.Wait == nil {
		dependencies.Wait = waitContext
	}
	go stream.run(ctx, config, dependencies)
	return stream
}

func (stream *Stream) Frames() <-chan Frame    { return stream.frames }
func (stream *Stream) Statuses() <-chan Update { return stream.statuses }
func (stream *Stream) Done() <-chan struct{}   { return stream.done }

func (stream *Stream) run(ctx context.Context, config Config, dependencies Dependencies) {
	defer stream.close()
	started := time.Now()
	attempt := 0
	sequence := uint64(0)
	for {
		if ctx.Err() != nil {
			stream.publishStatus(Update{Status: Stopped, Message: "视频流已停止", Attempt: attempt})
			return
		}
		stream.publishStatus(Update{Status: Connecting, Message: "正在连接视频流", Attempt: attempt})
		metadata, err := dependencies.Probe(ctx, config)
		if err == nil {
			stream.publishStatus(Update{Status: Connected, Message: "视频流已连接", Attempt: attempt})
			err = dependencies.Decode(ctx, config, metadata, func(bgr []byte) {
				sequence++
				publishLatest(stream.frames, Frame{
					BGR:       append([]byte(nil), bgr...),
					Width:     metadata.Width,
					Height:    metadata.Height,
					Timestamp: time.Since(started),
					Sequence:  sequence,
				})
			})
		}
		if ctx.Err() != nil {
			stream.publishStatus(Update{Status: Stopped, Message: "视频流已停止", Attempt: attempt})
			return
		}
		stream.publishStatus(Update{Status: Retrying, Message: "视频流已断开，正在重连", Attempt: attempt})
		if err := dependencies.Wait(ctx, retryDelay(attempt)); err != nil {
			stream.publishStatus(Update{Status: Stopped, Message: "视频流已停止", Attempt: attempt})
			return
		}
		attempt++
	}
}

func (stream *Stream) close() {
	stream.once.Do(func() {
		close(stream.frames)
		close(stream.statuses)
		close(stream.done)
	})
}

func (stream *Stream) publishStatus(update Update) {
	select {
	case stream.statuses <- update:
	default:
		// Status is advisory. A saturated observer must not block source cleanup.
	}
}

func publishLatest(frames chan Frame, frame Frame) {
	select {
	case frames <- frame:
		return
	default:
	}
	select {
	case <-frames:
	default:
	}
	select {
	case frames <- frame:
	default:
	}
}

func retryDelay(attempt int) time.Duration {
	if attempt <= 0 {
		return time.Second
	}
	if attempt >= 3 {
		return 8 * time.Second
	}
	return time.Second << attempt
}

func waitContext(ctx context.Context, duration time.Duration) error {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func probe(ctx context.Context, config Config) (Metadata, error) {
	if strings.TrimSpace(config.FFprobePath) == "" {
		return Metadata{}, errors.New("FFprobe executable is not configured")
	}
	if strings.TrimSpace(config.SourceURL) == "" {
		return Metadata{}, errors.New("video source is not configured")
	}
	args := []string{"-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height,r_frame_rate", "-of", "json"}
	args = appendInputOptions(args, config)
	args = append(args, "-i", config.SourceURL)
	output, err := exec.CommandContext(ctx, config.FFprobePath, args...).Output()
	if err != nil {
		return Metadata{}, errors.New("unable to probe video metadata")
	}
	var parsed struct {
		Streams []struct {
			Width     int    `json:"width"`
			Height    int    `json:"height"`
			FrameRate string `json:"r_frame_rate"`
		} `json:"streams"`
	}
	if err := json.Unmarshal(output, &parsed); err != nil || len(parsed.Streams) == 0 {
		return Metadata{}, errors.New("invalid video metadata")
	}
	metadata := Metadata{Width: parsed.Streams[0].Width, Height: parsed.Streams[0].Height}
	metadata.FPS = parseFrameRate(parsed.Streams[0].FrameRate)
	if metadata.Width <= 0 || metadata.Height <= 0 || metadata.FPS <= 0 {
		return Metadata{}, errors.New("invalid video dimensions or frame rate")
	}
	return metadata, nil
}

func decode(ctx context.Context, config Config, metadata Metadata, publish func([]byte)) error {
	if strings.TrimSpace(config.FFmpegPath) == "" {
		return errors.New("FFmpeg executable is not configured")
	}
	if strings.TrimSpace(config.SourceURL) == "" {
		return errors.New("video source is not configured")
	}
	if metadata.Width <= 0 || metadata.Height <= 0 {
		return errors.New("invalid video dimensions")
	}
	args := append([]string{"-hide_banner", "-loglevel", "error"}, appendInputOptions(nil, config)...)
	args = append(args, "-i", config.SourceURL, "-an", "-sn", "-dn", "-f", "rawvideo", "-pix_fmt", "bgr24", "-")
	command := exec.CommandContext(ctx, config.FFmpegPath, args...)
	command.Stderr = io.Discard
	stdout, err := command.StdoutPipe()
	if err != nil {
		return errors.New("unable to read video frames")
	}
	if err := command.Start(); err != nil {
		return errors.New("unable to start video decoder")
	}
	frameSize := metadata.Width * metadata.Height * 3
	frame := make([]byte, frameSize)
	for {
		_, err := io.ReadFull(stdout, frame)
		if err != nil {
			_ = command.Wait()
			if ctx.Err() != nil {
				return ctx.Err()
			}
			return errors.New("video decoder stopped")
		}
		publish(frame)
	}
}

func appendInputOptions(args []string, config Config) []string {
	transport := strings.ToLower(strings.TrimSpace(config.Transport))
	if transport == "tcp" || transport == "udp" {
		args = append(args, "-rtsp_transport", transport)
	}
	if config.Timeout > 0 {
		args = append(args, "-rw_timeout", strconv.FormatInt(config.Timeout.Microseconds(), 10))
	}
	if config.LowLatency {
		args = append(args, "-fflags", "nobuffer", "-flags", "low_delay")
	}
	return args
}

func parseFrameRate(value string) float64 {
	parts := strings.Split(value, "/")
	if len(parts) == 2 {
		numerator, numeratorErr := strconv.ParseFloat(parts[0], 64)
		denominator, denominatorErr := strconv.ParseFloat(parts[1], 64)
		if numeratorErr == nil && denominatorErr == nil && denominator > 0 {
			return numerator / denominator
		}
	}
	fps, err := strconv.ParseFloat(value, 64)
	if err != nil || fps <= 0 {
		return 0
	}
	return fps
}

func (metadata Metadata) String() string {
	return fmt.Sprintf("%dx%d@%.3f", metadata.Width, metadata.Height, metadata.FPS)
}