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")
procDWMShowWindow = user32.NewProc("ShowWindow")
procDWMUpdateWindow = user32.NewProc("UpdateWindow")
procDWMDestroyWindow = user32.NewProc("DestroyWindow")
backstageDWMStateMu sync.Mutex
backstageDWMHost uintptr
backstageDWMThumbnails map[uintptr]*backstageDWMThumbnail
@@ -460,11 +461,23 @@ func backstageCleanupDWMThumbnails() {
defer backstageDWMStateMu.Unlock()
backstageUnregisterDWMThumbnailsLocked()
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
}
func backstageAbandonDWMThumbnails() {
backstageDWMStateMu.Lock()
backstageDWMThumbnails = nil
backstageDWMOrder = nil
backstageDWMHost = 0
backstageDWMStateMu.Unlock()
}
func backstageUnregisterDWMThumbnailsLocked() {
for source, entry := range backstageDWMThumbnails {
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
import "testing"
import (
"testing"
"time"
)
func TestBackstageWindowCacheByteSizeAndLimit(t *testing.T) {
bytes, ok := backstageWindowCacheByteSize(3840, 2160)
@@ -52,3 +55,48 @@ func TestBackstagePrintWindowFallbackToggle(t *testing.T) {
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 (
backstageDesktopHandle uintptr
backstageDesktopMu sync.Mutex
backstageCaptureMu sync.Mutex
backstageDesktopName = "OverlordHiddenDesktop"
backstageInitialized bool
backstageOriginalDesktop uintptr
@@ -163,6 +164,7 @@ var (
backstageThreadErr error
backstageThreadReady chan struct{}
backstageThreadTasks chan backstageTask
backstageThreadDone chan struct{}
backstageWatchdogOnce sync.Once
backstageNoWindowLogNs atomic.Int64
backstageInputMu sync.Mutex
@@ -191,6 +193,7 @@ var (
backstageWinCachePrev []byte
backstageWinCacheBytes int64
backstageWinCacheSeq uint64
backstageHungWindows map[uintptr]struct{}
backstageCompHdcMem uintptr
backstageCompHbmp uintptr
@@ -215,6 +218,7 @@ const (
backstageTaskKeyUp
backstageTaskMouseWheel
backstageTaskAutoStartExplorer
backstageTaskShutdown
)
type backstageTask struct {
@@ -293,6 +297,7 @@ type backstageWinCacheEntry struct {
const (
backstageMaxWindowCacheEntries = 64
backstageMaxWindowCacheBytes = int64(256 << 20)
backstagePrintWindowTimeout = 250 * time.Millisecond
)
type keybdInput struct {
@@ -388,8 +393,11 @@ func CleanupbackstageDesktop() {
resetH264D3D11TextureEncoder("backstage")
backstageCleanupFrameReaders()
captureMu.Lock()
backstageCleanupDWMThumbnails()
dwmCleaned, workerStopped := shutdownbackstageThreadLocked(time.Second)
if !dwmCleaned {
backstageAbandonDWMThumbnails()
}
backstageCaptureMu.Lock()
backstageFreeCapCache()
@@ -407,7 +415,7 @@ func CleanupbackstageDesktop() {
backstageCompBits = nil
backstageCompW = 0
backstageCompH = 0
captureMu.Unlock()
backstageCaptureMu.Unlock()
backstageInputMu.Lock()
backstageShiftDown = false
@@ -420,14 +428,10 @@ func CleanupbackstageDesktop() {
backstageInputMu.Unlock()
backstageLastScale.Store(0)
if backstageDesktopHandle != 0 {
if backstageOriginalDesktop != 0 {
procSetThreadDesktop.Call(backstageOriginalDesktop)
}
if backstageDesktopHandle != 0 && workerStopped {
procCloseDesktop.Call(backstageDesktopHandle)
backstageDesktopHandle = 0
}
backstageDesktopHandle = 0
backstageInitialized = false
backstageExplorerStarted = false
@@ -436,16 +440,52 @@ func CleanupbackstageDesktop() {
resetWinUI3Cache()
resetInputSiteCache()
if backstageThreadTasks != nil {
close(backstageThreadTasks)
backstageThreadTasks = nil
}
backstageThreadTasks = nil
backstageThreadReady = nil
backstageThreadDone = nil
backstageThreadErr = nil
backstageThreadOnce = 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) {
backstageCursorEnabled = enabled
}
@@ -476,31 +516,38 @@ func ensurebackstageThread() error {
}
backstageThreadOnce.Do(func() {
backstageThreadReady = make(chan struct{})
backstageThreadTasks = make(chan backstageTask)
ready := make(chan struct{})
tasks := make(chan backstageTask)
done := make(chan struct{})
backstageThreadReady = ready
backstageThreadTasks = tasks
backstageThreadDone = done
backstageWatchdogOnce.Do(func() {
go func() {
defer recoverAndLog("backstage watchdog", nil)
backstageThreadWatchdog()
}()
})
go func(handle uintptr) {
go func(handle, originalDesktop uintptr) {
defer recoverAndLog("backstage desktop thread", nil)
runtime.LockOSThread()
defer runtime.UnlockOSThread()
defer func() {
if originalDesktop != 0 {
procSetThreadDesktop.Call(originalDesktop)
}
close(done)
runtime.UnlockOSThread()
}()
r, _, err := procSetThreadDesktop.Call(handle)
if r == 0 {
backstageThreadErr = fmt.Errorf("failed to set thread desktop: %v", err)
close(backstageThreadReady)
for task := range backstageThreadTasks {
task.resp <- backstageTaskResult{err: backstageThreadErr}
}
close(ready)
return
}
close(backstageThreadReady)
for task := range backstageThreadTasks {
close(ready)
for task := range tasks {
start := time.Now()
backstageCurrentTaskID.Store(task.id)
backstageCurrentTaskKind.Store(int64(task.kind))
@@ -530,6 +577,8 @@ func ensurebackstageThread() error {
result.err = backstageMouseWheelOnThread(task.delta)
case backstageTaskAutoStartExplorer:
result.err = backstageAutoStartExplorerOnThread()
case backstageTaskShutdown:
backstageCleanupDWMThumbnails()
default:
result.img, result.err = BackstageCaptureDisplayOnThread(task.display)
}
@@ -547,8 +596,11 @@ func ensurebackstageThread() error {
backstageCurrentTaskKind.Store(-1)
backstageCurrentTaskID.Store(0)
task.resp <- result
if task.kind == backstageTaskShutdown {
return
}
}
}(desktopHandle)
}(desktopHandle, backstageOriginalDesktop)
})
if backstageThreadReady != nil {
@@ -774,6 +826,8 @@ func backstageTaskKindName(kind backstageTaskKind) string {
return "mouse_wheel"
case backstageTaskAutoStartExplorer:
return "auto_start_explorer"
case backstageTaskShutdown:
return "shutdown"
default:
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) {
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
captureMu.Lock()
defer captureMu.Unlock()
backstageCaptureMu.Lock()
defer backstageCaptureMu.Unlock()
setDPIAware()
@@ -1720,6 +1774,11 @@ func drawbackstageWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, tar
backstageRemoveWindowCacheEntry(h)
}
}
for h := range backstageHungWindows {
if !alive[h] {
delete(backstageHungWindows, h)
}
}
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() {
for hwnd := range backstageWinCache {
backstageRemoveWindowCacheEntry(hwnd)
@@ -1834,6 +1905,8 @@ var backstageDXGIEnabled atomic.Bool
var backstageUIAEnabled atomic.Bool
var backstagePrintWindowFallbackEnabled atomic.Bool
var backstagePrintWindowFallbackLogNs atomic.Int64
var backstagePrintWindowTimeoutLogNs atomic.Int64
var backstagePrintWindowFn = printWindow
func init() {
backstageDXGIEnabled.Store(false) // disabled by default
@@ -1865,6 +1938,49 @@ func GetbackstagePrintWindowFallbackEnabled() bool {
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 {
if !backstageDXGIEnabled.Load() {
return false
@@ -1986,7 +2102,10 @@ func drawbackstageWindow(hdcScreen, hwnd uintptr, bounds image.Rectangle, target
return false
}
if !printWindow(hwnd, entry.hdcMem, PW_RENDERFULLCONTENT) {
if !backstagePrintWindowWithTimeout(hwnd, entry) {
if backstageWinCache[hwnd] != entry {
return false
}
entry.lastOK = false
entry.age++
return false
+35 -23
View File
@@ -530,14 +530,13 @@ const desktopHardwareH264FallbackCooldown = 30 * time.Second
var consecutiveCaptureFails atomic.Int64
var (
fpsWindowStart atomic.Int64
fpsCount atomic.Int64
fpsLatest atomic.Int64
lastFrameLog atomic.Int64
metricsOnce sync.Once
metricsEnabled bool
lastKeyframe atomic.Int64
fullNextFrames atomic.Int64
desktopFPSCounter frameRateCounter
backstageFPSCounter frameRateCounter
lastFrameLog atomic.Int64
metricsOnce sync.Once
metricsEnabled bool
lastKeyframe atomic.Int64
fullNextFrames atomic.Int64
desktopRecoveryAfterBackstageStart atomic.Bool
statFrames atomic.Int64
@@ -576,6 +575,12 @@ var (
backstageLastKeyframe atomic.Int64
)
type frameRateCounter struct {
windowStart atomic.Int64
count atomic.Int64
latest atomic.Int64
}
type prevImage struct {
w int
h int
@@ -604,28 +609,36 @@ func logCodecSupport() {
})
}
func frameFPS(now time.Time) int {
start := fpsWindowStart.Load()
func (c *frameRateCounter) fps(now time.Time) int {
start := c.windowStart.Load()
if start == 0 {
if fpsWindowStart.CompareAndSwap(0, now.UnixNano()) {
fpsCount.Store(1)
return int(fpsLatest.Load())
if c.windowStart.CompareAndSwap(0, now.UnixNano()) {
c.count.Store(1)
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)
if elapsed >= time.Second {
frames := fpsCount.Swap(0)
frames := c.count.Swap(0)
if frames > 0 {
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 {
@@ -693,7 +706,6 @@ func CleanupDesktopStream() {
ResetDesktopCapture()
ResetPrev()
resetHEVCEncoder()
ResetFrameSlots()
}
func ResetPrevbackstage() {
@@ -1903,7 +1915,7 @@ func captureAndSendVirtual(ctx context.Context, env *rt.Env) error {
}
now := time.Now()
fps := frameFPS(now)
fps := backstageFrameFPS(now)
if fps <= 0 {
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 {
now := time.Now()
fps := frameFPS(now)
fps := backstageFrameFPS(now)
if fps <= 0 {
fps = 1
}
@@ -2053,7 +2065,7 @@ func captureAndSendbackstage(ctx context.Context, env *rt.Env) error {
}
now := time.Now()
fps := frameFPS(now)
fps := backstageFrameFPS(now)
if fps <= 0 {
fps = 1
}
@@ -11,3 +11,9 @@ func CaptureDisplayRGBABitBlt(display int) (*image.RGBA, error) {
func CaptureDisplayRGBAPreferBitBlt(display int) (*image.RGBA, error) {
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)
}
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)
}
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) {
quality := desktopJPEGQuality()
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) {
_ = webrtcpub.ConsumeKeyframeRequest(webrtcpub.KindDesktop)
@@ -4,6 +4,7 @@ import (
"os"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
)
@@ -15,6 +16,8 @@ var (
frameAckSeen atomic.Bool
lastAckNano atomic.Int64
maxFrameSlots atomic.Int64
frameStreamsMu sync.Mutex
frameStreams = make(map[string]int)
)
func AcquireFrameSlot() bool {
@@ -54,6 +57,48 @@ func ResetFrameSlots() {
}
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
switch {
case fps >= 180:
@@ -75,7 +120,7 @@ func SetFrameFlowTargetFPS(fps int) {
}
}
}
maxFrameSlots.Store(limit)
return limit
}
func activeFrameSlotLimit() int64 {
@@ -139,11 +139,11 @@ func CleanupVirtualMode() {
virtualDupState.mu.Unlock()
ResetPrevbackstage()
resetH264D3D11TextureEncoder("backstage")
captureMu.Lock()
backstageCaptureMu.Lock()
backstageFreeCapCache()
backstageClearWindowCache()
virtualCapImg = nil
captureMu.Unlock()
backstageCaptureMu.Unlock()
virtualInputMu.Lock()
virtualShiftDown = false
@@ -2393,8 +2393,8 @@ var (
)
func VirtualCaptureDisplayFallback() (*image.RGBA, error) {
captureMu.Lock()
defer captureMu.Unlock()
backstageCaptureMu.Lock()
defer backstageCaptureMu.Unlock()
setDPIAware()
+47 -21
View File
@@ -165,9 +165,9 @@ var (
maxResHeight atomic.Int64 // 0 = default (1080), -1 = native, >0 = specific max height
lastCapScale atomic.Uint64
state capState
thumbnailState capState
captureMu sync.Mutex
stateCreatedAt time.Time
thumbnailMu sync.Mutex
stateRefreshInterval = 30 * time.Second
)
@@ -186,9 +186,18 @@ var captureDisplayFn = func(display int) (*image.RGBA, error) {
}
func captureDisplayBitBlt(display int) (*image.RGBA, error) {
captureMu.Lock()
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()
@@ -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)
dstH := int(float64(srcH) * userScale)
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) {
hdcMem, _, buf, stride, err := state.ensure(hdcScreen, capW, capH)
hdcMem, _, buf, stride, err := captureState.ensure(hdcScreen, capW, capH)
if err != nil {
return nil, 0, 0, 0, err
}
@@ -278,8 +290,8 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
captured = bitBltWithFallback(hdcMem, hdcScreen, bounds, capW, capH)
}
if !captured {
state.reset()
hdcMem, _, buf, stride, err = state.ensure(hdcScreen, capW, capH)
captureState.reset()
hdcMem, _, buf, stride, err = captureState.ensure(hdcScreen, capW, capH)
if err != nil {
return nil, 0, 0, 0, err
}
@@ -328,13 +340,13 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
return img, bitDur, dibDur, convDur, nil
}
if !stateCreatedAt.IsZero() && time.Since(stateCreatedAt) > stateRefreshInterval {
state.reset()
if !captureState.createdAt.IsZero() && time.Since(captureState.createdAt) > stateRefreshInterval {
captureState.reset()
}
img, bitDur, dibDur, convDur, err := captureWithDC()
if err != nil && fromCreateDC {
state.reset()
captureState.reset()
closeDC()
hdcScreen = getDC(0)
fromCreateDC = false
@@ -344,7 +356,7 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) {
img, bitDur, dibDur, convDur, err = captureWithDC()
}
if err != nil {
state.reset()
captureState.reset()
return nil, err
}
@@ -447,14 +459,15 @@ func resizeNearest(src *image.RGBA, w, h int) *image.RGBA {
}
type capState struct {
mu sync.Mutex
hdcMem uintptr
hbmp uintptr
buf []byte
stride int
w int
h int
bits unsafe.Pointer
mu sync.Mutex
hdcMem uintptr
hbmp uintptr
buf []byte
stride int
w int
h int
bits unsafe.Pointer
createdAt time.Time
}
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 {
return 0, 0, nil, 0, syscall.EINVAL
}
stateCreatedAt = time.Now()
s.createdAt = time.Now()
}
if s.w != w || s.h != h || s.hbmp == 0 {
@@ -519,7 +532,7 @@ func (s *capState) reset() {
s.w = 0
s.h = 0
s.bits = nil
stateCreatedAt = time.Time{}
s.createdAt = time.Time{}
}
func captureScale() float64 {
@@ -554,6 +567,19 @@ func effectiveScale(srcW, srcH int) float64 {
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) {
lastCapScale.Store(math.Float64bits(s))
}
@@ -33,13 +33,15 @@ func backstageStart(ctx context.Context, env *rt.Env, autoStartExplorer bool) er
fps := activebackstageTargetFPS()
interval := time.Second / time.Duration(fps)
capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
log.Printf("backstage: starting stream (target fps %d)", fps)
if err := capture.InitializebackstageDesktop(); err != nil {
log.Printf("backstage: failed to initialize hidden desktop: %v", err)
return err
}
capture.StartFrameFlowStream("backstage", fps)
defer capture.StopFrameFlowStream("backstage")
defer capture.CleanupbackstageDesktop()
if autoStartExplorer {
goSafe("backstage auto-start explorer", nil, func() {
@@ -56,21 +58,19 @@ func backstageStart(ctx context.Context, env *rt.Env, autoStartExplorer bool) er
select {
case <-ctx.Done():
log.Printf("backstage: stopping stream")
capture.CleanupbackstageDesktop()
return nil
case <-ticker.C:
fps := activebackstageTargetFPS()
if fps != currentFPS {
currentFPS = fps
capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
capture.UpdateFrameFlowStream("backstage", fps)
ticker.Reset(time.Second / time.Duration(fps))
log.Printf("backstage: target fps changed to %d", fps)
}
if err := capture.Nowbackstage(ctx, env); err != nil {
if ctx.Err() != nil {
log.Printf("backstage: stopping stream")
capture.CleanupbackstageDesktop()
return nil
}
log.Printf("backstage: capture error: %v", err)
@@ -91,7 +91,7 @@ func SetbackstageTargetFPS(fps int) int {
fps = clampDesktopTargetFPS(fps)
backstageTargetFPS.Store(int64(fps))
capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
capture.UpdateFrameFlowStream("backstage", 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":
payload, _ := envelope["payload"].(map[string]interface{})
allDisplays := false
dashboardThumbnail := false
if payload != nil {
mode, _ := payload["mode"].(string)
dashboardThumbnail = mode == "thumbnail"
if v, ok := payload["allDisplays"].(bool); ok && v {
allDisplays = true
} else if mode, ok := payload["mode"].(string); ok && mode == "notification" {
} else if mode == "notification" || mode == "thumbnail" {
allDisplays = true
}
}
if goruntime.GOOS == "windows" {
allDisplays = true
}
return HandleScreenshot(ctx, env, cmdID, allDisplays)
return HandleScreenshot(ctx, env, cmdID, allDisplays, dashboardThumbnail)
case "plugin_load":
payload, _ := envelope["payload"].(map[string]interface{})
if payload == nil {
@@ -1137,6 +1140,22 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
}
capture.RequestDesktopRecoveryAfterBackstageStart()
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()
if env.VirtualCancel != nil {
env.VirtualCancel()
@@ -1165,6 +1184,12 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
sendCommandResultSafe(env, cmdID, true, "")
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()
if env.BackstageCancel != nil {
env.BackstageCancel()
@@ -34,7 +34,8 @@ func DesktopStart(ctx context.Context, env *rt.Env) error {
fps := activeDesktopTargetFPS()
interval := time.Second / time.Duration(fps)
capture.SetDesktopH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
capture.StartFrameFlowStream("desktop", fps)
defer capture.StopFrameFlowStream("desktop")
log.Printf("desktop: starting stream (target fps %d)", fps)
ticker := time.NewTicker(interval)
defer ticker.Stop()
@@ -49,7 +50,7 @@ func DesktopStart(ctx context.Context, env *rt.Env) error {
if fps != currentFPS {
currentFPS = fps
capture.SetDesktopH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
capture.UpdateFrameFlowStream("desktop", fps)
ticker.Reset(time.Second / time.Duration(fps))
log.Printf("desktop: target fps changed to %d", fps)
}
@@ -76,7 +77,7 @@ func SetDesktopTargetFPS(fps int) int {
fps = clampDesktopTargetFPS(fps)
desktopTargetFPS.Store(int64(fps))
capture.SetDesktopH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
capture.UpdateFrameFlowStream("desktop", fps)
return fps
}
@@ -9,6 +9,9 @@ import (
"image/jpeg"
"log"
goruntime "runtime"
"runtime/debug"
"sync"
"time"
"overlord-client/cmd/agent/capture"
rt "overlord-client/cmd/agent/runtime"
@@ -18,12 +21,18 @@ import (
)
var (
monitorCountFn = capture.MonitorCount
displayBoundsFn = capture.DisplayBounds
captureRectFn = screenshot.CaptureRect
monitorCountFn = capture.MonitorCount
displayBoundsFn = capture.DisplayBounds
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
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 {
log.Printf("screenshot: capture failed: %v", err)
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
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 {
log.Printf("screenshot: jpeg encode failed: %v", err)
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) {
if goruntime.GOOS == "windows" {
if img, displayIndex, bounds, err := captureScreenshotImageWindows(allDisplays); err == nil && img != nil {
@@ -136,7 +136,7 @@ func TestHandleScreenshot_RepeatedRequestsDoNotEmitLiveFrames(t *testing.T) {
ctx := context.Background()
writer := &testWriter{}
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)
}
}
@@ -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()
interval := time.Second / time.Duration(fps)
capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
log.Printf("virtual: starting stream (target fps %d)", fps)
if err := capture.InitializeVirtualMode(); err != nil {
log.Printf("virtual: failed to initialize virtual mode: %v", err)
return err
}
capture.StartFrameFlowStream("virtual", fps)
defer capture.StopFrameFlowStream("virtual")
defer capture.CleanupVirtualMode()
ticker := time.NewTicker(interval)
@@ -39,7 +40,7 @@ func VirtualStart(ctx context.Context, env *rt.Env) error {
if fps != currentFPS {
currentFPS = fps
capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
capture.UpdateFrameFlowStream("virtual", fps)
ticker.Reset(time.Second / time.Duration(fps))
log.Printf("virtual: target fps changed to %d", fps)
}
@@ -66,7 +67,7 @@ func SetVirtualTargetFPS(fps int) int {
fps = clampDesktopTargetFPS(fps)
virtualTargetFPS.Store(int64(fps))
capture.SetBackstageH264TargetFPS(fps)
capture.SetFrameFlowTargetFPS(fps)
capture.UpdateFrameFlowStream("virtual", fps)
return fps
}