perf: improve H264 streaming client handling

Use client snapshots for stream fanout, reduce frame queue latency, expose backend ICE server configuration to the WebRTC client, and clean up disconnected WebRTC clients more aggressively.

server/service/stream/direct/h264.go

@@ -12,6 +12,7 @@ import (
 var (
 	streamer = newStreamer()
 	upgrader = websocket.Upgrader{
+		WriteBufferSize: 256 * 1024,
 		CheckOrigin: func(r *http.Request) bool {
 			return true
 		},
@@ -36,7 +37,7 @@ func Connect(c *gin.Context) {
 	defer streamer.removeClient(ws)
 
 	for {
-		if _, _, err := ws.ReadMessage(); err != nil {
+		if _, _, err := ws.NextReader(); err != nil {
 			log.Debugf("failed to read message (client disconnected): %s", err)
 			return
 		}

server/service/stream/direct/streamer.go

@@ -14,20 +14,25 @@ import (
 )
 
 type Streamer struct {
-	mutex   sync.RWMutex
-	clients map[*websocket.Conn]bool
-	running int32
+	mutex          sync.Mutex
+	clients        map[*websocket.Conn]bool
+	clientSnapshot atomic.Pointer[[]*websocket.Conn]
+	running        int32
 }
 
 func newStreamer() *Streamer {
-	return &Streamer{
+	s := &Streamer{
 		clients: make(map[*websocket.Conn]bool),
 	}
+	s.updateClientSnapshotLocked()
+
+	return s
 }
 
 func (s *Streamer) addClient(ws *websocket.Conn) {
 	s.mutex.Lock()
 	s.clients[ws] = true
+	s.updateClientSnapshotLocked()
 	s.mutex.Unlock()
 
 	if atomic.CompareAndSwapInt32(&s.running, 0, 1) {
@@ -39,31 +44,47 @@ func (s *Streamer) addClient(ws *websocket.Conn) {
 func (s *Streamer) removeClient(ws *websocket.Conn) {
 	s.mutex.Lock()
 	delete(s.clients, ws)
+	count := s.updateClientSnapshotLocked()
 	s.mutex.Unlock()
 
-	log.Debugf("h264 websocket disconnected, remaining clients: %d", len(s.clients))
+	log.Debugf("h264 websocket disconnected, remaining clients: %d", count)
 }
 
-func (s *Streamer) getClientCount() int {
-	s.mutex.RLock()
-	defer s.mutex.RUnlock()
+func (s *Streamer) updateClientSnapshotLocked() int {
+	clients := make([]*websocket.Conn, 0, len(s.clients))
+	for client := range s.clients {
+		clients = append(clients, client)
+	}
+	s.clientSnapshot.Store(&clients)
 
-	return len(s.clients)
+	return len(clients)
+}
+
+func (s *Streamer) getClients() []*websocket.Conn {
+	clients := s.clientSnapshot.Load()
+	if clients == nil {
+		return nil
+	}
+
+	return *clients
 }
 
 func (s *Streamer) run() {
 	defer atomic.StoreInt32(&s.running, 0)
 
-	duration := time.Second / time.Duration(120)
-	ticker := time.NewTicker(duration)
+	screen := common.GetScreen()
+	common.CheckScreen()
+	fps := screen.FPS
+
+	ticker := time.NewTicker(time.Second / time.Duration(fps))
 	defer ticker.Stop()
 
-	screen := common.GetScreen()
 	vision := common.GetKvmVision()
 	startTime := time.Now()
 
 	for range ticker.C {
-		if s.getClientCount() == 0 {
+		clients := s.getClients()
+		if len(clients) == 0 {
 			log.Debug("h264 stream stopped due to no clients")
 			return
 		}
@@ -81,28 +102,34 @@ func (s *Streamer) run() {
 
 		timestamp := time.Since(startTime).Microseconds()
 
-		if err := s.send(isKeyFrame, timestamp, data); err != nil {
+		if err := s.send(clients, isKeyFrame, timestamp, data); err != nil {
 			continue
 		}
 
+		if screen.FPS != fps && screen.FPS != 0 {
+			fps = screen.FPS
+			ticker.Reset(time.Second / time.Duration(fps))
+		}
+
 		stream.GetFrameRateCounter().Update()
 	}
 }
 
-func (s *Streamer) send(isKeyFrame byte, timestamp int64, data []byte) error {
+func (s *Streamer) send(clients []*websocket.Conn, isKeyFrame byte, timestamp int64, data []byte) error {
 	buf := BufferPool.Get().(*bytes.Buffer)
 	defer BufferPool.Put(buf)
 
 	buf.Reset()
+	buf.Grow(1 + 8 + len(data))
 
 	if err := buf.WriteByte(isKeyFrame); err != nil {
 		log.Errorf("failed to write keyframe flag: %s", err)
 		return err
 	}
 
-	tsBytes := make([]byte, 8)
-	binary.LittleEndian.PutUint64(tsBytes, uint64(timestamp))
-	if _, err := buf.Write(tsBytes); err != nil {
+	var tsBytes [8]byte
+	binary.LittleEndian.PutUint64(tsBytes[:], uint64(timestamp))
+	if _, err := buf.Write(tsBytes[:]); err != nil {
 		log.Errorf("failed to write timestamp: %s", err)
 		return err
 	}
@@ -112,7 +139,7 @@ func (s *Streamer) send(isKeyFrame byte, timestamp int64, data []byte) error {
 		return err
 	}
 
-	for client := range s.clients {
+	for _, client := range clients {
 		if err := client.WriteMessage(websocket.BinaryMessage, buf.Bytes()); err != nil {
 			log.Errorf("failed to write message to client %s: %s.", client.RemoteAddr(), err)
 

server/service/stream/frame_rate.go

@@ -16,9 +16,8 @@ var (
 )
 
 type FrameRateCounter struct {
-	frameCount int32
-	fps        int32
-	mutex      sync.Mutex
+	frameCount atomic.Int32
+	fps        atomic.Int32
 }
 
 func GetFrameRateCounter() *FrameRateCounter {
@@ -30,16 +29,10 @@ func GetFrameRateCounter() *FrameRateCounter {
 			defer ticker.Stop()
 
 			for range ticker.C {
-				counter.mutex.Lock()
+				currentCount := counter.frameCount.Swap(0)
+				counter.fps.Store(currentCount / 3)
 
-				currentCount := atomic.LoadInt32(&counter.frameCount)
-
-				counter.fps = currentCount / 3
-				atomic.StoreInt32(&counter.frameCount, 0)
-
-				counter.mutex.Unlock()
-
-				data := fmt.Sprintf("%d", counter.fps)
+				data := fmt.Sprintf("%d", counter.fps.Load())
 				err := os.WriteFile("/kvmapp/kvm/now_fps", []byte(data), 0o666)
 				if err != nil {
 					log.Errorf("failed to write fps: %s", err)
@@ -52,12 +45,9 @@ func GetFrameRateCounter() *FrameRateCounter {
 }
 
 func (f *FrameRateCounter) Update() {
-	atomic.AddInt32(&f.frameCount, 1)
+	f.frameCount.Add(1)
 }
 
 func (f *FrameRateCounter) GetFPS() int32 {
-	f.mutex.Lock()
-	defer f.mutex.Unlock()
-
-	return f.fps
+	return f.fps.Load()
 }

server/service/stream/mjpeg/streamer.go

@@ -13,24 +13,31 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+var crlf = []byte("\r\n")
+
 type Streamer struct {
-	mutex       sync.RWMutex
-	clients     map[*gin.Context]bool
-	running     int32
-	frameMutex  sync.RWMutex
-	latestFrame LatestFrame
-	cacheRefs   int32
+	mutex          sync.Mutex
+	clients        map[*gin.Context]bool
+	clientSnapshot atomic.Pointer[[]*gin.Context]
+	running        int32
+	frameMutex     sync.RWMutex
+	latestFrame    LatestFrame
+	cacheRefs      int32
 }
 
 func NewStreamer() *Streamer {
-	return &Streamer{
+	s := &Streamer{
 		clients: make(map[*gin.Context]bool),
 	}
+	s.updateClientSnapshotLocked()
+
+	return s
 }
 
 func (s *Streamer) AddClient(c *gin.Context) {
 	s.mutex.Lock()
 	s.clients[c] = true
+	s.updateClientSnapshotLocked()
 	s.mutex.Unlock()
 
 	if atomic.CompareAndSwapInt32(&s.running, 0, 1) {
@@ -42,28 +49,29 @@ func (s *Streamer) AddClient(c *gin.Context) {
 func (s *Streamer) RemoveClient(c *gin.Context) {
 	s.mutex.Lock()
 	delete(s.clients, c)
+	count := s.updateClientSnapshotLocked()
 	s.mutex.Unlock()
 
-	log.Debugf("mjpeg connection removed, remaining clients: %d", len(s.clients))
+	log.Debugf("mjpeg connection removed, remaining clients: %d", count)
 }
 
-func (s *Streamer) getClients() []*gin.Context {
-	s.mutex.RLock()
-	defer s.mutex.RUnlock()
-
+func (s *Streamer) updateClientSnapshotLocked() int {
 	clients := make([]*gin.Context, 0, len(s.clients))
 	for c := range s.clients {
 		clients = append(clients, c)
 	}
+	s.clientSnapshot.Store(&clients)
 
-	return clients
+	return len(clients)
 }
 
-func (s *Streamer) getClientCount() int {
-	s.mutex.RLock()
-	defer s.mutex.RUnlock()
+func (s *Streamer) getClients() []*gin.Context {
+	clients := s.clientSnapshot.Load()
+	if clients == nil {
+		return nil
+	}
 
-	return len(s.clients)
+	return *clients
 }
 
 func (s *Streamer) run() {
@@ -79,7 +87,8 @@ func (s *Streamer) run() {
 	defer ticker.Stop()
 
 	for range ticker.C {
-		if s.getClientCount() == 0 {
+		clients := s.getClients()
+		if len(clients) == 0 {
 			log.Debug("mjpeg stream stopped due to no clients")
 			return
 		}
@@ -94,7 +103,6 @@ func (s *Streamer) run() {
 			s.setLatestFrame(data, screen.Width, screen.Height)
 		}
 
-		clients := s.getClients()
 		for _, client := range clients {
 			if err := writeFrame(client, data); err != nil {
 				log.Errorf("failed to write mjpeg frame for client %s: %s", client.Request.RemoteAddr, err)
@@ -195,7 +203,7 @@ func writeFrame(c *gin.Context, data []byte) (err error) {
 		return err
 	}
 
-	if _, err = c.Writer.Write([]byte("\r\n")); err != nil {
+	if _, err = c.Writer.Write(crlf); err != nil {
 		return err
 	}
 

server/service/stream/webrtc/client.go

@@ -21,6 +21,20 @@ func NewClient(ws *websocket.Conn, videoConn *webrtc.PeerConnection) *Client {
 	}
 }
 
+func (c *Client) Close() {
+	if c.video != nil {
+		if err := c.video.Close(); err != nil {
+			log.Debugf("failed to close video peer connection: %s", err)
+		}
+	}
+
+	if c.ws != nil {
+		if err := c.ws.Close(); err != nil {
+			log.Debugf("failed to close websocket: %s", err)
+		}
+	}
+}
+
 func (c *Client) WriteMessage(event string, data string) error {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()

server/service/stream/webrtc/h264.go

@@ -2,6 +2,7 @@ package webrtc
 
 import (
 	"NanoKVM-Server/config"
+	"encoding/json"
 	"net/http"
 	"sync"
 	"time"
@@ -15,6 +16,7 @@ import (
 
 var (
 	upgrader = websocket.Upgrader{
+		WriteBufferSize: 256 * 1024,
 		CheckOrigin: func(r *http.Request) bool {
 			return true
 		},
@@ -79,6 +81,10 @@ func Connect(c *gin.Context) {
 	// handle signaling
 	signalingHandler := NewSignalingHandler(client)
 	signalingHandler.RegisterCallbacks()
+	if err := sendICEServers(client, iceServers); err != nil {
+		log.Errorf("failed to send ICE servers: %s", err)
+		return
+	}
 
 	// read and wait
 	for {
@@ -116,6 +122,30 @@ func createICEServers() []webrtc.ICEServer {
 	return iceServers
 }
 
+type clientICEServer struct {
+	URLs       []string    `json:"urls"`
+	Username   string      `json:"username,omitempty"`
+	Credential interface{} `json:"credential,omitempty"`
+}
+
+func sendICEServers(client *Client, iceServers []webrtc.ICEServer) error {
+	clientServers := make([]clientICEServer, 0, len(iceServers))
+	for _, server := range iceServers {
+		clientServers = append(clientServers, clientICEServer{
+			URLs:       server.URLs,
+			Username:   server.Username,
+			Credential: server.Credential,
+		})
+	}
+
+	data, err := json.Marshal(clientServers)
+	if err != nil {
+		return err
+	}
+
+	return client.WriteMessage("ice-servers", string(data))
+}
+
 func createMediaEngine() (*webrtc.MediaEngine, error) {
 	mediaEngine := &webrtc.MediaEngine{}
 

server/service/stream/webrtc/manager.go

@@ -3,7 +3,6 @@ package webrtc
 import (
 	"NanoKVM-Server/common"
 	"NanoKVM-Server/service/stream"
-	"sync"
 	"sync/atomic"
 	"time"
 
@@ -13,11 +12,13 @@ import (
 )
 
 func NewWebRTCManager() *WebRTCManager {
-	return &WebRTCManager{
+	m := &WebRTCManager{
 		clients:      make(map[*websocket.Conn]*Client),
 		videoSending: 0,
-		mutex:        sync.RWMutex{},
 	}
+	m.updateClientSnapshotLocked()
+
+	return m
 }
 
 func (m *WebRTCManager) AddClient(ws *websocket.Conn, client *Client) {
@@ -25,24 +26,42 @@ func (m *WebRTCManager) AddClient(ws *websocket.Conn, client *Client) {
 
 	m.mutex.Lock()
 	m.clients[ws] = client
+	count := m.updateClientSnapshotLocked()
 	m.mutex.Unlock()
 
-	log.Debugf("added client %s, total clients: %d", ws.RemoteAddr(), len(m.clients))
+	log.Debugf("added client %s, total clients: %d", ws.RemoteAddr(), count)
 }
 
 func (m *WebRTCManager) RemoveClient(ws *websocket.Conn) {
 	m.mutex.Lock()
 	delete(m.clients, ws)
+	count := m.updateClientSnapshotLocked()
 	m.mutex.Unlock()
 
-	log.Debugf("removed client %s, total clients: %d", ws.RemoteAddr(), len(m.clients))
+	log.Debugf("removed client %s, total clients: %d", ws.RemoteAddr(), count)
 }
 
 func (m *WebRTCManager) GetClientCount() int {
-	m.mutex.RLock()
-	defer m.mutex.RUnlock()
+	return len(m.getClients())
+}
 
-	return len(m.clients)
+func (m *WebRTCManager) updateClientSnapshotLocked() int {
+	clients := make([]*Client, 0, len(m.clients))
+	for _, client := range m.clients {
+		clients = append(clients, client)
+	}
+	m.clientSnapshot.Store(&clients)
+
+	return len(clients)
+}
+
+func (m *WebRTCManager) getClients() []*Client {
+	clients := m.clientSnapshot.Load()
+	if clients == nil {
+		return nil
+	}
+
+	return *clients
 }
 
 func (m *WebRTCManager) StartVideoStream() {
@@ -66,7 +85,8 @@ func (m *WebRTCManager) sendVideoStream() {
 	defer ticker.Stop()
 
 	for range ticker.C {
-		if m.GetClientCount() == 0 {
+		clients := m.getClients()
+		if len(clients) == 0 {
 			log.Debugf("stop sending h264 stream")
 			return
 		}
@@ -82,8 +102,12 @@ func (m *WebRTCManager) sendVideoStream() {
 			Duration: duration,
 		}
 
-		for _, client := range m.clients {
-			client.track.writeVideo(sample)
+		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 {

server/service/stream/webrtc/track.go

@@ -11,7 +11,7 @@ func (t *Track) updateExtension() {
 		t.playoutDelayExtensionID = 5
 	}
 
-	if t.playoutDelayExtensionData == nil || len(t.playoutDelayExtensionData) == 0 {
+	if len(t.playoutDelayExtensionData) == 0 {
 		playoutDelay := &rtp.PlayoutDelayExtension{
 			MinDelay: 0,
 			MaxDelay: 0,
@@ -44,10 +44,3 @@ func (t *Track) writeVideoSample(sample media.Sample) error {
 
 	return nil
 }
-
-func (t *Track) writeVideo(sample media.Sample) {
-	err := t.writeVideoSample(sample)
-	if err != nil {
-		log.Errorf("failed to write h264 video: %s", err)
-	}
-}

server/service/stream/webrtc/types.go

@@ -2,6 +2,7 @@ package webrtc
 
 import (
 	"sync"
+	"sync/atomic"
 
 	"github.com/gorilla/websocket"
 	"github.com/pion/rtp"
@@ -9,9 +10,10 @@ import (
 )
 
 type WebRTCManager struct {
-	clients      map[*websocket.Conn]*Client
-	videoSending int32
-	mutex        sync.RWMutex
+	clients        map[*websocket.Conn]*Client
+	clientSnapshot atomic.Pointer[[]*Client]
+	videoSending   int32
+	mutex          sync.Mutex
 }
 
 type Client struct {
@@ -21,6 +23,10 @@ type Client struct {
 	mutex sync.Mutex
 }
 
+func (c *Client) WsConn() *websocket.Conn {
+	return c.ws
+}
+
 type SignalingHandler struct {
 	client *Client
 }