AI-tianyan/internal/stream/ingestor.go

package stream

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"os/exec"
	"sync"
	"time"

	"tianyan-edge/internal/config"
)

type Frame struct {
	StreamID int
	DeviceID string
	URL      string
	JPEG     []byte
	TS       float64
}

type Ingestor struct {
	cfg    *config.Config
	frames chan<- Frame
}

func NewIngestor(cfg *config.Config, frames chan<- Frame) *Ingestor {
	return &Ingestor{cfg: cfg, frames: frames}
}

func (ing *Ingestor) Run(ctx context.Context) {
	if len(ing.cfg.RTSPURLs) == 0 {
		log.Println("stream: no RTSP URLs configured, idle")
		<-ctx.Done()
		return
	}
	var wg sync.WaitGroup
	for i, url := range ing.cfg.RTSPURLs {
		wg.Add(1)
		go func(idx int, u string) {
			defer wg.Done()
			ing.streamLoop(ctx, idx, u)
		}(i, url)
	}
	wg.Wait()
}

func (ing *Ingestor) streamLoop(ctx context.Context, idx int, url string) {
	deviceID := fmt.Sprintf("cam-%03d", idx)
	fps := ing.cfg.InferFPS
	if fps <= 0 {
		fps = 5
	}
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		cmd := exec.CommandContext(ctx, "ffmpeg",
			"-rtsp_transport", "tcp", "-i", url,
			"-vf", fmt.Sprintf("fps=%d", fps),
			"-f", "image2pipe", "-vcodec", "mjpeg", "-q:v", "5", "-",
		)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Printf("stream[%d] stdout pipe failed: %v, retry 10s", idx, err)
			time.Sleep(10 * time.Second)
			continue
		}
		if err := cmd.Start(); err != nil {
			log.Printf("stream[%d] ffmpeg start failed: %v, retry 10s", idx, err)
			time.Sleep(10 * time.Second)
			continue
		}
		log.Printf("stream[%d] connected  url=%s", idx, url)
		ing.readFrames(ctx, stdout, idx, deviceID, url)
		cmd.Wait()
		select {
		case <-ctx.Done():
			return
		default:
			log.Printf("stream[%d] lost, retry 5s", idx)
			time.Sleep(5 * time.Second)
		}
	}
}

func (ing *Ingestor) readFrames(ctx context.Context, r io.Reader, idx int, deviceID, url string) {
	soi, eoi := []byte{0xFF, 0xD8}, []byte{0xFF, 0xD9}
	buf, tmp := make([]byte, 0, 1<<20), make([]byte, 32768)
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		n, err := r.Read(tmp)
		if n > 0 {
			buf = append(buf, tmp[:n]...)
			for {
				s := bytes.Index(buf, soi)
				if s < 0 {
					buf = buf[:0]
					break
				}
				buf = buf[s:]
				e := bytes.Index(buf[2:], eoi)
				if e < 0 {
					break
				}
				end := e + 4
				frame := make([]byte, end)
				copy(frame, buf[:end])
				buf = buf[end:]
				select {
				case ing.frames <- Frame{StreamID: idx, DeviceID: deviceID, URL: url,
					JPEG: frame, TS: float64(time.Now().UnixMilli()) / 1000.0}:
				default:
				}
			}
		}
		if err != nil {
			return
		}
	}
}