WebRTC cursor handler extracted + cache.

internal/webrtc/cursor/image.go

@@ -0,0 +1,126 @@
+package cursor
+
+import (
+	"reflect"
+	"sync"
+
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+
+	"demodesk/neko/internal/types"
+	"demodesk/neko/internal/utils"
+)
+
+func NewImage(desktop types.DesktopManager) *ImageCtx {
+	return &ImageCtx{
+		logger:    log.With().Str("module", "cursor-image").Logger(),
+		desktop:   desktop,
+		listeners: map[uintptr]*func(entry *ImageEntry){},
+		cache:     map[uint64]*ImageEntry{},
+	}
+}
+
+type ImageCtx struct {
+	logger    zerolog.Logger
+	desktop   types.DesktopManager
+	emitMu    sync.Mutex
+	listeners map[uintptr]*func(entry *ImageEntry)
+	cacheMu   sync.Mutex
+	cache     map[uint64]*ImageEntry
+	current   *ImageEntry
+}
+
+type ImageEntry struct {
+	Cursor *types.CursorImage
+	Image  []byte
+}
+
+func (manager *ImageCtx) Start() {
+	manager.desktop.OnCursorChanged(func(serial uint64) {
+		entry, err := manager.GetCached(serial)
+		if err != nil {
+			manager.logger.Warn().Err(err).Msg("failed to get cursor image")
+			return
+		}
+
+		for _, emit := range manager.listeners {
+			(*emit)(entry)
+		}
+	})
+}
+
+func (manager *ImageCtx) Shutdown() {
+	manager.logger.Info().Msgf("shutting down")
+
+	manager.emitMu.Lock()
+	for key := range manager.listeners {
+		delete(manager.listeners, key)
+	}
+	manager.emitMu.Unlock()
+}
+
+func (manager *ImageCtx) GetCached(serial uint64) (*ImageEntry, error) {
+	manager.cacheMu.Lock()
+	entry, ok := manager.cache[serial]
+	manager.cacheMu.Unlock()
+
+	if ok {
+		return entry, nil
+	}
+
+	entry, err := manager.fetchEntry()
+	if err != nil {
+		return nil, err
+	}
+
+	manager.cacheMu.Lock()
+	manager.cache[serial] = entry
+	manager.cacheMu.Unlock()
+
+	return entry, nil
+}
+
+func (manager *ImageCtx) GetCurrent() (*ImageEntry, error) {
+	if manager.current != nil {
+		return manager.current, nil
+	}
+
+	return manager.fetchEntry()
+}
+
+func (manager *ImageCtx) AddListener(listener *func(entry *ImageEntry)) {
+	manager.emitMu.Lock()
+	defer manager.emitMu.Unlock()
+
+	if listener != nil {
+		ptr := reflect.ValueOf(listener).Pointer()
+		manager.listeners[ptr] = listener
+	}
+}
+
+func (manager *ImageCtx) RemoveListener(listener *func(entry *ImageEntry)) {
+	manager.emitMu.Lock()
+	defer manager.emitMu.Unlock()
+
+	if listener != nil {
+		ptr := reflect.ValueOf(listener).Pointer()
+		delete(manager.listeners, ptr)
+	}
+}
+
+func (manager *ImageCtx) fetchEntry() (*ImageEntry, error) {
+	cur := manager.desktop.GetCursorImage()
+
+	img, err := utils.CreatePNGImage(cur.Image)
+	if err != nil {
+		return nil, err
+	}
+
+	entry := &ImageEntry{
+		Cursor: cur,
+		Image:  img,
+	}
+
+	manager.current = entry
+	return entry, nil
+}

internal/webrtc/cursor/position.go

@@ -0,0 +1,68 @@
+package cursor
+
+import (
+	"reflect"
+	"sync"
+
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+
+	"demodesk/neko/internal/types"
+)
+
+func NewPosition(desktop types.DesktopManager) *PositionCtx {
+	return &PositionCtx{
+		logger:    log.With().Str("module", "cursor-position").Logger(),
+		desktop:   desktop,
+		listeners: map[uintptr]*func(x, y int){},
+	}
+}
+
+type PositionCtx struct {
+	logger    zerolog.Logger
+	desktop   types.DesktopManager
+	emitMu    sync.Mutex
+	listeners map[uintptr]*func(x, y int)
+}
+
+func (manager *PositionCtx) Start() {
+	manager.desktop.OnCursorPosition(func(x, y int) {
+		for _, emit := range manager.listeners {
+			(*emit)(x, y)
+		}
+	})
+}
+
+func (manager *PositionCtx) Shutdown() {
+	manager.logger.Info().Msgf("shutting down")
+
+	manager.emitMu.Lock()
+	for key := range manager.listeners {
+		delete(manager.listeners, key)
+	}
+	manager.emitMu.Unlock()
+}
+
+func (manager *PositionCtx) GetCurrent() (x, y int) {
+	return manager.desktop.GetCursorPosition()
+}
+
+func (manager *PositionCtx) AddListener(listener *func(x, y int)) {
+	manager.emitMu.Lock()
+	defer manager.emitMu.Unlock()
+
+	if listener != nil {
+		ptr := reflect.ValueOf(listener).Pointer()
+		manager.listeners[ptr] = listener
+	}
+}
+
+func (manager *PositionCtx) RemoveListener(listener *func(x, y int)) {
+	manager.emitMu.Lock()
+	defer manager.emitMu.Unlock()
+
+	if listener != nil {
+		ptr := reflect.ValueOf(listener).Pointer()
+		delete(manager.listeners, ptr)
+	}
+}

internal/webrtc/manager.go

@@ -3,7 +3,6 @@ package webrtc
 import (
 	"fmt"
 	"io"
-	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -17,7 +16,7 @@ import (
 	"demodesk/neko/internal/types"
 	"demodesk/neko/internal/types/event"
 	"demodesk/neko/internal/types/message"
-	"demodesk/neko/internal/utils"
+	"demodesk/neko/internal/webrtc/cursor"
 )
 
 // how long is can take between sending offer and connecting
@@ -39,9 +38,8 @@ func New(desktop types.DesktopManager, capture types.CaptureManager, config *con
 		capture:      capture,
 		config:       config,
 		participants: 0,
-		// TODO: Refactor.
-		curImgListeners: map[uintptr]*func(cur *types.CursorImage, img []byte){},
-		curPosListeners: map[uintptr]*func(x, y int){},
+		curImage:     cursor.NewImage(desktop),
+		curPosition:  cursor.NewPosition(desktop),
 	}
 }
 
@@ -54,9 +52,8 @@ type WebRTCManagerCtx struct {
 	capture      types.CaptureManager
 	config       *config.WebRTC
 	participants uint32
-	// TODO: Refactor.
-	curImgListeners map[uintptr]*func(cur *types.CursorImage, img []byte)
-	curPosListeners map[uintptr]*func(x, y int)
+	curImage     *cursor.ImageCtx
+	curPosition  *cursor.PositionCtx
 }
 
 func (manager *WebRTCManagerCtx) Start() {
@@ -88,32 +85,16 @@ func (manager *WebRTCManagerCtx) Start() {
 		Str("nat_ips", strings.Join(manager.config.NAT1To1IPs, ",")).
 		Msgf("webrtc starting")
 
-	// TODO: Refactor.
-	manager.desktop.OnCursorChanged(func(serial uint64) {
-		cur := manager.desktop.GetCursorImage()
-
-		img, err := utils.CreatePNGImage(cur.Image)
-		if err != nil {
-			manager.logger.Warn().Err(err).Msg("failed to create cursor image")
-			return
-		}
-
-		for _, emit := range manager.curImgListeners {
-			(*emit)(cur, img)
-		}
-	})
-
-	// TODO: Refactor.
-	manager.desktop.OnCursorPosition(func(x, y int) {
-		for _, emit := range manager.curPosListeners {
-			(*emit)(x, y)
-		}
-	})
+	manager.curImage.Start()
+	manager.curPosition.Start()
 }
 
 func (manager *WebRTCManagerCtx) Shutdown() error {
 	manager.logger.Info().Msgf("webrtc shutting down")
 
+	manager.curImage.Shutdown()
+	manager.curPosition.Shutdown()
+
 	manager.audioStop()
 	return nil
 }
@@ -277,8 +258,8 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
 		dataChannel: dataChannel,
 	}
 
-	cursorChange := func(cur *types.CursorImage, img []byte) {
-		if err := peer.SendCursorImage(cur, img); err != nil {
+	cursorImage := func(entry *cursor.ImageEntry) {
+		if err := peer.SendCursorImage(entry.Cursor, entry.Image); err != nil {
 			manager.logger.Warn().Err(err).Msg("could not send cursor image")
 		}
 	}
@@ -293,10 +274,6 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
 		}
 	}
 
-	// TODO: Refactor.
-	cursorChangePtr := reflect.ValueOf(&cursorChange).Pointer()
-	cursorPositionPtr := reflect.ValueOf(&cursorPosition).Pointer()
-
 	connection.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
 		switch state {
 		case webrtc.PeerConnectionStateConnected:
@@ -328,33 +305,32 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
 				}
 			}
 
-			// TODO: Refactor.
-			delete(manager.curImgListeners, cursorChangePtr)
-			delete(manager.curPosListeners, cursorPositionPtr)
-
 			manager.mu.Unlock()
 		}
 	})
 
 	dataChannel.OnOpen(func() {
-		// TODO: Refactor.
-		manager.curImgListeners[cursorChangePtr] = &cursorChange
-		manager.curPosListeners[cursorPositionPtr] = &cursorPosition
+		manager.curImage.AddListener(&cursorImage)
+		manager.curPosition.AddListener(&cursorPosition)
 
 		// send initial cursor image
-		cur := manager.desktop.GetCursorImage()
-		img, err := utils.CreatePNGImage(cur.Image)
+		entry, err := manager.curImage.GetCurrent()
 		if err == nil {
-			cursorChange(cur, img)
+			cursorImage(entry)
 		} else {
-			manager.logger.Warn().Err(err).Msg("failed to create cursor image")
+			manager.logger.Warn().Err(err).Msg("failed to get cursor image")
 		}
 
 		// send initial cursor position
-		x, y := manager.desktop.GetCursorPosition()
+		x, y := manager.curPosition.GetCurrent()
 		cursorPosition(x, y)
 	})
 
+	dataChannel.OnClose(func() {
+		manager.curImage.RemoveListener(&cursorImage)
+		manager.curPosition.RemoveListener(&cursorPosition)
+	})
+
 	dataChannel.OnMessage(func(message webrtc.DataChannelMessage) {
 		if !session.IsHost() {
 			return