perf(stream): share H264 capture frames

Route Direct and WebRTC consumers through one H264 capture source to avoid duplicate VENC reads. Overlap WebRTC capture with ordered RTP/SRTP writes and packetize each access unit once for all WebRTC clients.

Add libopencv_video.so.409 into dl_lib.

server/go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/mervick/aes-everywhere/go/aes256 v0.0.0-20240803013625-6759956693c0
 	github.com/modelcontextprotocol/go-sdk v1.6.1
 	github.com/pion/dtls/v3 v3.1.2
+	github.com/pion/interceptor v0.1.39
 	github.com/pion/rtp v1.8.18
 	github.com/pion/webrtc/v4 v4.0.1
 	github.com/rs/cors/wrapper/gin v0.0.0-20240830163046-1084d89a1692
@@ -48,7 +49,6 @@ require (
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
 	github.com/pion/datachannel v1.5.9 // indirect
 	github.com/pion/ice/v4 v4.0.2 // indirect
-	github.com/pion/interceptor v0.1.39 // indirect
 	github.com/pion/logging v0.2.4 // indirect
 	github.com/pion/mdns/v2 v2.0.7 // indirect
 	github.com/pion/randutil v0.1.0 // indirect

server/service/stream/direct/streamer.go

@@ -1,12 +1,10 @@
 package direct
 
 import (
-	"NanoKVM-Server/common"
 	"NanoKVM-Server/service/stream"
 	"NanoKVM-Server/service/vm"
 	"sync"
 	"sync/atomic"
-	"time"
 
 	log "github.com/sirupsen/logrus"
 )
@@ -87,17 +85,14 @@ func (s *Streamer) getClients() []*client {
 }
 
 func (s *Streamer) run() {
-	screen := common.GetScreen()
-	common.CheckScreen()
-	fps := screen.FPS
+	subscription := stream.SubscribeH264()
+	defer subscription.Close()
 
-	ticker := time.NewTicker(time.Second / time.Duration(fps))
-	defer ticker.Stop()
-
-	vision := common.GetKvmVision()
-	startTime := time.Now()
-
-	for range ticker.C {
+	for {
+		frame, ok := subscription.Next()
+		if !ok {
+			return
+		}
 		clients := s.getClients()
 		if len(clients) == 0 {
 			if s.stopIfIdle() {
@@ -107,28 +102,19 @@ func (s *Streamer) run() {
 			continue
 		}
 
-		if screen.FPS != fps && screen.FPS != 0 {
-			fps = screen.FPS
-			ticker.Reset(time.Second / time.Duration(fps))
-		}
-
 		if !hasCaptureDemand(clients) {
 			continue
 		}
 
-		data, result := vision.ReadH264(screen.Width, screen.Height, screen.BitRate)
-		stream.UpdateCaptureStatus(stream.CaptureModeDirect, result)
-		if result < 0 || len(data) == 0 {
+		stream.UpdateCaptureStatus(stream.CaptureModeDirect, frame.Result)
+		if frame.Result < 0 || len(frame.Data) == 0 {
 			continue
 		}
 
-		timestamp := time.Since(startTime).Microseconds()
-		frame := newOutboundFrame(result == 3, timestamp, data)
+		outbound := newOutboundFrame(frame.Result == 3, frame.Timestamp, frame.Data)
 		for _, client := range clients {
-			client.offer(frame)
+			client.offer(outbound)
 		}
-
-		stream.GetFrameRateCounter().Update()
 	}
 }
 

server/service/stream/h264_source.go

@@ -0,0 +1,157 @@
+package stream
+
+import (
+	"NanoKVM-Server/common"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type H264Frame struct {
+	Data      []byte
+	Result    int
+	Duration  time.Duration
+	Timestamp int64
+}
+
+type H264Subscription struct {
+	frames chan H264Frame
+	done   chan struct{}
+	closed atomic.Bool
+	once   sync.Once
+}
+
+type H264Source struct {
+	mutex       sync.Mutex
+	subscribers map[*H264Subscription]struct{}
+	running     bool
+}
+
+var defaultH264Source = &H264Source{
+	subscribers: make(map[*H264Subscription]struct{}),
+}
+
+func SubscribeH264() *H264Subscription {
+	return defaultH264Source.subscribe()
+}
+
+func (s *H264Source) subscribe() *H264Subscription {
+	subscription := &H264Subscription{
+		frames: make(chan H264Frame, 4),
+		done:   make(chan struct{}),
+	}
+
+	s.mutex.Lock()
+	s.subscribers[subscription] = struct{}{}
+	start := !s.running
+	if start {
+		s.running = true
+	}
+	s.mutex.Unlock()
+
+	if start {
+		go s.run()
+	}
+
+	return subscription
+}
+
+func (s *H264Subscription) Next() (H264Frame, bool) {
+	select {
+	case frame := <-s.frames:
+		return frame, true
+	case <-s.done:
+		return H264Frame{}, false
+	}
+}
+
+func (s *H264Subscription) Close() {
+	s.once.Do(func() {
+		s.closed.Store(true)
+		close(s.done)
+		defaultH264Source.remove(s)
+	})
+}
+
+func (s *H264Source) remove(subscription *H264Subscription) {
+	s.mutex.Lock()
+	delete(s.subscribers, subscription)
+	s.mutex.Unlock()
+}
+
+func (s *H264Source) run() {
+	screen := common.GetScreen()
+	common.CheckScreen()
+	fps := screen.FPS
+	ticker := time.NewTicker(time.Second / time.Duration(fps))
+	defer ticker.Stop()
+
+	vision := common.GetKvmVision()
+	startTime := time.Now()
+
+	for range ticker.C {
+		subscribers := s.snapshot()
+		if len(subscribers) == 0 {
+			s.mutex.Lock()
+			if len(s.subscribers) == 0 {
+				s.running = false
+				s.mutex.Unlock()
+				return
+			}
+			s.mutex.Unlock()
+			continue
+		}
+
+		if screen.FPS != fps && screen.FPS != 0 {
+			fps = screen.FPS
+			ticker.Reset(time.Second / time.Duration(fps))
+		}
+
+		data, result := vision.ReadH264(screen.Width, screen.Height, screen.BitRate)
+		if result < 0 {
+			frame := H264Frame{Result: result}
+			for _, subscription := range subscribers {
+				subscription.send(frame)
+			}
+			continue
+		}
+		if len(data) == 0 {
+			continue
+		}
+
+		frame := H264Frame{
+			Data:      data,
+			Result:    result,
+			Duration:  time.Second / time.Duration(fps),
+			Timestamp: time.Since(startTime).Microseconds(),
+		}
+		for _, subscription := range subscribers {
+			subscription.send(frame)
+		}
+		GetFrameRateCounter().Update()
+	}
+}
+
+func (s *H264Source) snapshot() []*H264Subscription {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	subscribers := make([]*H264Subscription, 0, len(s.subscribers))
+	for subscription := range s.subscribers {
+		subscribers = append(subscribers, subscription)
+	}
+	return subscribers
+}
+
+func (s *H264Subscription) send(frame H264Frame) bool {
+	if s.closed.Load() {
+		return false
+	}
+
+	select {
+	case s.frames <- frame:
+		return true
+	case <-s.done:
+		return false
+	}
+}

server/service/stream/webrtc/client.go

@@ -2,17 +2,16 @@ package webrtc
 
 import (
 	"encoding/json"
-	"errors"
 
 	"github.com/gorilla/websocket"
-	"github.com/pion/rtp"
-	"github.com/pion/rtp/codecs"
 	"github.com/pion/webrtc/v4"
 	log "github.com/sirupsen/logrus"
 
 	"sync"
 )
 
+const h264SDPFmtpLine = "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=4d001f"
+
 func NewClient(ws *websocket.Conn, videoConn *webrtc.PeerConnection) *Client {
 	return &Client{
 		ws:    ws,
@@ -72,7 +71,11 @@ func (c *Client) ReadMessage() (*Message, error) {
 func (c *Client) AddTrack() error {
 	// video track
 	videoTrack, err := webrtc.NewTrackLocalStaticRTP(
-		webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeH264},
+		webrtc.RTPCodecCapability{
+			MimeType:    webrtc.MimeTypeH264,
+			ClockRate:   90000,
+			SDPFmtpLine: h264SDPFmtpLine,
+		},
 		"video",
 		"pion-video",
 	)
@@ -81,20 +84,6 @@ func (c *Client) AddTrack() error {
 		return err
 	}
 
-	videoPacketizer := rtp.NewPacketizer(
-		1200,
-		100,
-		0x1234ABCD,
-		&codecs.H264Payloader{},
-		rtp.NewRandomSequencer(),
-		90000,
-	)
-	if videoPacketizer == nil {
-		err := errors.New("failed to create rtp packetizer")
-		log.Error(err)
-		return err
-	}
-
 	videoSender, err := c.video.AddTrack(videoTrack)
 	if err != nil {
 		log.Errorf("failed to add video track: %s", err)
@@ -103,10 +92,8 @@ func (c *Client) AddTrack() error {
 	go startRTCPReader(videoSender)
 
 	track := &Track{
-		videoPacketizer: videoPacketizer,
-		video:           videoTrack,
+		video: videoTrack,
 	}
-	track.updateExtension()
 
 	c.mutex.Lock()
 	c.track = track

server/service/stream/webrtc/h264.go

@@ -10,6 +10,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
 	"github.com/pion/dtls/v3"
+	"github.com/pion/interceptor"
 	"github.com/pion/webrtc/v4"
 	log "github.com/sirupsen/logrus"
 )
@@ -151,14 +152,6 @@ func createMediaEngine() (*webrtc.MediaEngine, error) {
 		return nil, err
 	}
 
-	if err := mediaEngine.RegisterHeaderExtension(
-		webrtc.RTPHeaderExtensionCapability{URI: "http://www.webrtc.org/experiments/rtp-hdrext/playout-delay"},
-		webrtc.RTPCodecTypeVideo,
-	); err != nil {
-		log.Errorf("failed to register header extension: %s", err)
-		return nil, err
-	}
-
 	return mediaEngine, nil
 }
 
@@ -171,6 +164,7 @@ func createPeerConnection(iceServers []webrtc.ICEServer, mediaEngine *webrtc.Med
 
 	apiOptions := []func(api *webrtc.API){
 		webrtc.WithSettingEngine(settingEngine),
+		webrtc.WithInterceptorRegistry(&interceptor.Registry{}),
 	}
 	if mediaEngine != nil {
 		apiOptions = append(apiOptions, webrtc.WithMediaEngine(mediaEngine))

server/service/stream/webrtc/manager.go

@@ -1,19 +1,27 @@
 package webrtc
 
 import (
-	"NanoKVM-Server/common"
 	"NanoKVM-Server/service/stream"
 	"NanoKVM-Server/service/vm"
-	"time"
 
 	"github.com/gorilla/websocket"
+	"github.com/pion/rtp"
+	"github.com/pion/rtp/codecs"
 	"github.com/pion/webrtc/v4/pkg/media"
 	log "github.com/sirupsen/logrus"
 )
 
 func NewWebRTCManager() *WebRTCManager {
 	m := &WebRTCManager{
-		clients:      make(map[*websocket.Conn]*Client),
+		clients: make(map[*websocket.Conn]*Client),
+		videoPacketizer: rtp.NewPacketizer(
+			1450,
+			100,
+			0x1234ABCD,
+			&codecs.H264Payloader{},
+			rtp.NewRandomSequencer(),
+			90000,
+		),
 		videoSending: false,
 	}
 	m.updateClientSnapshotLocked()
@@ -22,8 +30,6 @@ func NewWebRTCManager() *WebRTCManager {
 }
 
 func (m *WebRTCManager) AddClient(ws *websocket.Conn, client *Client) {
-	client.track.updateExtension()
-
 	m.mutex.Lock()
 	m.clients[ws] = client
 	count := m.updateClientSnapshotLocked()
@@ -96,52 +102,62 @@ func (m *WebRTCManager) stopVideoStreamIfIdle() bool {
 }
 
 func (m *WebRTCManager) sendVideoStream() {
-	screen := common.GetScreen()
-	common.CheckScreen()
-	fps := screen.FPS
-	duration := time.Second / time.Duration(fps)
+	subscription := stream.SubscribeH264()
+	defer subscription.Close()
+	samples, writerDone := m.startVideoWriter()
 
-	vision := common.GetKvmVision()
-
-	ticker := time.NewTicker(duration)
-	defer ticker.Stop()
-
-	for range ticker.C {
+	for {
+		frame, ok := subscription.Next()
+		if !ok {
+			close(samples)
+			<-writerDone
+			return
+		}
 		clients := m.getClients()
 		if len(clients) == 0 {
+			close(samples)
+			<-writerDone
 			if m.stopVideoStreamIfIdle() {
 				log.Debugf("stop sending h264 stream")
 				return
 			}
+			samples, writerDone = m.startVideoWriter()
 
 			continue
 		}
 
-		data, result := vision.ReadH264(screen.Width, screen.Height, screen.BitRate)
-		stream.UpdateCaptureStatus(stream.CaptureModeH264, result)
-		if result < 0 || len(data) == 0 {
+		stream.UpdateCaptureStatus(stream.CaptureModeH264, frame.Result)
+		if frame.Result < 0 || len(frame.Data) == 0 {
 			continue
 		}
 
 		sample := media.Sample{
-			Data:     data,
-			Duration: duration,
+			Data:     frame.Data,
+			Duration: frame.Duration,
 		}
 
-		for _, client := range clients {
-			if err := client.track.writeVideoSample(sample); err != nil {
-				log.Errorf("failed to write h264 video to client: %s", err)
-				m.RemoveClient(client.WsConn())
-				client.Close()
-			}
-		}
-
-		if screen.FPS != fps && screen.FPS != 0 {
-			fps = screen.FPS
-			duration = time.Second / time.Duration(fps)
-			ticker.Reset(duration)
-		}
-
-		stream.GetFrameRateCounter().Update()
+		samples <- sample
 	}
 }
+
+func (m *WebRTCManager) startVideoWriter() (chan media.Sample, <-chan struct{}) {
+	samples := make(chan media.Sample, 1)
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+		for sample := range samples {
+			packets := m.videoPacketizer.Packetize(sample.Data, uint32(sample.Duration.Seconds()*90000))
+			for _, client := range m.getClients() {
+				err := client.track.writeVideoPackets(packets)
+				if err != nil {
+					log.Errorf("failed to write h264 video to client: %s", err)
+					m.RemoveClient(client.WsConn())
+					client.Close()
+				}
+			}
+		}
+	}()
+
+	return samples, done
+}

server/service/stream/webrtc/signaling.go

@@ -157,11 +157,6 @@ func (s *SignalingHandler) handleVideoOffer(data string) error {
 		return err
 	}
 
-	if err := s.updateHeaderExtensionID(); err != nil {
-		log.Errorf("could not update header extension ID: %v", err)
-		return err
-	}
-
 	answerByte, err := json.Marshal(answer)
 	if err != nil {
 		log.Errorf("failed to marshal answer: %s", err)
@@ -171,30 +166,6 @@ func (s *SignalingHandler) handleVideoOffer(data string) error {
 	return s.client.WriteMessage("video-answer", string(answerByte))
 }
 
-// set extension ID
-func (s *SignalingHandler) updateHeaderExtensionID() error {
-	receivers := s.client.video.GetReceivers()
-	if len(receivers) == 0 {
-		return errors.New("no RTP receiver found for video")
-	}
-
-	params := receivers[0].GetParameters()
-	if len(params.HeaderExtensions) == 0 {
-		return errors.New("no header extensions found in negotiated parameters")
-	}
-
-	for _, ext := range params.HeaderExtensions {
-		if ext.URI == "http://www.webrtc.org/experiments/rtp-hdrext/playout-delay" {
-			s.client.track.playoutDelayExtensionID = uint8(ext.ID)
-			log.Debugf("found and set playout delay extension ID to: %d", ext.ID)
-			return nil
-		}
-	}
-
-	log.Warnf("no track extension found in negotiated parameters, use default value 5")
-	return nil
-}
-
 // handle video candidate
 func (s *SignalingHandler) handleVideoCandidate(data string) error {
 	candidate := webrtc.ICECandidateInit{}

server/service/stream/webrtc/track.go

@@ -2,42 +2,11 @@ package webrtc
 
 import (
 	"github.com/pion/rtp"
-	"github.com/pion/webrtc/v4/pkg/media"
-	log "github.com/sirupsen/logrus"
 )
 
-func (t *Track) updateExtension() {
-	if t.playoutDelayExtensionID == 0 {
-		t.playoutDelayExtensionID = 5
-	}
-
-	if len(t.playoutDelayExtensionData) == 0 {
-		playoutDelay := &rtp.PlayoutDelayExtension{
-			MinDelay: 0,
-			MaxDelay: 0,
-		}
-		playoutDelayExtensionData, err := playoutDelay.Marshal()
-		if err == nil {
-			t.playoutDelayExtensionData = playoutDelayExtensionData
-		}
-	}
-}
-
-func (t *Track) writeVideoSample(sample media.Sample) error {
-	samples := uint32(sample.Duration.Seconds() * 90000)
-	packets := t.videoPacketizer.Packetize(sample.Data, samples)
-
-	for _, p := range packets {
-		p.Header.Extension = true
-		p.Header.ExtensionProfile = 0xBEDE
-
-		if err := p.Header.SetExtension(t.playoutDelayExtensionID, t.playoutDelayExtensionData); err != nil {
-			log.Errorf("Failed to set extension: %v", err)
-			return err
-		}
-
-		if err := t.video.WriteRTP(p); err != nil {
-			log.Errorf("failed to write RTP: %v", err)
+func (t *Track) writeVideoPackets(packets []*rtp.Packet) error {
+	for _, packet := range packets {
+		if err := t.video.WriteRTP(packet); err != nil {
 			return err
 		}
 	}

server/service/stream/webrtc/types.go

@@ -10,11 +10,12 @@ import (
 )
 
 type WebRTCManager struct {
-	clients        map[*websocket.Conn]*Client
-	clientSnapshot atomic.Pointer[[]*Client]
-	videoSending   bool
-	mutex          sync.Mutex
-	viewerVersion  uint64
+	clients         map[*websocket.Conn]*Client
+	clientSnapshot  atomic.Pointer[[]*Client]
+	videoPacketizer rtp.Packetizer
+	videoSending    bool
+	mutex           sync.Mutex
+	viewerVersion   uint64
 }
 
 type Client struct {
@@ -36,10 +37,7 @@ type SignalingHandler struct {
 }
 
 type Track struct {
-	playoutDelayExtensionID   uint8
-	playoutDelayExtensionData []byte
-	videoPacketizer           rtp.Packetizer
-	video                     *webrtc.TrackLocalStaticRTP
+	video *webrtc.TrackLocalStaticRTP
 }
 
 type Message struct {