perf: reduce latency in 60 Hz mode (#844)

- Add decode-driven flow control and bounded GOP-aware queues to prevent stale H.264 frames from accumulating in TCP and WebSocket buffers.
- Isolate slow clients with dedicated writer goroutines, write deadlines, and safer connection lifecycle handling.
- Move the Direct H.264 socket into the worker, simplify the low-latency decode path, and improve reconnect behavior.

server/service/stream/direct/client.go

@@ -0,0 +1,348 @@
+package direct
+
+import (
+	"encoding/binary"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	frameAckMessage     byte = 2
+	streamResyncMessage byte = 3
+
+	defaultQueueFrames = 8
+	defaultQueueBytes  = 2 * 1024 * 1024
+	maxFlowWindow      = 8
+	writeWait          = 2 * time.Second
+	pingPeriod         = 15 * time.Second
+)
+
+type outboundFrame struct {
+	key       bool
+	timestamp int64
+	payload   []byte
+}
+
+type frameQueue struct {
+	mutex sync.Mutex
+	wake  chan struct{}
+
+	frames      []*outboundFrame
+	queuedBytes int
+	inFlight    []int64
+
+	maxFrames int
+	maxBytes  int
+	window    int
+
+	flowControlled     bool
+	waitingForKeyframe bool
+	closed             bool
+}
+
+func newFrameQueue(maxFrames int, maxBytes int) *frameQueue {
+	return &frameQueue{
+		wake:               make(chan struct{}, 1),
+		maxFrames:          maxFrames,
+		maxBytes:           maxBytes,
+		waitingForKeyframe: true,
+	}
+}
+
+func (q *frameQueue) enableFlowControl(window int) {
+	if window < 1 {
+		window = 1
+	}
+	if window > maxFlowWindow {
+		window = maxFlowWindow
+	}
+
+	q.mutex.Lock()
+	q.clearFramesLocked()
+	q.inFlight = q.inFlight[:0]
+	q.window = window
+	q.flowControlled = true
+	q.waitingForKeyframe = true
+	q.mutex.Unlock()
+}
+
+func (q *frameQueue) canAdvanceStream() bool {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	if q.closed {
+		return false
+	}
+	if !q.flowControlled {
+		return true
+	}
+
+	if q.waitingForKeyframe {
+		return len(q.inFlight) < q.window
+	}
+
+	return len(q.frames)+len(q.inFlight) < q.window
+}
+
+func (q *frameQueue) captureState() (active bool, flowControlled bool, canAdvance bool) {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	if q.closed {
+		return false, q.flowControlled, false
+	}
+	if !q.flowControlled {
+		return true, false, true
+	}
+	if q.waitingForKeyframe {
+		return true, true, len(q.inFlight) < q.window
+	}
+
+	return true, true, len(q.frames)+len(q.inFlight) < q.window
+}
+
+func (q *frameQueue) offer(frame *outboundFrame) bool {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	if q.closed {
+		return false
+	}
+
+	if frame.key {
+		q.clearFramesLocked()
+		if q.flowControlled && len(q.inFlight) >= q.window {
+			q.waitingForKeyframe = true
+			return false
+		}
+		q.waitingForKeyframe = false
+		q.pushLocked(frame)
+		q.signalLocked()
+		return true
+	}
+
+	if q.waitingForKeyframe {
+		return false
+	}
+
+	if q.flowControlled {
+		if len(q.frames)+len(q.inFlight) >= q.window {
+			q.clearFramesLocked()
+			q.waitingForKeyframe = true
+			return false
+		}
+	} else if len(q.frames) >= q.maxFrames || q.queuedBytes+len(frame.payload) > q.maxBytes {
+		q.clearFramesLocked()
+		q.waitingForKeyframe = true
+		return false
+	}
+
+	q.pushLocked(frame)
+	q.signalLocked()
+	return true
+}
+
+func (q *frameQueue) popForWrite() *outboundFrame {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	if q.closed || len(q.frames) == 0 {
+		return nil
+	}
+
+	frame := q.frames[0]
+	q.frames[0] = nil
+	q.frames = q.frames[1:]
+	q.queuedBytes -= len(frame.payload)
+	if q.flowControlled {
+		q.inFlight = append(q.inFlight, frame.timestamp)
+	}
+
+	return frame
+}
+
+func (q *frameQueue) acknowledge(timestamp int64) {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	acknowledged := 0
+	for acknowledged < len(q.inFlight) && q.inFlight[acknowledged] <= timestamp {
+		acknowledged++
+	}
+	if acknowledged == 0 {
+		return
+	}
+
+	clear(q.inFlight[:acknowledged])
+	q.inFlight = q.inFlight[acknowledged:]
+}
+
+func (q *frameQueue) requestResync() {
+	q.mutex.Lock()
+	q.clearFramesLocked()
+	q.inFlight = q.inFlight[:0]
+	q.waitingForKeyframe = true
+	q.mutex.Unlock()
+}
+
+func (q *frameQueue) close() {
+	q.mutex.Lock()
+	q.closed = true
+	q.clearFramesLocked()
+	q.inFlight = q.inFlight[:0]
+	q.mutex.Unlock()
+}
+
+func (q *frameQueue) pushLocked(frame *outboundFrame) {
+	q.frames = append(q.frames, frame)
+	q.queuedBytes += len(frame.payload)
+}
+
+func (q *frameQueue) clearFramesLocked() {
+	clear(q.frames)
+	q.frames = q.frames[:0]
+	q.queuedBytes = 0
+}
+
+func (q *frameQueue) signalLocked() {
+	select {
+	case q.wake <- struct{}{}:
+	default:
+	}
+}
+
+type client struct {
+	conn  *websocket.Conn
+	queue *frameQueue
+
+	done       chan struct{}
+	writerDone chan struct{}
+	closeOnce  sync.Once
+}
+
+func newClient(conn *websocket.Conn) *client {
+	return &client{
+		conn:       conn,
+		queue:      newFrameQueue(defaultQueueFrames, defaultQueueBytes),
+		done:       make(chan struct{}),
+		writerDone: make(chan struct{}),
+	}
+}
+
+func (c *client) start() {
+	go c.writeLoop()
+}
+
+func (c *client) close() {
+	c.closeOnce.Do(func() {
+		close(c.done)
+		c.queue.close()
+		_ = c.conn.Close()
+	})
+}
+
+func (c *client) wait() {
+	<-c.writerDone
+}
+
+func (c *client) captureState() (active bool, flowControlled bool, canAdvance bool) {
+	return c.queue.captureState()
+}
+
+func hasCaptureDemand(clients []*client) bool {
+	for _, client := range clients {
+		active, flowControlled, canAdvance := client.captureState()
+		if !active {
+			continue
+		}
+		if !flowControlled || canAdvance {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (c *client) offer(frame *outboundFrame) {
+	c.queue.offer(frame)
+}
+
+func (c *client) handleControl(messageType int, data []byte) {
+	if messageType != websocket.BinaryMessage || len(data) == 0 {
+		return
+	}
+
+	switch data[0] {
+	case frameAckMessage:
+		if len(data) == 9 {
+			timestamp := int64(binary.LittleEndian.Uint64(data[1:]))
+			c.queue.acknowledge(timestamp)
+		}
+	case streamResyncMessage:
+		if len(data) == 1 {
+			c.queue.requestResync()
+		}
+	}
+}
+
+func newOutboundFrame(isKeyFrame bool, timestamp int64, data []byte) *outboundFrame {
+	payload := make([]byte, 9+len(data))
+	if isKeyFrame {
+		payload[0] = 1
+	}
+	binary.LittleEndian.PutUint64(payload[1:9], uint64(timestamp))
+	copy(payload[9:], data)
+
+	return &outboundFrame{
+		key:       isKeyFrame,
+		timestamp: timestamp,
+		payload:   payload,
+	}
+}
+
+func (c *client) writeLoop() {
+	defer close(c.writerDone)
+
+	pingTicker := time.NewTicker(pingPeriod)
+	defer pingTicker.Stop()
+
+	for {
+		select {
+		case <-pingTicker.C:
+			deadline := time.Now().Add(writeWait)
+			if err := c.conn.WriteControl(websocket.PingMessage, nil, deadline); err != nil {
+				c.close()
+				return
+			}
+		default:
+		}
+
+		if frame := c.queue.popForWrite(); frame != nil {
+			if err := c.conn.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
+				c.close()
+				return
+			}
+			if err := c.conn.WriteMessage(websocket.BinaryMessage, frame.payload); err != nil {
+				log.Debugf("failed to write h264 frame to %s: %s", c.conn.RemoteAddr(), err)
+				c.close()
+				return
+			}
+			continue
+		}
+
+		select {
+		case <-c.done:
+			return
+		case <-c.queue.wake:
+		case <-pingTicker.C:
+			deadline := time.Now().Add(writeWait)
+			if err := c.conn.WriteControl(websocket.PingMessage, nil, deadline); err != nil {
+				c.close()
+				return
+			}
+		}
+	}
+}

server/service/stream/direct/h264.go

@@ -3,6 +3,7 @@ package direct
 import (
 	"NanoKVM-Server/service/stream"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -10,6 +11,8 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+const pongWait = 30 * time.Second
+
 var (
 	streamer = newStreamer()
 	upgrader = websocket.Upgrader{
@@ -26,24 +29,36 @@ func Connect(c *gin.Context) {
 		log.Errorf("failed to upgrade to websocket: %s", err)
 		return
 	}
+	client := newClient(ws)
+	if flowWindow, err := strconv.Atoi(c.Query("flow")); err == nil && flowWindow > 0 {
+		client.queue.enableFlowControl(flowWindow)
+	}
 	defer func() {
-		_ = ws.Close()
+		streamer.removeClient(client)
+		client.close()
+		client.wait()
 		log.Debugf("h264 websocket disconnected: %s", ws.RemoteAddr())
 	}()
 	log.Debugf("h264 websocket connected: %s", ws.RemoteAddr())
 
-	_ = ws.SetReadDeadline(time.Time{})
+	ws.SetReadLimit(64)
+	_ = ws.SetReadDeadline(time.Now().Add(pongWait))
+	ws.SetPongHandler(func(string) error {
+		return ws.SetReadDeadline(time.Now().Add(pongWait))
+	})
 
-	streamer.addClient(ws)
-	defer streamer.removeClient(ws)
+	streamer.addClient(client)
 
 	unregisterMode := stream.RegisterH264Mode(stream.H264ModeDirect)
 	defer unregisterMode()
 
 	for {
-		if _, _, err := ws.NextReader(); err != nil {
+		messageType, data, err := ws.ReadMessage()
+		if err != nil {
 			log.Debugf("failed to read message (client disconnected): %s", err)
 			return
 		}
+		_ = ws.SetReadDeadline(time.Now().Add(pongWait))
+		client.handleControl(messageType, data)
 	}
 }

server/service/stream/direct/pool.go

@@ -1,12 +0,0 @@
-package direct
-
-import (
-	"bytes"
-	"sync"
-)
-
-var BufferPool = sync.Pool{
-	New: func() interface{} {
-		return new(bytes.Buffer)
-	},
-}

server/service/stream/direct/streamer.go

@@ -4,62 +4,71 @@ import (
 	"NanoKVM-Server/common"
 	"NanoKVM-Server/service/stream"
 	"NanoKVM-Server/service/vm"
-	"bytes"
-	"encoding/binary"
 	"sync"
 	"sync/atomic"
 	"time"
 
-	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
 )
 
 type Streamer struct {
 	mutex          sync.Mutex
-	clients        map[*websocket.Conn]bool
-	clientSnapshot atomic.Pointer[[]*websocket.Conn]
-	running        int32
+	clients        map[*client]struct{}
+	clientSnapshot atomic.Pointer[[]*client]
+	running        bool
 	viewerVersion  uint64
 }
 
 func newStreamer() *Streamer {
 	s := &Streamer{
-		clients: make(map[*websocket.Conn]bool),
+		clients: make(map[*client]struct{}),
 	}
 	s.updateClientSnapshotLocked()
 
 	return s
 }
 
-func (s *Streamer) addClient(ws *websocket.Conn) {
+func (s *Streamer) addClient(client *client) {
+	client.start()
+
 	s.mutex.Lock()
-	s.clients[ws] = true
+	s.clients[client] = struct{}{}
 	count := s.updateClientSnapshotLocked()
 	s.viewerVersion++
 	version := s.viewerVersion
+	start := !s.running
+	if start {
+		s.running = true
+	}
 	s.mutex.Unlock()
 	vm.UpdateHdmiViewerSnapshot("direct", count, version)
 
-	if atomic.CompareAndSwapInt32(&s.running, 0, 1) {
+	if start {
 		go s.run()
 		log.Debug("h264 stream started")
 	}
 }
 
-func (s *Streamer) removeClient(ws *websocket.Conn) {
+func (s *Streamer) removeClient(client *client) {
 	s.mutex.Lock()
-	delete(s.clients, ws)
+	if _, exists := s.clients[client]; !exists {
+		s.mutex.Unlock()
+		return
+	}
+
+	delete(s.clients, client)
 	count := s.updateClientSnapshotLocked()
 	s.viewerVersion++
 	version := s.viewerVersion
 	s.mutex.Unlock()
+	client.close()
 	vm.UpdateHdmiViewerSnapshot("direct", count, version)
 
 	log.Debugf("h264 websocket disconnected, remaining clients: %d", count)
 }
 
 func (s *Streamer) updateClientSnapshotLocked() int {
-	clients := make([]*websocket.Conn, 0, len(s.clients))
+	clients := make([]*client, 0, len(s.clients))
 	for client := range s.clients {
 		clients = append(clients, client)
 	}
@@ -68,7 +77,7 @@ func (s *Streamer) updateClientSnapshotLocked() int {
 	return len(clients)
 }
 
-func (s *Streamer) getClients() []*websocket.Conn {
+func (s *Streamer) getClients() []*client {
 	clients := s.clientSnapshot.Load()
 	if clients == nil {
 		return nil
@@ -78,8 +87,6 @@ func (s *Streamer) getClients() []*websocket.Conn {
 }
 
 func (s *Streamer) run() {
-	defer atomic.StoreInt32(&s.running, 0)
-
 	screen := common.GetScreen()
 	common.CheckScreen()
 	fps := screen.FPS
@@ -93,24 +100,10 @@ func (s *Streamer) run() {
 	for range ticker.C {
 		clients := s.getClients()
 		if len(clients) == 0 {
-			log.Debug("h264 stream stopped due to no clients")
-			return
-		}
-
-		data, result := vision.ReadH264(screen.Width, screen.Height, screen.BitRate)
-		stream.UpdateCaptureStatus(stream.CaptureModeDirect, result)
-		if result < 0 || len(data) == 0 {
-			continue
-		}
-
-		isKeyFrame := byte(0)
-		if result == 3 {
-			isKeyFrame = byte(1)
-		}
-
-		timestamp := time.Since(startTime).Microseconds()
-
-		if err := s.send(clients, isKeyFrame, timestamp, data); err != nil {
+			if s.stopIfIdle() {
+				log.Debug("h264 stream stopped due to no clients")
+				return
+			}
 			continue
 		}
 
@@ -119,42 +112,34 @@ func (s *Streamer) run() {
 			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 {
+			continue
+		}
+
+		timestamp := time.Since(startTime).Microseconds()
+		frame := newOutboundFrame(result == 3, timestamp, data)
+		for _, client := range clients {
+			client.offer(frame)
+		}
+
 		stream.GetFrameRateCounter().Update()
 	}
 }
 
-func (s *Streamer) send(clients []*websocket.Conn, isKeyFrame byte, timestamp int64, data []byte) error {
-	buf := BufferPool.Get().(*bytes.Buffer)
-	defer BufferPool.Put(buf)
+func (s *Streamer) stopIfIdle() bool {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
 
-	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
+	if len(s.clients) > 0 {
+		return false
 	}
 
-	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
-	}
-
-	if _, err := buf.Write(data); err != nil {
-		log.Errorf("failed to write h264 data: %s", err)
-		return err
-	}
-
-	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)
-
-			s.removeClient(client)
-			_ = client.Close()
-		}
-	}
-
-	return nil
+	s.running = false
+	return true
 }