fix: warn about mixed H264 stream modes (#843)

server/service/stream/direct/h264.go

@@ -1,6 +1,7 @@
 package direct
 
 import (
+	"NanoKVM-Server/service/stream"
 	"net/http"
 	"time"
 
@@ -36,6 +37,9 @@ func Connect(c *gin.Context) {
 	streamer.addClient(ws)
 	defer streamer.removeClient(ws)
 
+	unregisterMode := stream.RegisterH264Mode(stream.H264ModeDirect)
+	defer unregisterMode()
+
 	for {
 		if _, _, err := ws.NextReader(); err != nil {
 			log.Debugf("failed to read message (client disconnected): %s", err)

server/service/stream/direct/streamer.go

@@ -152,6 +152,7 @@ func (s *Streamer) send(clients []*websocket.Conn, isKeyFrame byte, timestamp in
 			log.Errorf("failed to write message to client %s: %s.", client.RemoteAddr(), err)
 
 			s.removeClient(client)
+			_ = client.Close()
 		}
 	}
 

server/service/stream/h264_mode.go

@@ -0,0 +1,172 @@
+package stream
+
+import (
+	"sync"
+
+	"github.com/google/uuid"
+)
+
+const (
+	H264ModeStatusEvent = "h264-mode-status"
+	H264ModeDirect      = "direct"
+	H264ModeWebRTC      = "webrtc"
+)
+
+type H264ModeStatus struct {
+	Generation string `json:"generation"`
+	Revision   uint64 `json:"revision"`
+	Direct     int    `json:"direct"`
+	WebRTC     int    `json:"webrtc"`
+	Mixed      bool   `json:"mixed"`
+}
+
+type H264ModeStatusSubscriber func(H264ModeStatus)
+
+var defaultH264ModeStore = newH264ModeStore()
+
+func RegisterH264Mode(mode string) func() {
+	return defaultH264ModeStore.register(mode)
+}
+
+func SubscribeH264ModeStatus(subscriber H264ModeStatusSubscriber) func() {
+	return defaultH264ModeStore.subscribe(subscriber)
+}
+
+func CurrentH264ModeStatus() H264ModeStatus {
+	return defaultH264ModeStore.current()
+}
+
+type h264ModeStore struct {
+	mutex       sync.Mutex
+	counts      map[string]int
+	status      H264ModeStatus
+	subscribers map[int]H264ModeStatusSubscriber
+	nextID      int
+	pending     []h264ModeStatusNotification
+	notify      chan struct{}
+}
+
+type h264ModeStatusNotification struct {
+	status      H264ModeStatus
+	subscribers []H264ModeStatusSubscriber
+}
+
+func newH264ModeStore() *h264ModeStore {
+	store := &h264ModeStore{
+		counts:      make(map[string]int),
+		subscribers: make(map[int]H264ModeStatusSubscriber),
+		status: H264ModeStatus{
+			Generation: uuid.NewString(),
+		},
+		notify: make(chan struct{}, 1),
+	}
+	go store.run()
+
+	return store
+}
+
+func (s *h264ModeStore) register(mode string) func() {
+	if !isH264Mode(mode) {
+		return func() {}
+	}
+
+	s.update(mode, 1)
+
+	var once sync.Once
+	return func() {
+		once.Do(func() {
+			s.update(mode, -1)
+		})
+	}
+}
+
+func (s *h264ModeStore) subscribe(subscriber H264ModeStatusSubscriber) func() {
+	s.mutex.Lock()
+	id := s.nextID
+	s.nextID++
+	s.subscribers[id] = subscriber
+	s.mutex.Unlock()
+
+	return func() {
+		s.mutex.Lock()
+		delete(s.subscribers, id)
+		s.mutex.Unlock()
+	}
+}
+
+func (s *h264ModeStore) current() H264ModeStatus {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	return s.status
+}
+
+func (s *h264ModeStore) update(mode string, delta int) {
+	s.mutex.Lock()
+	previous := s.status
+	s.counts[mode] += delta
+	if s.counts[mode] < 0 {
+		s.counts[mode] = 0
+	}
+
+	next := H264ModeStatus{
+		Generation: previous.Generation,
+		Revision:   previous.Revision + 1,
+		Direct:     s.counts[H264ModeDirect],
+		WebRTC:     s.counts[H264ModeWebRTC],
+	}
+	next.Mixed = next.Direct > 0 && next.WebRTC > 0
+	s.status = next
+
+	if previous.Mixed == next.Mixed {
+		s.mutex.Unlock()
+		return
+	}
+
+	subscribers := make([]H264ModeStatusSubscriber, 0, len(s.subscribers))
+	for _, subscriber := range s.subscribers {
+		subscribers = append(subscribers, subscriber)
+	}
+	s.pending = append(s.pending, h264ModeStatusNotification{
+		status:      next,
+		subscribers: subscribers,
+	})
+	s.mutex.Unlock()
+
+	select {
+	case s.notify <- struct{}{}:
+	default:
+	}
+}
+
+func (s *h264ModeStore) run() {
+	for range s.notify {
+		for {
+			notification, ok := s.takePending()
+			if !ok {
+				break
+			}
+
+			for _, subscriber := range notification.subscribers {
+				subscriber(notification.status)
+			}
+		}
+	}
+}
+
+func (s *h264ModeStore) takePending() (h264ModeStatusNotification, bool) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	if len(s.pending) == 0 {
+		return h264ModeStatusNotification{}, false
+	}
+
+	notification := s.pending[0]
+	s.pending = s.pending[1:]
+	return notification, true
+}
+
+func isH264Mode(mode string) bool {
+	return mode == H264ModeDirect || mode == H264ModeWebRTC
+}

server/service/stream/webrtc/h264.go

@@ -74,12 +74,9 @@ func Connect(c *gin.Context) {
 		return
 	}
 
-	manager := getManager()
-	manager.AddClient(wsConn, client)
-	defer manager.RemoveClient(wsConn)
-
 	// handle signaling
 	signalingHandler := NewSignalingHandler(client)
+	defer signalingHandler.Close()
 	signalingHandler.RegisterCallbacks()
 	if err := sendICEServers(client, iceServers); err != nil {
 		log.Errorf("failed to send ICE servers: %s", err)

server/service/stream/webrtc/manager.go

@@ -4,7 +4,6 @@ import (
 	"NanoKVM-Server/common"
 	"NanoKVM-Server/service/stream"
 	"NanoKVM-Server/service/vm"
-	"sync/atomic"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -15,7 +14,7 @@ import (
 func NewWebRTCManager() *WebRTCManager {
 	m := &WebRTCManager{
 		clients:      make(map[*websocket.Conn]*Client),
-		videoSending: 0,
+		videoSending: false,
 	}
 	m.updateClientSnapshotLocked()
 
@@ -72,15 +71,31 @@ func (m *WebRTCManager) getClients() []*Client {
 }
 
 func (m *WebRTCManager) StartVideoStream() {
-	if atomic.CompareAndSwapInt32(&m.videoSending, 0, 1) {
-		go m.sendVideoStream()
-		log.Debugf("start sending h264 stream")
+	m.mutex.Lock()
+	if m.videoSending || len(m.clients) == 0 {
+		m.mutex.Unlock()
+		return
 	}
+	m.videoSending = true
+	m.mutex.Unlock()
+
+	go m.sendVideoStream()
+	log.Debugf("start sending h264 stream")
+}
+
+func (m *WebRTCManager) stopVideoStreamIfIdle() bool {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if len(m.clients) > 0 {
+		return false
+	}
+
+	m.videoSending = false
+	return true
 }
 
 func (m *WebRTCManager) sendVideoStream() {
-	defer atomic.StoreInt32(&m.videoSending, 0)
-
 	screen := common.GetScreen()
 	common.CheckScreen()
 	fps := screen.FPS
@@ -94,8 +109,12 @@ func (m *WebRTCManager) sendVideoStream() {
 	for range ticker.C {
 		clients := m.getClients()
 		if len(clients) == 0 {
-			log.Debugf("stop sending h264 stream")
-			return
+			if m.stopVideoStreamIfIdle() {
+				log.Debugf("stop sending h264 stream")
+				return
+			}
+
+			continue
 		}
 
 		data, result := vision.ReadH264(screen.Width, screen.Height, screen.BitRate)

server/service/stream/webrtc/signaling.go

@@ -1,6 +1,7 @@
 package webrtc
 
 import (
+	"NanoKVM-Server/service/stream"
 	"encoding/json"
 	"errors"
 
@@ -33,18 +34,85 @@ func (s *SignalingHandler) RegisterCallbacks() {
 		}
 	})
 
-	manager := getManager()
-
 	// video connection state change
 	s.client.video.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
-		if state == webrtc.ICEConnectionStateConnected {
-			manager.StartVideoStream()
-		}
+		s.updateVideoStreamState(state)
 
 		log.Debugf("video connection state changed to %s", state.String())
 	})
 }
 
+func (s *SignalingHandler) Close() {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	if s.closed {
+		return
+	}
+
+	s.closed = true
+	s.unregisterH264ModeLocked()
+	if s.client != nil && s.client.WsConn() != nil {
+		getManager().RemoveClient(s.client.WsConn())
+	}
+}
+
+func (s *SignalingHandler) updateVideoStreamState(state webrtc.ICEConnectionState) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	if s.closed {
+		return
+	}
+
+	manager := getManager()
+	if s.updateH264ModeLocked(state) {
+		manager.AddClient(s.client.WsConn(), s.client)
+		manager.StartVideoStream()
+		return
+	}
+
+	manager.RemoveClient(s.client.WsConn())
+}
+
+func (s *SignalingHandler) updateH264Mode(state webrtc.ICEConnectionState) bool {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	if s.closed {
+		return false
+	}
+
+	return s.updateH264ModeLocked(state)
+}
+
+func (s *SignalingHandler) updateH264ModeLocked(state webrtc.ICEConnectionState) bool {
+	if state == webrtc.ICEConnectionStateConnected || state == webrtc.ICEConnectionStateCompleted {
+		s.registerH264ModeLocked()
+		return true
+	}
+
+	s.unregisterH264ModeLocked()
+	return false
+}
+
+func (s *SignalingHandler) registerH264ModeLocked() {
+	if s.unregisterMode != nil {
+		return
+	}
+
+	s.unregisterMode = stream.RegisterH264Mode(stream.H264ModeWebRTC)
+}
+
+func (s *SignalingHandler) unregisterH264ModeLocked() {
+	unregisterMode := s.unregisterMode
+	s.unregisterMode = nil
+
+	if unregisterMode != nil {
+		unregisterMode()
+	}
+}
+
 // HandleMessage handle the received message
 func (s *SignalingHandler) HandleMessage(message *Message) error {
 	switch message.Event {

server/service/stream/webrtc/types.go

@@ -12,7 +12,7 @@ import (
 type WebRTCManager struct {
 	clients        map[*websocket.Conn]*Client
 	clientSnapshot atomic.Pointer[[]*Client]
-	videoSending   int32
+	videoSending   bool
 	mutex          sync.Mutex
 	viewerVersion  uint64
 }
@@ -29,7 +29,10 @@ func (c *Client) WsConn() *websocket.Conn {
 }
 
 type SignalingHandler struct {
-	client *Client
+	client         *Client
+	mutex          sync.Mutex
+	unregisterMode func()
+	closed         bool
 }
 
 type Track struct {

server/service/ws/h264_mode.go

@@ -0,0 +1,93 @@
+package ws
+
+import (
+	"encoding/json"
+	"sync"
+
+	"NanoKVM-Server/service/stream"
+	log "github.com/sirupsen/logrus"
+)
+
+var h264ModeStatusBroadcasterInstance = newH264ModeStatusBroadcaster(broadcastH264ModeStatus)
+
+func init() {
+	go h264ModeStatusBroadcasterInstance.Run()
+
+	stream.SubscribeH264ModeStatus(func(status stream.H264ModeStatus) {
+		h264ModeStatusBroadcasterInstance.Enqueue(status)
+	})
+}
+
+type h264ModeStatusBroadcaster struct {
+	mutex     sync.Mutex
+	pending   *stream.H264ModeStatus
+	notify    chan struct{}
+	broadcast func(stream.H264ModeStatus)
+}
+
+func newH264ModeStatusBroadcaster(broadcast func(stream.H264ModeStatus)) *h264ModeStatusBroadcaster {
+	return &h264ModeStatusBroadcaster{
+		notify:    make(chan struct{}, 1),
+		broadcast: broadcast,
+	}
+}
+
+func (b *h264ModeStatusBroadcaster) Enqueue(status stream.H264ModeStatus) {
+	b.mutex.Lock()
+	b.pending = &status
+	b.mutex.Unlock()
+
+	select {
+	case b.notify <- struct{}{}:
+	default:
+	}
+}
+
+func (b *h264ModeStatusBroadcaster) Run() {
+	for range b.notify {
+		for {
+			status, ok := b.takePending()
+			if !ok {
+				break
+			}
+
+			b.broadcast(status)
+		}
+	}
+}
+
+func (b *h264ModeStatusBroadcaster) takePending() (stream.H264ModeStatus, bool) {
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+
+	if b.pending == nil {
+		return stream.H264ModeStatus{}, false
+	}
+
+	status := *b.pending
+	b.pending = nil
+	return status, true
+}
+
+func sendH264ModeStatusSnapshot(client *Client) {
+	if err := sendH264ModeStatus(client, stream.CurrentH264ModeStatus()); err != nil {
+		log.Errorf("failed to send h264 mode status snapshot: %s", err)
+	}
+}
+
+func broadcastH264ModeStatus(status stream.H264ModeStatus) {
+	for _, client := range GetManager().GetClients() {
+		if err := sendH264ModeStatus(client, status); err != nil {
+			log.Errorf("failed to send h264 mode status: %s", err)
+		}
+	}
+}
+
+func sendH264ModeStatus(client *Client, status stream.H264ModeStatus) error {
+	payload, err := json.Marshal(status)
+	if err != nil {
+		return err
+	}
+
+	return client.Write(stream.H264ModeStatusEvent, string(payload))
+}

server/service/ws/service.go

@@ -38,6 +38,7 @@ func (s *Service) Connect(c *gin.Context) {
 	defer manager.RemoveClient(ws)
 
 	sendCaptureStatusSnapshot(client)
+	sendH264ModeStatusSnapshot(client)
 
 	client.Start()
 }