aiops/backend/controller/rdp_live.go

package controller

import (
	"sync"

	rdpengine "ezops/pkg/rdp"
)

// rdpLiveBroadcaster 将 RDP 主会话帧广播给审计实时查看端，慢消费者会丢旧帧保最新。 // Logan.Li
type rdpLiveBroadcaster struct {
	mu      sync.Mutex
	closed  bool
	clients map[chan rdpengine.ServerFrame]struct{}
}

func newRDPLiveBroadcaster() *rdpLiveBroadcaster {
	return &rdpLiveBroadcaster{clients: make(map[chan rdpengine.ServerFrame]struct{})}
}

func (b *rdpLiveBroadcaster) subscribe() chan rdpengine.ServerFrame {
	ch := make(chan rdpengine.ServerFrame, 8)
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.closed {
		close(ch)
		return ch
	}
	b.clients[ch] = struct{}{}
	return ch
}

func (b *rdpLiveBroadcaster) unsubscribe(ch chan rdpengine.ServerFrame) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if _, ok := b.clients[ch]; ok {
		delete(b.clients, ch)
		close(ch)
	}
}

func (b *rdpLiveBroadcaster) publish(frame rdpengine.ServerFrame) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.closed {
		return
	}
	for ch := range b.clients {
		select {
		case ch <- frame:
		default:
			select {
			case <-ch:
			default:
			}
			select {
			case ch <- frame:
			default:
			}
		}
	}
}

func (b *rdpLiveBroadcaster) close() {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.closed {
		return
	}
	b.closed = true
	for ch := range b.clients {
		delete(b.clients, ch)
		close(ch)
	}
}