fix screenshot mem usage + fix backstage not starting after being stopped once

This commit is contained in:
KingKDot
2026-08-01 00:46:03 -05:00
parent b68a306708
commit 7b249457a1
18 changed files with 675 additions and 99 deletions
@@ -23,6 +23,7 @@ var (
procDWMCreateWindowExW = user32.NewProc("CreateWindowExW") procDWMCreateWindowExW = user32.NewProc("CreateWindowExW")
procDWMShowWindow = user32.NewProc("ShowWindow") procDWMShowWindow = user32.NewProc("ShowWindow")
procDWMUpdateWindow = user32.NewProc("UpdateWindow") procDWMUpdateWindow = user32.NewProc("UpdateWindow")
procDWMDestroyWindow = user32.NewProc("DestroyWindow")
backstageDWMStateMu sync.Mutex backstageDWMStateMu sync.Mutex
backstageDWMHost uintptr backstageDWMHost uintptr
backstageDWMThumbnails map[uintptr]*backstageDWMThumbnail backstageDWMThumbnails map[uintptr]*backstageDWMThumbnail
@@ -460,11 +461,23 @@ func backstageCleanupDWMThumbnails() {
defer backstageDWMStateMu.Unlock() defer backstageDWMStateMu.Unlock()
backstageUnregisterDWMThumbnailsLocked() backstageUnregisterDWMThumbnailsLocked()
if backstageDWMHost != 0 { if backstageDWMHost != 0 {
procDWMShowWindow.Call(backstageDWMHost, 0) // This runs on the worker that created the host. Destroying the HWND is
// required before that OS thread can be moved back to its original
// desktop; merely hiding it leaves SetThreadDesktop failing with
// ERROR_BUSY on the next lifecycle.
procDWMDestroyWindow.Call(backstageDWMHost)
} }
backstageDWMHost = 0 backstageDWMHost = 0
} }
func backstageAbandonDWMThumbnails() {
backstageDWMStateMu.Lock()
backstageDWMThumbnails = nil
backstageDWMOrder = nil
backstageDWMHost = 0
backstageDWMStateMu.Unlock()
}
func backstageUnregisterDWMThumbnailsLocked() { func backstageUnregisterDWMThumbnailsLocked() {
for source, entry := range backstageDWMThumbnails { for source, entry := range backstageDWMThumbnails {
procDwmUnregisterThumbnail.Call(entry.handle) procDwmUnregisterThumbnail.Call(entry.handle)
@@ -0,0 +1,25 @@
//go:build windows
package capture
import "testing"
func TestBackstageDesktopCanRestartAfterCleanup(t *testing.T) {
t.Cleanup(CleanupbackstageDesktop)
for cycle := 1; cycle <= 2; cycle++ {
if err := InitializebackstageDesktop(); err != nil {
t.Fatalf("cycle %d initialize: %v", cycle, err)
}
if err := ensurebackstageThread(); err != nil {
t.Fatalf("cycle %d worker: %v", cycle, err)
}
img, err := BackstageCaptureDisplay(0)
if err != nil {
t.Fatalf("cycle %d capture: %v", cycle, err)
}
if img == nil {
t.Fatalf("cycle %d capture returned no image", cycle)
}
CleanupbackstageDesktop()
}
}
@@ -2,7 +2,10 @@
package capture package capture
import "testing" import (
"testing"
"time"
)
func TestBackstageWindowCacheByteSizeAndLimit(t *testing.T) { func TestBackstageWindowCacheByteSizeAndLimit(t *testing.T) {
bytes, ok := backstageWindowCacheByteSize(3840, 2160) bytes, ok := backstageWindowCacheByteSize(3840, 2160)
@@ -52,3 +55,48 @@ func TestBackstagePrintWindowFallbackToggle(t *testing.T) {
t.Fatal("PrintWindow fallback did not re-enable") t.Fatal("PrintWindow fallback did not re-enable")
} }
} }
func TestBackstagePrintWindowTimeoutQuarantinesHungWindow(t *testing.T) {
savedFn := backstagePrintWindowFn
savedCache := backstageWinCache
savedBytes := backstageWinCacheBytes
savedHung := backstageHungWindows
release := make(chan struct{})
returned := make(chan struct{})
backstagePrintWindowFn = func(_, _ uintptr, _ uint32) bool {
<-release
close(returned)
return true
}
entry := &backstageWinCacheEntry{bytes: 64}
const hwnd = uintptr(0x1234)
backstageWinCache = map[uintptr]*backstageWinCacheEntry{hwnd: entry}
backstageWinCacheBytes = entry.bytes
backstageHungWindows = make(map[uintptr]struct{})
t.Cleanup(func() {
close(release)
select {
case <-returned:
case <-time.After(time.Second):
t.Error("timed-out PrintWindow worker did not exit after release")
}
backstagePrintWindowFn = savedFn
backstageWinCache = savedCache
backstageWinCacheBytes = savedBytes
backstageHungWindows = savedHung
})
start := time.Now()
if backstagePrintWindowWithTimeout(hwnd, entry) {
t.Fatal("hung PrintWindow unexpectedly succeeded")
}
if elapsed := time.Since(start); elapsed > time.Second {
t.Fatalf("PrintWindow timeout took too long: %s", elapsed)
}
if backstageWinCache[hwnd] != nil || backstageWinCacheBytes != 0 {
t.Fatal("timed-out PrintWindow cache entry was not detached")
}
if _, ok := backstageHungWindows[hwnd]; !ok {
t.Fatal("timed-out window was not quarantined")
}
}
@@ -155,6 +155,7 @@ const (
var ( var (
backstageDesktopHandle uintptr backstageDesktopHandle uintptr
backstageDesktopMu sync.Mutex backstageDesktopMu sync.Mutex
backstageCaptureMu sync.Mutex
backstageDesktopName = "OverlordHiddenDesktop" backstageDesktopName = "OverlordHiddenDesktop"
backstageInitialized bool backstageInitialized bool
backstageOriginalDesktop uintptr backstageOriginalDesktop uintptr
@@ -163,6 +164,7 @@ var (
backstageThreadErr error backstageThreadErr error
backstageThreadReady chan struct{} backstageThreadReady chan struct{}
backstageThreadTasks chan backstageTask backstageThreadTasks chan backstageTask
backstageThreadDone chan struct{}
backstageWatchdogOnce sync.Once backstageWatchdogOnce sync.Once
backstageNoWindowLogNs atomic.Int64 backstageNoWindowLogNs atomic.Int64
backstageInputMu sync.Mutex backstageInputMu sync.Mutex
@@ -191,6 +193,7 @@ var (
backstageWinCachePrev []byte backstageWinCachePrev []byte
backstageWinCacheBytes int64 backstageWinCacheBytes int64
backstageWinCacheSeq uint64 backstageWinCacheSeq uint64
backstageHungWindows map[uintptr]struct{}
backstageCompHdcMem uintptr backstageCompHdcMem uintptr
backstageCompHbmp uintptr backstageCompHbmp uintptr
@@ -215,6 +218,7 @@ const (
backstageTaskKeyUp backstageTaskKeyUp
backstageTaskMouseWheel backstageTaskMouseWheel
backstageTaskAutoStartExplorer backstageTaskAutoStartExplorer
backstageTaskShutdown
) )
type backstageTask struct { type backstageTask struct {
@@ -293,6 +297,7 @@ type backstageWinCacheEntry struct {
const ( const (
backstageMaxWindowCacheEntries = 64 backstageMaxWindowCacheEntries = 64
backstageMaxWindowCacheBytes = int64(256 << 20) backstageMaxWindowCacheBytes = int64(256 << 20)
backstagePrintWindowTimeout = 250 * time.Millisecond
) )
type keybdInput struct { type keybdInput struct {
@@ -388,8 +393,11 @@ func CleanupbackstageDesktop() {
resetH264D3D11TextureEncoder("backstage") resetH264D3D11TextureEncoder("backstage")
backstageCleanupFrameReaders() backstageCleanupFrameReaders()
captureMu.Lock() dwmCleaned, workerStopped := shutdownbackstageThreadLocked(time.Second)
backstageCleanupDWMThumbnails() if !dwmCleaned {
backstageAbandonDWMThumbnails()
}
backstageCaptureMu.Lock()
backstageFreeCapCache() backstageFreeCapCache()
@@ -407,7 +415,7 @@ func CleanupbackstageDesktop() {
backstageCompBits = nil backstageCompBits = nil
backstageCompW = 0 backstageCompW = 0
backstageCompH = 0 backstageCompH = 0
captureMu.Unlock() backstageCaptureMu.Unlock()
backstageInputMu.Lock() backstageInputMu.Lock()
backstageShiftDown = false backstageShiftDown = false
@@ -420,14 +428,10 @@ func CleanupbackstageDesktop() {
backstageInputMu.Unlock() backstageInputMu.Unlock()
backstageLastScale.Store(0) backstageLastScale.Store(0)
if backstageDesktopHandle != 0 { if backstageDesktopHandle != 0 && workerStopped {
if backstageOriginalDesktop != 0 {
procSetThreadDesktop.Call(backstageOriginalDesktop)
}
procCloseDesktop.Call(backstageDesktopHandle) procCloseDesktop.Call(backstageDesktopHandle)
backstageDesktopHandle = 0
} }
backstageDesktopHandle = 0
backstageInitialized = false backstageInitialized = false
backstageExplorerStarted = false backstageExplorerStarted = false
@@ -436,16 +440,52 @@ func CleanupbackstageDesktop() {
resetWinUI3Cache() resetWinUI3Cache()
resetInputSiteCache() resetInputSiteCache()
if backstageThreadTasks != nil { backstageThreadTasks = nil
close(backstageThreadTasks)
backstageThreadTasks = nil
}
backstageThreadReady = nil backstageThreadReady = nil
backstageThreadDone = nil
backstageThreadErr = nil backstageThreadErr = nil
backstageThreadOnce = sync.Once{} backstageThreadOnce = sync.Once{}
backstageWatchdogOnce = sync.Once{} backstageWatchdogOnce = sync.Once{}
} }
func shutdownbackstageThreadLocked(timeout time.Duration) (dwmCleaned, stopped bool) {
tasks := backstageThreadTasks
done := backstageThreadDone
if tasks == nil || done == nil {
return false, true
}
if timeout <= 0 {
timeout = time.Second
}
resp := make(chan backstageTaskResult, 1)
task := backstageTask{kind: backstageTaskShutdown, resp: resp, id: backstageTaskSeq.Add(1), queuedAt: time.Now()}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case tasks <- task:
case <-done:
return false, true
case <-timer.C:
log.Printf("backstage cleanup: worker shutdown enqueue timed out")
return false, false
}
select {
case result := <-resp:
dwmCleaned = result.err == nil
case <-done:
case <-timer.C:
log.Printf("backstage cleanup: worker shutdown timed out")
return false, false
}
select {
case <-done:
return dwmCleaned, true
case <-timer.C:
log.Printf("backstage cleanup: worker exit timed out")
return dwmCleaned, false
}
}
func SetbackstageCursorCapture(enabled bool) { func SetbackstageCursorCapture(enabled bool) {
backstageCursorEnabled = enabled backstageCursorEnabled = enabled
} }
@@ -476,31 +516,38 @@ func ensurebackstageThread() error {
} }
backstageThreadOnce.Do(func() { backstageThreadOnce.Do(func() {
backstageThreadReady = make(chan struct{}) ready := make(chan struct{})
backstageThreadTasks = make(chan backstageTask) tasks := make(chan backstageTask)
done := make(chan struct{})
backstageThreadReady = ready
backstageThreadTasks = tasks
backstageThreadDone = done
backstageWatchdogOnce.Do(func() { backstageWatchdogOnce.Do(func() {
go func() { go func() {
defer recoverAndLog("backstage watchdog", nil) defer recoverAndLog("backstage watchdog", nil)
backstageThreadWatchdog() backstageThreadWatchdog()
}() }()
}) })
go func(handle uintptr) { go func(handle, originalDesktop uintptr) {
defer recoverAndLog("backstage desktop thread", nil) defer recoverAndLog("backstage desktop thread", nil)
runtime.LockOSThread() runtime.LockOSThread()
defer runtime.UnlockOSThread() defer func() {
if originalDesktop != 0 {
procSetThreadDesktop.Call(originalDesktop)
}
close(done)
runtime.UnlockOSThread()
}()
r, _, err := procSetThreadDesktop.Call(handle) r, _, err := procSetThreadDesktop.Call(handle)
if r == 0 { if r == 0 {
backstageThreadErr = fmt.Errorf("failed to set thread desktop: %v", err) backstageThreadErr = fmt.Errorf("failed to set thread desktop: %v", err)
close(backstageThreadReady) close(ready)
for task := range backstageThreadTasks {
task.resp <- backstageTaskResult{err: backstageThreadErr}
}
return return
} }
close(backstageThreadReady) close(ready)
for task := range backstageThreadTasks { for task := range tasks {
start := time.Now() start := time.Now()
backstageCurrentTaskID.Store(task.id) backstageCurrentTaskID.Store(task.id)
backstageCurrentTaskKind.Store(int64(task.kind)) backstageCurrentTaskKind.Store(int64(task.kind))
@@ -530,6 +577,8 @@ func ensurebackstageThread() error {
result.err = backstageMouseWheelOnThread(task.delta) result.err = backstageMouseWheelOnThread(task.delta)
case backstageTaskAutoStartExplorer: case backstageTaskAutoStartExplorer:
result.err = backstageAutoStartExplorerOnThread() result.err = backstageAutoStartExplorerOnThread()
case backstageTaskShutdown:
backstageCleanupDWMThumbnails()
default: default:
result.img, result.err = BackstageCaptureDisplayOnThread(task.display) result.img, result.err = BackstageCaptureDisplayOnThread(task.display)
} }
@@ -547,8 +596,11 @@ func ensurebackstageThread() error {
backstageCurrentTaskKind.Store(-1) backstageCurrentTaskKind.Store(-1)
backstageCurrentTaskID.Store(0) backstageCurrentTaskID.Store(0)
task.resp <- result task.resp <- result
if task.kind == backstageTaskShutdown {
return
}
} }
}(desktopHandle) }(desktopHandle, backstageOriginalDesktop)
}) })
if backstageThreadReady != nil { if backstageThreadReady != nil {
@@ -774,6 +826,8 @@ func backstageTaskKindName(kind backstageTaskKind) string {
return "mouse_wheel" return "mouse_wheel"
case backstageTaskAutoStartExplorer: case backstageTaskAutoStartExplorer:
return "auto_start_explorer" return "auto_start_explorer"
case backstageTaskShutdown:
return "shutdown"
default: default:
return fmt.Sprintf("unknown(%d)", kind) return fmt.Sprintf("unknown(%d)", kind)
} }
@@ -874,8 +928,8 @@ func backstageEnsureCapCache(w, h int) (uintptr, uintptr, []byte, bool) {
func BackstageCaptureDisplayOnThread(display int) (*image.RGBA, error) { func BackstageCaptureDisplayOnThread(display int) (*image.RGBA, error) {
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2 //garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
captureMu.Lock() backstageCaptureMu.Lock()
defer captureMu.Unlock() defer backstageCaptureMu.Unlock()
setDPIAware() setDPIAware()
@@ -1720,6 +1774,11 @@ func drawbackstageWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, tar
backstageRemoveWindowCacheEntry(h) backstageRemoveWindowCacheEntry(h)
} }
} }
for h := range backstageHungWindows {
if !alive[h] {
delete(backstageHungWindows, h)
}
}
return drawn return drawn
} }
@@ -1811,6 +1870,18 @@ func backstageRemoveWindowCacheEntry(hwnd uintptr) {
} }
} }
func backstageDetachWindowCacheEntry(hwnd uintptr, entry *backstageWinCacheEntry) bool {
if backstageWinCache[hwnd] != entry {
return false
}
delete(backstageWinCache, hwnd)
backstageWinCacheBytes -= entry.bytes
if backstageWinCacheBytes < 0 {
backstageWinCacheBytes = 0
}
return true
}
func backstageClearWindowCache() { func backstageClearWindowCache() {
for hwnd := range backstageWinCache { for hwnd := range backstageWinCache {
backstageRemoveWindowCacheEntry(hwnd) backstageRemoveWindowCacheEntry(hwnd)
@@ -1834,6 +1905,8 @@ var backstageDXGIEnabled atomic.Bool
var backstageUIAEnabled atomic.Bool var backstageUIAEnabled atomic.Bool
var backstagePrintWindowFallbackEnabled atomic.Bool var backstagePrintWindowFallbackEnabled atomic.Bool
var backstagePrintWindowFallbackLogNs atomic.Int64 var backstagePrintWindowFallbackLogNs atomic.Int64
var backstagePrintWindowTimeoutLogNs atomic.Int64
var backstagePrintWindowFn = printWindow
func init() { func init() {
backstageDXGIEnabled.Store(false) // disabled by default backstageDXGIEnabled.Store(false) // disabled by default
@@ -1865,6 +1938,49 @@ func GetbackstagePrintWindowFallbackEnabled() bool {
return backstagePrintWindowFallbackEnabled.Load() return backstagePrintWindowFallbackEnabled.Load()
} }
func backstagePrintWindowWithTimeout(hwnd uintptr, entry *backstageWinCacheEntry) bool {
if backstageHungWindows == nil {
backstageHungWindows = make(map[uintptr]struct{})
}
if _, quarantined := backstageHungWindows[hwnd]; quarantined {
return false
}
var ownership atomic.Int32 // 0=racing, 1=caller keeps cache, 2=worker frees detached cache
done := make(chan bool, 1)
go func() {
ok := false
func() {
defer func() { _ = recover() }()
ok = backstagePrintWindowFn(hwnd, entry.hdcMem, PW_RENDERFULLCONTENT)
}()
if ownership.CompareAndSwap(0, 1) {
done <- ok
return
}
backstageFreeCacheEntry(entry)
}()
timer := time.NewTimer(backstagePrintWindowTimeout)
defer timer.Stop()
select {
case ok := <-done:
return ok
case <-timer.C:
if !ownership.CompareAndSwap(0, 2) {
return <-done
}
backstageDetachWindowCacheEntry(hwnd, entry)
backstageHungWindows[hwnd] = struct{}{}
now := time.Now().UnixNano()
last := backstagePrintWindowTimeoutLogNs.Load()
if now-last > int64(5*time.Second) && backstagePrintWindowTimeoutLogNs.CompareAndSwap(last, now) {
log.Printf("backstage capture: PrintWindow timed out for hwnd=0x%x; quarantining window", hwnd)
}
return false
}
}
func drawbackstageWindowFromDXGI(hwnd uintptr, winLeft, winTop, winW, winH int, bounds image.Rectangle, target []byte, targetStride int) bool { func drawbackstageWindowFromDXGI(hwnd uintptr, winLeft, winTop, winW, winH int, bounds image.Rectangle, target []byte, targetStride int) bool {
if !backstageDXGIEnabled.Load() { if !backstageDXGIEnabled.Load() {
return false return false
@@ -1986,7 +2102,10 @@ func drawbackstageWindow(hdcScreen, hwnd uintptr, bounds image.Rectangle, target
return false return false
} }
if !printWindow(hwnd, entry.hdcMem, PW_RENDERFULLCONTENT) { if !backstagePrintWindowWithTimeout(hwnd, entry) {
if backstageWinCache[hwnd] != entry {
return false
}
entry.lastOK = false entry.lastOK = false
entry.age++ entry.age++
return false return false
+35 -23
View File
@@ -530,14 +530,13 @@ const desktopHardwareH264FallbackCooldown = 30 * time.Second
var consecutiveCaptureFails atomic.Int64 var consecutiveCaptureFails atomic.Int64
var ( var (
fpsWindowStart atomic.Int64 desktopFPSCounter frameRateCounter
fpsCount atomic.Int64 backstageFPSCounter frameRateCounter
fpsLatest atomic.Int64 lastFrameLog atomic.Int64
lastFrameLog atomic.Int64 metricsOnce sync.Once
metricsOnce sync.Once metricsEnabled bool
metricsEnabled bool lastKeyframe atomic.Int64
lastKeyframe atomic.Int64 fullNextFrames atomic.Int64
fullNextFrames atomic.Int64
desktopRecoveryAfterBackstageStart atomic.Bool desktopRecoveryAfterBackstageStart atomic.Bool
statFrames atomic.Int64 statFrames atomic.Int64
@@ -576,6 +575,12 @@ var (
backstageLastKeyframe atomic.Int64 backstageLastKeyframe atomic.Int64
) )
type frameRateCounter struct {
windowStart atomic.Int64
count atomic.Int64
latest atomic.Int64
}
type prevImage struct { type prevImage struct {
w int w int
h int h int
@@ -604,28 +609,36 @@ func logCodecSupport() {
}) })
} }
func frameFPS(now time.Time) int { func (c *frameRateCounter) fps(now time.Time) int {
start := fpsWindowStart.Load() start := c.windowStart.Load()
if start == 0 { if start == 0 {
if fpsWindowStart.CompareAndSwap(0, now.UnixNano()) { if c.windowStart.CompareAndSwap(0, now.UnixNano()) {
fpsCount.Store(1) c.count.Store(1)
return int(fpsLatest.Load()) return int(c.latest.Load())
} }
start = fpsWindowStart.Load() start = c.windowStart.Load()
} }
fpsCount.Add(1) c.count.Add(1)
elapsed := time.Duration(now.UnixNano() - start) elapsed := time.Duration(now.UnixNano() - start)
if elapsed >= time.Second { if elapsed >= time.Second {
frames := fpsCount.Swap(0) frames := c.count.Swap(0)
if frames > 0 { if frames > 0 {
fps := int(float64(frames) / elapsed.Seconds()) fps := int(float64(frames) / elapsed.Seconds())
fpsLatest.Store(int64(fps)) c.latest.Store(int64(fps))
} }
fpsWindowStart.Store(now.UnixNano()) c.windowStart.Store(now.UnixNano())
} }
return int(fpsLatest.Load()) return int(c.latest.Load())
}
func frameFPS(now time.Time) int {
return desktopFPSCounter.fps(now)
}
func backstageFrameFPS(now time.Time) int {
return backstageFPSCounter.fps(now)
} }
func shouldLogFrame(now time.Time) bool { func shouldLogFrame(now time.Time) bool {
@@ -693,7 +706,6 @@ func CleanupDesktopStream() {
ResetDesktopCapture() ResetDesktopCapture()
ResetPrev() ResetPrev()
resetHEVCEncoder() resetHEVCEncoder()
ResetFrameSlots()
} }
func ResetPrevbackstage() { func ResetPrevbackstage() {
@@ -1903,7 +1915,7 @@ func captureAndSendVirtual(ctx context.Context, env *rt.Env) error {
} }
now := time.Now() now := time.Now()
fps := frameFPS(now) fps := backstageFrameFPS(now)
if fps <= 0 { if fps <= 0 {
fps = 1 fps = 1
} }
@@ -1954,7 +1966,7 @@ func captureAndSendVirtual(ctx context.Context, env *rt.Env) error {
func virtualSendCompletedFrame(ctx context.Context, env *rt.Env, frame wire.Frame, t0 time.Time, captureDur, encodeDur time.Duration) error { func virtualSendCompletedFrame(ctx context.Context, env *rt.Env, frame wire.Frame, t0 time.Time, captureDur, encodeDur time.Duration) error {
now := time.Now() now := time.Now()
fps := frameFPS(now) fps := backstageFrameFPS(now)
if fps <= 0 { if fps <= 0 {
fps = 1 fps = 1
} }
@@ -2053,7 +2065,7 @@ func captureAndSendbackstage(ctx context.Context, env *rt.Env) error {
} }
now := time.Now() now := time.Now()
fps := frameFPS(now) fps := backstageFrameFPS(now)
if fps <= 0 { if fps <= 0 {
fps = 1 fps = 1
} }
@@ -11,3 +11,9 @@ func CaptureDisplayRGBABitBlt(display int) (*image.RGBA, error) {
func CaptureDisplayRGBAPreferBitBlt(display int) (*image.RGBA, error) { func CaptureDisplayRGBAPreferBitBlt(display int) (*image.RGBA, error) {
return CaptureDisplayRGBA(display) return CaptureDisplayRGBA(display)
} }
func CaptureDisplayRGBAThumbnail(display, _ int) (*image.RGBA, error) {
return CaptureDisplayRGBA(display)
}
func CleanupThumbnailCapture() {}
@@ -15,3 +15,13 @@ func CaptureDisplayRGBAPreferBitBlt(display int) (*image.RGBA, error) {
} }
return CaptureDisplayRGBA(display) return CaptureDisplayRGBA(display)
} }
func CaptureDisplayRGBAThumbnail(display, maxHeight int) (*image.RGBA, error) {
return captureDisplayThumbnailBitBlt(display, maxHeight)
}
func CleanupThumbnailCapture() {
thumbnailMu.Lock()
thumbnailState.reset()
thumbnailMu.Unlock()
}
@@ -361,6 +361,27 @@ func TestFrameFPS(t *testing.T) {
t.Logf("Calculated FPS: %d", fps) t.Logf("Calculated FPS: %d", fps)
} }
func TestDesktopAndBackstageFPSCountersAreIndependent(t *testing.T) {
var desktop frameRateCounter
var backstage frameRateCounter
start := time.Unix(100, 0)
desktop.fps(start)
backstage.fps(start)
for i := 1; i <= 59; i++ {
desktop.fps(start.Add(time.Duration(i) * time.Second / 60))
}
desktopFPS := desktop.fps(start.Add(time.Second))
backstageFPS := backstage.fps(start.Add(time.Second))
if desktopFPS < 55 || desktopFPS > 65 {
t.Fatalf("expected desktop counter near 60 fps, got %d", desktopFPS)
}
if backstageFPS != 2 {
t.Fatalf("expected backstage counter to contain only its two samples, got %d fps", backstageFPS)
}
}
func TestJPEGQuality(t *testing.T) { func TestJPEGQuality(t *testing.T) {
quality := desktopJPEGQuality() quality := desktopJPEGQuality()
if quality < 1 || quality > 100 { if quality < 1 || quality > 100 {
@@ -570,6 +591,60 @@ func TestSetFrameFlowTargetFPSEnvOverride(t *testing.T) {
} }
} }
func TestStoppingDesktopPreservesBackstageFrameFlow(t *testing.T) {
t.Setenv("OVERLORD_DESKTOP_IN_FLIGHT_FRAMES", "")
frameStreamsMu.Lock()
frameStreams = make(map[string]int)
frameStreamsMu.Unlock()
ResetFrameSlots()
t.Cleanup(func() {
StopFrameFlowStream("desktop")
StopFrameFlowStream("backstage")
})
StartFrameFlowStream("desktop", 60)
StartFrameFlowStream("backstage", 60)
if got := activeFrameSlotLimit(); got != 8 {
t.Fatalf("expected combined 120 fps slot limit 8, got %d", got)
}
inFlightFrames.Store(2)
frameAckSeen.Store(true)
StopFrameFlowStream("desktop")
if got := activeFrameSlotLimit(); got != 4 {
t.Fatalf("expected remaining backstage 60 fps slot limit 4, got %d", got)
}
if got := inFlightFrames.Load(); got != 2 {
t.Fatalf("desktop stop reset %d backstage in-flight frames", got)
}
if !frameAckSeen.Load() {
t.Fatal("desktop stop disabled ACK flow control while backstage remained active")
}
StopFrameFlowStream("backstage")
if got := inFlightFrames.Load(); got != 0 {
t.Fatalf("last stream stop left %d in-flight frames", got)
}
if frameAckSeen.Load() {
t.Fatal("last stream stop did not reset ACK flow state")
}
}
func TestInactiveStreamTargetDoesNotChangeActiveFrameFlow(t *testing.T) {
t.Setenv("OVERLORD_DESKTOP_IN_FLIGHT_FRAMES", "")
frameStreamsMu.Lock()
frameStreams = make(map[string]int)
frameStreamsMu.Unlock()
ResetFrameSlots()
t.Cleanup(func() { StopFrameFlowStream("backstage") })
StartFrameFlowStream("backstage", 60)
UpdateFrameFlowStream("desktop", 240)
if got := activeFrameSlotLimit(); got != 4 {
t.Fatalf("inactive desktop target changed backstage slot limit to %d", got)
}
}
func TestRequestDesktopFullFrameRequestsWebRTCKeyframe(t *testing.T) { func TestRequestDesktopFullFrameRequestsWebRTCKeyframe(t *testing.T) {
_ = webrtcpub.ConsumeKeyframeRequest(webrtcpub.KindDesktop) _ = webrtcpub.ConsumeKeyframeRequest(webrtcpub.KindDesktop)
@@ -4,6 +4,7 @@ import (
"os" "os"
"strconv" "strconv"
"strings" "strings"
"sync"
"sync/atomic" "sync/atomic"
"time" "time"
) )
@@ -15,6 +16,8 @@ var (
frameAckSeen atomic.Bool frameAckSeen atomic.Bool
lastAckNano atomic.Int64 lastAckNano atomic.Int64
maxFrameSlots atomic.Int64 maxFrameSlots atomic.Int64
frameStreamsMu sync.Mutex
frameStreams = make(map[string]int)
) )
func AcquireFrameSlot() bool { func AcquireFrameSlot() bool {
@@ -54,6 +57,48 @@ func ResetFrameSlots() {
} }
func SetFrameFlowTargetFPS(fps int) { func SetFrameFlowTargetFPS(fps int) {
maxFrameSlots.Store(frameSlotLimitForFPS(fps))
}
func StartFrameFlowStream(stream string, fps int) {
frameStreamsMu.Lock()
frameStreams[stream] = fps
applyActiveFrameStreamLimitLocked()
frameStreamsMu.Unlock()
}
func UpdateFrameFlowStream(stream string, fps int) {
frameStreamsMu.Lock()
if _, active := frameStreams[stream]; active {
frameStreams[stream] = fps
applyActiveFrameStreamLimitLocked()
}
frameStreamsMu.Unlock()
}
func StopFrameFlowStream(stream string) {
frameStreamsMu.Lock()
delete(frameStreams, stream)
if len(frameStreams) == 0 {
ResetFrameSlots()
maxFrameSlots.Store(defaultMaxInFlightFrames)
} else {
applyActiveFrameStreamLimitLocked()
}
frameStreamsMu.Unlock()
}
func applyActiveFrameStreamLimitLocked() {
combinedFPS := 0
for _, fps := range frameStreams {
if fps > 0 {
combinedFPS += fps
}
}
maxFrameSlots.Store(frameSlotLimitForFPS(combinedFPS))
}
func frameSlotLimitForFPS(fps int) int64 {
limit := defaultMaxInFlightFrames limit := defaultMaxInFlightFrames
switch { switch {
case fps >= 180: case fps >= 180:
@@ -75,7 +120,7 @@ func SetFrameFlowTargetFPS(fps int) {
} }
} }
} }
maxFrameSlots.Store(limit) return limit
} }
func activeFrameSlotLimit() int64 { func activeFrameSlotLimit() int64 {
@@ -139,11 +139,11 @@ func CleanupVirtualMode() {
virtualDupState.mu.Unlock() virtualDupState.mu.Unlock()
ResetPrevbackstage() ResetPrevbackstage()
resetH264D3D11TextureEncoder("backstage") resetH264D3D11TextureEncoder("backstage")
captureMu.Lock() backstageCaptureMu.Lock()
backstageFreeCapCache() backstageFreeCapCache()
backstageClearWindowCache() backstageClearWindowCache()
virtualCapImg = nil virtualCapImg = nil
captureMu.Unlock() backstageCaptureMu.Unlock()
virtualInputMu.Lock() virtualInputMu.Lock()
virtualShiftDown = false virtualShiftDown = false
@@ -2393,8 +2393,8 @@ var (
) )
func VirtualCaptureDisplayFallback() (*image.RGBA, error) { func VirtualCaptureDisplayFallback() (*image.RGBA, error) {
captureMu.Lock() backstageCaptureMu.Lock()
defer captureMu.Unlock() defer backstageCaptureMu.Unlock()
setDPIAware() setDPIAware()
+47 -21
View File
@@ -165,9 +165,9 @@ var (
maxResHeight atomic.Int64 // 0 = default (1080), -1 = native, >0 = specific max height maxResHeight atomic.Int64 // 0 = default (1080), -1 = native, >0 = specific max height
lastCapScale atomic.Uint64 lastCapScale atomic.Uint64
state capState state capState
thumbnailState capState
captureMu sync.Mutex captureMu sync.Mutex
thumbnailMu sync.Mutex
stateCreatedAt time.Time
stateRefreshInterval = 30 * time.Second stateRefreshInterval = 30 * time.Second
) )
@@ -186,9 +186,18 @@ var captureDisplayFn = func(display int) (*image.RGBA, error) {
} }
func captureDisplayBitBlt(display int) (*image.RGBA, error) { func captureDisplayBitBlt(display int) (*image.RGBA, error) {
captureMu.Lock() captureMu.Lock()
defer captureMu.Unlock() defer captureMu.Unlock()
return captureDisplayBitBltWithState(display, &state, 0)
}
func captureDisplayThumbnailBitBlt(display, maxHeight int) (*image.RGBA, error) {
thumbnailMu.Lock()
defer thumbnailMu.Unlock()
return captureDisplayBitBltWithState(display, &thumbnailState, maxHeight)
}
func captureDisplayBitBltWithState(display int, captureState *capState, forcedMaxHeight int) (*image.RGBA, error) {
setDPIAware() setDPIAware()
@@ -223,7 +232,10 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
} }
} }
userScale := effectiveScale(srcW, srcH) userScale := scaleForMaxHeight(srcH, forcedMaxHeight)
if forcedMaxHeight <= 0 {
userScale = effectiveScale(srcW, srcH)
}
dstW := int(float64(srcW) * userScale) dstW := int(float64(srcW) * userScale)
dstH := int(float64(srcH) * userScale) dstH := int(float64(srcH) * userScale)
if dstW <= 0 || dstH <= 0 { if dstW <= 0 || dstH <= 0 {
@@ -265,7 +277,7 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
} }
captureWithDC := func() (*image.RGBA, time.Duration, time.Duration, time.Duration, error) { captureWithDC := func() (*image.RGBA, time.Duration, time.Duration, time.Duration, error) {
hdcMem, _, buf, stride, err := state.ensure(hdcScreen, capW, capH) hdcMem, _, buf, stride, err := captureState.ensure(hdcScreen, capW, capH)
if err != nil { if err != nil {
return nil, 0, 0, 0, err return nil, 0, 0, 0, err
} }
@@ -278,8 +290,8 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
captured = bitBltWithFallback(hdcMem, hdcScreen, bounds, capW, capH) captured = bitBltWithFallback(hdcMem, hdcScreen, bounds, capW, capH)
} }
if !captured { if !captured {
state.reset() captureState.reset()
hdcMem, _, buf, stride, err = state.ensure(hdcScreen, capW, capH) hdcMem, _, buf, stride, err = captureState.ensure(hdcScreen, capW, capH)
if err != nil { if err != nil {
return nil, 0, 0, 0, err return nil, 0, 0, 0, err
} }
@@ -328,13 +340,13 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
return img, bitDur, dibDur, convDur, nil return img, bitDur, dibDur, convDur, nil
} }
if !stateCreatedAt.IsZero() && time.Since(stateCreatedAt) > stateRefreshInterval { if !captureState.createdAt.IsZero() && time.Since(captureState.createdAt) > stateRefreshInterval {
state.reset() captureState.reset()
} }
img, bitDur, dibDur, convDur, err := captureWithDC() img, bitDur, dibDur, convDur, err := captureWithDC()
if err != nil && fromCreateDC { if err != nil && fromCreateDC {
state.reset() captureState.reset()
closeDC() closeDC()
hdcScreen = getDC(0) hdcScreen = getDC(0)
fromCreateDC = false fromCreateDC = false
@@ -344,7 +356,7 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
img, bitDur, dibDur, convDur, err = captureWithDC() img, bitDur, dibDur, convDur, err = captureWithDC()
} }
if err != nil { if err != nil {
state.reset() captureState.reset()
return nil, err return nil, err
} }
@@ -447,14 +459,15 @@ func resizeNearest(src *image.RGBA, w, h int) *image.RGBA {
} }
type capState struct { type capState struct {
mu sync.Mutex mu sync.Mutex
hdcMem uintptr hdcMem uintptr
hbmp uintptr hbmp uintptr
buf []byte buf []byte
stride int stride int
w int w int
h int h int
bits unsafe.Pointer bits unsafe.Pointer
createdAt time.Time
} }
func (s *capState) ensure(hdcScreen uintptr, w, h int) (uintptr, uintptr, []byte, int, error) { func (s *capState) ensure(hdcScreen uintptr, w, h int) (uintptr, uintptr, []byte, int, error) {
@@ -466,7 +479,7 @@ func (s *capState) ensure(hdcScreen uintptr, w, h int) (uintptr, uintptr, []byte
if s.hdcMem == 0 { if s.hdcMem == 0 {
return 0, 0, nil, 0, syscall.EINVAL return 0, 0, nil, 0, syscall.EINVAL
} }
stateCreatedAt = time.Now() s.createdAt = time.Now()
} }
if s.w != w || s.h != h || s.hbmp == 0 { if s.w != w || s.h != h || s.hbmp == 0 {
@@ -519,7 +532,7 @@ func (s *capState) reset() {
s.w = 0 s.w = 0
s.h = 0 s.h = 0
s.bits = nil s.bits = nil
stateCreatedAt = time.Time{} s.createdAt = time.Time{}
} }
func captureScale() float64 { func captureScale() float64 {
@@ -554,6 +567,19 @@ func effectiveScale(srcW, srcH int) float64 {
return s return s
} }
func scaleForMaxHeight(srcH, maxHeight int) float64 {
s := captureScale()
if s > 1 {
s = 1
}
if maxHeight > 0 && srcH > maxHeight {
if capScale := float64(maxHeight) / float64(srcH); capScale < s {
s = capScale
}
}
return s
}
func storeLastScale(s float64) { func storeLastScale(s float64) {
lastCapScale.Store(math.Float64bits(s)) lastCapScale.Store(math.Float64bits(s))
} }
@@ -33,13 +33,15 @@ func backstageStart(ctx context.Context, env *rt.Env, autoStartExplorer bool) er
fps := activebackstageTargetFPS() fps := activebackstageTargetFPS()
interval := time.Second / time.Duration(fps) interval := time.Second / time.Duration(fps)
capture.SetBackstageH264TargetFPS(fps) capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
log.Printf("backstage: starting stream (target fps %d)", fps) log.Printf("backstage: starting stream (target fps %d)", fps)
if err := capture.InitializebackstageDesktop(); err != nil { if err := capture.InitializebackstageDesktop(); err != nil {
log.Printf("backstage: failed to initialize hidden desktop: %v", err) log.Printf("backstage: failed to initialize hidden desktop: %v", err)
return err return err
} }
capture.StartFrameFlowStream("backstage", fps)
defer capture.StopFrameFlowStream("backstage")
defer capture.CleanupbackstageDesktop()
if autoStartExplorer { if autoStartExplorer {
goSafe("backstage auto-start explorer", nil, func() { goSafe("backstage auto-start explorer", nil, func() {
@@ -56,21 +58,19 @@ func backstageStart(ctx context.Context, env *rt.Env, autoStartExplorer bool) er
select { select {
case <-ctx.Done(): case <-ctx.Done():
log.Printf("backstage: stopping stream") log.Printf("backstage: stopping stream")
capture.CleanupbackstageDesktop()
return nil return nil
case <-ticker.C: case <-ticker.C:
fps := activebackstageTargetFPS() fps := activebackstageTargetFPS()
if fps != currentFPS { if fps != currentFPS {
currentFPS = fps currentFPS = fps
capture.SetBackstageH264TargetFPS(fps) capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps) capture.UpdateFrameFlowStream("backstage", fps)
ticker.Reset(time.Second / time.Duration(fps)) ticker.Reset(time.Second / time.Duration(fps))
log.Printf("backstage: target fps changed to %d", fps) log.Printf("backstage: target fps changed to %d", fps)
} }
if err := capture.Nowbackstage(ctx, env); err != nil { if err := capture.Nowbackstage(ctx, env); err != nil {
if ctx.Err() != nil { if ctx.Err() != nil {
log.Printf("backstage: stopping stream") log.Printf("backstage: stopping stream")
capture.CleanupbackstageDesktop()
return nil return nil
} }
log.Printf("backstage: capture error: %v", err) log.Printf("backstage: capture error: %v", err)
@@ -91,7 +91,7 @@ func SetbackstageTargetFPS(fps int) int {
fps = clampDesktopTargetFPS(fps) fps = clampDesktopTargetFPS(fps)
backstageTargetFPS.Store(int64(fps)) backstageTargetFPS.Store(int64(fps))
capture.SetBackstageH264TargetFPS(fps) capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps) capture.UpdateFrameFlowStream("backstage", fps)
return fps return fps
} }
+27 -2
View File
@@ -595,17 +595,20 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
case "screenshot": case "screenshot":
payload, _ := envelope["payload"].(map[string]interface{}) payload, _ := envelope["payload"].(map[string]interface{})
allDisplays := false allDisplays := false
dashboardThumbnail := false
if payload != nil { if payload != nil {
mode, _ := payload["mode"].(string)
dashboardThumbnail = mode == "thumbnail"
if v, ok := payload["allDisplays"].(bool); ok && v { if v, ok := payload["allDisplays"].(bool); ok && v {
allDisplays = true allDisplays = true
} else if mode, ok := payload["mode"].(string); ok && mode == "notification" { } else if mode == "notification" || mode == "thumbnail" {
allDisplays = true allDisplays = true
} }
} }
if goruntime.GOOS == "windows" { if goruntime.GOOS == "windows" {
allDisplays = true allDisplays = true
} }
return HandleScreenshot(ctx, env, cmdID, allDisplays) return HandleScreenshot(ctx, env, cmdID, allDisplays, dashboardThumbnail)
case "plugin_load": case "plugin_load":
payload, _ := envelope["payload"].(map[string]interface{}) payload, _ := envelope["payload"].(map[string]interface{})
if payload == nil { if payload == nil {
@@ -1137,6 +1140,22 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
} }
capture.RequestDesktopRecoveryAfterBackstageStart() capture.RequestDesktopRecoveryAfterBackstageStart()
if VirtualMode { if VirtualMode {
// The hidden-desktop and virtual-display implementations share the
// backstage encoder, frame history, and fallback capture caches. Stop
// the other implementation before switching; ordinary remote desktop
// remains independent and continues running.
env.BackstageMu.Lock()
if env.BackstageCancel != nil {
env.BackstageCancel()
waitStreamStop(env.BackstageDone, "backstage")
env.BackstageCancel = nil
env.BackstageDone = nil
}
env.BackstageMouseControl = false
env.BackstageKeyboardControl = false
env.BackstageCursorCapture = false
clearbackstageInputQueue()
env.BackstageMu.Unlock()
env.VirtualMu.Lock() env.VirtualMu.Lock()
if env.VirtualCancel != nil { if env.VirtualCancel != nil {
env.VirtualCancel() env.VirtualCancel()
@@ -1165,6 +1184,12 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
sendCommandResultSafe(env, cmdID, true, "") sendCommandResultSafe(env, cmdID, true, "")
return nil return nil
} }
env.VirtualMu.Lock()
if env.VirtualCancel != nil {
log.Printf("hidden: stopping virtual stream before starting backstage desktop")
stopVirtualStreamLocked(env)
}
env.VirtualMu.Unlock()
env.BackstageMu.Lock() env.BackstageMu.Lock()
if env.BackstageCancel != nil { if env.BackstageCancel != nil {
env.BackstageCancel() env.BackstageCancel()
@@ -34,7 +34,8 @@ func DesktopStart(ctx context.Context, env *rt.Env) error {
fps := activeDesktopTargetFPS() fps := activeDesktopTargetFPS()
interval := time.Second / time.Duration(fps) interval := time.Second / time.Duration(fps)
capture.SetDesktopH264TargetFPS(fps) capture.SetDesktopH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps) capture.StartFrameFlowStream("desktop", fps)
defer capture.StopFrameFlowStream("desktop")
log.Printf("desktop: starting stream (target fps %d)", fps) log.Printf("desktop: starting stream (target fps %d)", fps)
ticker := time.NewTicker(interval) ticker := time.NewTicker(interval)
defer ticker.Stop() defer ticker.Stop()
@@ -49,7 +50,7 @@ func DesktopStart(ctx context.Context, env *rt.Env) error {
if fps != currentFPS { if fps != currentFPS {
currentFPS = fps currentFPS = fps
capture.SetDesktopH264TargetFPS(fps) capture.SetDesktopH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps) capture.UpdateFrameFlowStream("desktop", fps)
ticker.Reset(time.Second / time.Duration(fps)) ticker.Reset(time.Second / time.Duration(fps))
log.Printf("desktop: target fps changed to %d", fps) log.Printf("desktop: target fps changed to %d", fps)
} }
@@ -76,7 +77,7 @@ func SetDesktopTargetFPS(fps int) int {
fps = clampDesktopTargetFPS(fps) fps = clampDesktopTargetFPS(fps)
desktopTargetFPS.Store(int64(fps)) desktopTargetFPS.Store(int64(fps))
capture.SetDesktopH264TargetFPS(fps) capture.SetDesktopH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps) capture.UpdateFrameFlowStream("desktop", fps)
return fps return fps
} }
@@ -9,6 +9,9 @@ import (
"image/jpeg" "image/jpeg"
"log" "log"
goruntime "runtime" goruntime "runtime"
"runtime/debug"
"sync"
"time"
"overlord-client/cmd/agent/capture" "overlord-client/cmd/agent/capture"
rt "overlord-client/cmd/agent/runtime" rt "overlord-client/cmd/agent/runtime"
@@ -18,12 +21,18 @@ import (
) )
var ( var (
monitorCountFn = capture.MonitorCount monitorCountFn = capture.MonitorCount
displayBoundsFn = capture.DisplayBounds displayBoundsFn = capture.DisplayBounds
captureRectFn = screenshot.CaptureRect captureRectFn = screenshot.CaptureRect
thumbnailCaptureDisplayFn = capture.CaptureDisplayRGBAThumbnail
thumbnailCaptureCleanupFn = capture.CleanupThumbnailCapture
screenshotScavengeMu sync.Mutex
screenshotScavengeTimer *time.Timer
) )
func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplays bool) error { const dashboardThumbnailMaxHeight = 1080
func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplays, dashboardThumbnail bool) error {
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2 //garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
if allDisplays { if allDisplays {
@@ -44,7 +53,16 @@ func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplay
} }
}() }()
img, _, bounds, err := captureScreenshotImage(allDisplays) var img *image.RGBA
var bounds image.Rectangle
var err error
if dashboardThumbnail && goruntime.GOOS == "windows" {
img, _, bounds, err = captureDashboardThumbnailWindows(allDisplays)
thumbnailCaptureCleanupFn()
defer scheduleScreenshotHeapScavenge()
} else {
img, _, bounds, err = captureScreenshotImage(allDisplays)
}
if err != nil { if err != nil {
log.Printf("screenshot: capture failed: %v", err) log.Printf("screenshot: capture failed: %v", err)
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{ return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
@@ -56,7 +74,11 @@ func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplay
} }
var buf bytes.Buffer var buf bytes.Buffer
opts := &jpeg.Options{Quality: 92} quality := 92
if dashboardThumbnail {
quality = 85
}
opts := &jpeg.Options{Quality: quality}
if err := jpeg.Encode(&buf, img, opts); err != nil { if err := jpeg.Encode(&buf, img, opts); err != nil {
log.Printf("screenshot: jpeg encode failed: %v", err) log.Printf("screenshot: jpeg encode failed: %v", err)
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{ return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
@@ -91,6 +113,110 @@ func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplay
}) })
} }
func scheduleScreenshotHeapScavenge() {
screenshotScavengeMu.Lock()
defer screenshotScavengeMu.Unlock()
if screenshotScavengeTimer != nil {
screenshotScavengeTimer.Stop()
}
screenshotScavengeTimer = time.AfterFunc(time.Second, func() {
screenshotScavengeMu.Lock()
screenshotScavengeTimer = nil
screenshotScavengeMu.Unlock()
debug.FreeOSMemory()
})
}
func captureDashboardThumbnailWindows(allDisplays bool) (*image.RGBA, int, image.Rectangle, error) {
monitorCount := monitorCountFn()
if monitorCount <= 0 {
return nil, 0, image.Rectangle{}, errors.New("no active displays available")
}
if !allDisplays {
img, err := thumbnailCaptureDisplayFn(0, dashboardThumbnailMaxHeight)
if err != nil {
return nil, 0, image.Rectangle{}, err
}
if img == nil {
return nil, 0, image.Rectangle{}, errors.New("dashboard thumbnail capture returned no image")
}
return img, 0, img.Rect, nil
}
type thumbnailPart struct {
bounds image.Rectangle
img *image.RGBA
}
parts := make([]thumbnailPart, 0, monitorCount)
minX, minY := int(1e9), int(1e9)
monitorBounds := make([]image.Rectangle, monitorCount)
maxMonitorHeight := 0
for i := 0; i < monitorCount; i++ {
bounds := displayBoundsFn(i)
monitorBounds[i] = bounds
if bounds.Dy() > maxMonitorHeight {
maxMonitorHeight = bounds.Dy()
}
if bounds.Min.X < minX {
minX = bounds.Min.X
}
if bounds.Min.Y < minY {
minY = bounds.Min.Y
}
}
if maxMonitorHeight <= 0 {
return nil, 0, image.Rectangle{}, errors.New("invalid dashboard monitor bounds")
}
layoutScale := 1.0
if maxMonitorHeight > dashboardThumbnailMaxHeight {
layoutScale = float64(dashboardThumbnailMaxHeight) / float64(maxMonitorHeight)
}
for i, bounds := range monitorBounds {
monitorMaxHeight := int(float64(bounds.Dy()) * layoutScale)
if monitorMaxHeight < 1 {
monitorMaxHeight = 1
}
img, err := thumbnailCaptureDisplayFn(i, monitorMaxHeight)
if err != nil {
return nil, 0, image.Rectangle{}, err
}
if img == nil {
return nil, 0, image.Rectangle{}, errors.New("dashboard thumbnail capture returned no image")
}
parts = append(parts, thumbnailPart{bounds: bounds, img: img})
}
type positionedPart struct {
img *image.RGBA
offX int
offY int
}
positioned := make([]positionedPart, len(parts))
canvasW, canvasH := 0, 0
for i, part := range parts {
positioned[i] = positionedPart{
img: part.img,
offX: int(float64(part.bounds.Min.X-minX) * layoutScale),
offY: int(float64(part.bounds.Min.Y-minY) * layoutScale),
}
if right := positioned[i].offX + part.img.Rect.Dx(); right > canvasW {
canvasW = right
}
if bottom := positioned[i].offY + part.img.Rect.Dy(); bottom > canvasH {
canvasH = bottom
}
}
if canvasW <= 0 || canvasH <= 0 {
return nil, 0, image.Rectangle{}, errors.New("invalid dashboard thumbnail bounds")
}
canvas := image.NewRGBA(image.Rect(0, 0, canvasW, canvasH))
for _, part := range positioned {
dst := image.Rect(part.offX, part.offY, part.offX+part.img.Rect.Dx(), part.offY+part.img.Rect.Dy())
draw.Draw(canvas, dst, part.img, part.img.Rect.Min, draw.Src)
}
return canvas, 0, canvas.Rect, nil
}
func captureScreenshotImage(allDisplays bool) (*image.RGBA, int, image.Rectangle, error) { func captureScreenshotImage(allDisplays bool) (*image.RGBA, int, image.Rectangle, error) {
if goruntime.GOOS == "windows" { if goruntime.GOOS == "windows" {
if img, displayIndex, bounds, err := captureScreenshotImageWindows(allDisplays); err == nil && img != nil { if img, displayIndex, bounds, err := captureScreenshotImageWindows(allDisplays); err == nil && img != nil {
@@ -136,7 +136,7 @@ func TestHandleScreenshot_RepeatedRequestsDoNotEmitLiveFrames(t *testing.T) {
ctx := context.Background() ctx := context.Background()
writer := &testWriter{} writer := &testWriter{}
for i, commandID := range []string{"cmd-1", "cmd-2"} { for i, commandID := range []string{"cmd-1", "cmd-2"} {
if err := HandleScreenshot(ctx, &rt.Env{Conn: writer}, commandID, false); err != nil { if err := HandleScreenshot(ctx, &rt.Env{Conn: writer}, commandID, false, false); err != nil {
t.Fatalf("HandleScreenshot request %d failed: %v", i+1, err) t.Fatalf("HandleScreenshot request %d failed: %v", i+1, err)
} }
} }
@@ -162,3 +162,47 @@ func TestHandleScreenshot_RepeatedRequestsDoNotEmitLiveFrames(t *testing.T) {
} }
} }
} }
func TestDashboardThumbnailCapsMultiMonitorCaptureBeforeStitching(t *testing.T) {
origMonitorCountFn := monitorCountFn
origBoundsFn := displayBoundsFn
origThumbnailCaptureFn := thumbnailCaptureDisplayFn
t.Cleanup(func() {
monitorCountFn = origMonitorCountFn
displayBoundsFn = origBoundsFn
thumbnailCaptureDisplayFn = origThumbnailCaptureFn
})
monitorCountFn = func() int { return 2 }
displayBoundsFn = func(idx int) image.Rectangle {
if idx == 0 {
return image.Rect(0, 0, 3840, 2160)
}
return image.Rect(3840, 235, 4920, 2155)
}
requestedHeights := make([]int, 0, 2)
thumbnailCaptureDisplayFn = func(display, maxHeight int) (*image.RGBA, error) {
requestedHeights = append(requestedHeights, maxHeight)
if display == 0 {
return image.NewRGBA(image.Rect(0, 0, 1920, 1080)), nil
}
return image.NewRGBA(image.Rect(0, 0, 540, 960)), nil
}
img, _, bounds, err := captureDashboardThumbnailWindows(true)
if err != nil {
t.Fatalf("capture dashboard thumbnail: %v", err)
}
if img == nil {
t.Fatal("expected dashboard thumbnail image")
}
if len(requestedHeights) != 2 || requestedHeights[0] != 1080 || requestedHeights[1] != 960 {
t.Fatalf("unexpected per-monitor capture caps: %v", requestedHeights)
}
if bounds.Dx() != 2460 || bounds.Dy() != 1080 {
t.Fatalf("expected bounded 2460x1080 canvas, got %dx%d", bounds.Dx(), bounds.Dy())
}
if got := int64(len(img.Pix)); got >= int64(4920*2160*4) {
t.Fatalf("thumbnail canvas retained native virtual-desktop allocation: %d bytes", got)
}
}
@@ -17,13 +17,14 @@ func VirtualStart(ctx context.Context, env *rt.Env) error {
fps := activeVirtualTargetFPS() fps := activeVirtualTargetFPS()
interval := time.Second / time.Duration(fps) interval := time.Second / time.Duration(fps)
capture.SetBackstageH264TargetFPS(fps) capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
log.Printf("virtual: starting stream (target fps %d)", fps) log.Printf("virtual: starting stream (target fps %d)", fps)
if err := capture.InitializeVirtualMode(); err != nil { if err := capture.InitializeVirtualMode(); err != nil {
log.Printf("virtual: failed to initialize virtual mode: %v", err) log.Printf("virtual: failed to initialize virtual mode: %v", err)
return err return err
} }
capture.StartFrameFlowStream("virtual", fps)
defer capture.StopFrameFlowStream("virtual")
defer capture.CleanupVirtualMode() defer capture.CleanupVirtualMode()
ticker := time.NewTicker(interval) ticker := time.NewTicker(interval)
@@ -39,7 +40,7 @@ func VirtualStart(ctx context.Context, env *rt.Env) error {
if fps != currentFPS { if fps != currentFPS {
currentFPS = fps currentFPS = fps
capture.SetBackstageH264TargetFPS(fps) capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps) capture.UpdateFrameFlowStream("virtual", fps)
ticker.Reset(time.Second / time.Duration(fps)) ticker.Reset(time.Second / time.Duration(fps))
log.Printf("virtual: target fps changed to %d", fps) log.Printf("virtual: target fps changed to %d", fps)
} }
@@ -66,7 +67,7 @@ func SetVirtualTargetFPS(fps int) int {
fps = clampDesktopTargetFPS(fps) fps = clampDesktopTargetFPS(fps)
virtualTargetFPS.Store(int64(fps)) virtualTargetFPS.Store(int64(fps))
capture.SetBackstageH264TargetFPS(fps) capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps) capture.UpdateFrameFlowStream("virtual", fps)
return fps return fps
} }
@@ -314,7 +314,7 @@ export async function handleClientRoutes(
const beforeVersion = getThumbnailVersion(clientId); const beforeVersion = getThumbnailVersion(clientId);
if (target?.online) { if (target?.online) {
const commandId = uuidv4(); const commandId = uuidv4();
target.ws.send(encodeMessage({ type: "command", commandType: "screenshot", id: commandId, payload: { mode: "notification", allDisplays: true } })); target.ws.send(encodeMessage({ type: "command", commandType: "screenshot", id: commandId, payload: { mode: "thumbnail", allDisplays: true } }));
metrics.recordCommand("screenshot"); metrics.recordCommand("screenshot");
} }
const fresh = target?.online ? await waitForThumbnail(clientId, 2500) : false; const fresh = target?.online ? await waitForThumbnail(clientId, 2500) : false;