better cleaning up

This commit is contained in:
KingKDot
2026-08-01 00:04:41 -05:00
parent 6b83e65a6c
commit b68a306708
12 changed files with 266 additions and 27 deletions
+1 -1
View File
@@ -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) # Server source (Overlord-Server/dist-clients may carry pre-built MSVC DLLs from CI)
COPY Overlord-Server/ ./ 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 BackstageInjection/ ./BackstageInjection/
COPY scripts/build-backstage-dll.sh ./scripts/ COPY scripts/build-backstage-dll.sh ./scripts/
COPY BackstageCapture/ ./BackstageCapture/ COPY BackstageCapture/ ./BackstageCapture/
@@ -4,6 +4,38 @@ package capture
import "testing" 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) { func TestBackstagePrintWindowFallbackToggle(t *testing.T) {
SetbackstagePrintWindowFallbackEnabled(true) SetbackstagePrintWindowFallbackEnabled(true)
t.Cleanup(func() { SetbackstagePrintWindowFallbackEnabled(true) }) t.Cleanup(func() { SetbackstagePrintWindowFallbackEnabled(true) })
@@ -187,8 +187,10 @@ var (
backstageLastScale atomic.Uint64 // float64 bits — scale used by last backstage capture backstageLastScale atomic.Uint64 // float64 bits — scale used by last backstage capture
// Capture cache: pooled DC/DIB per window to avoid per-frame allocation // Capture cache: pooled DC/DIB per window to avoid per-frame allocation
backstageWinCache map[uintptr]*backstageWinCacheEntry backstageWinCache map[uintptr]*backstageWinCacheEntry
backstageWinCachePrev []byte backstageWinCachePrev []byte
backstageWinCacheBytes int64
backstageWinCacheSeq uint64
backstageCompHdcMem uintptr backstageCompHdcMem uintptr
backstageCompHbmp uintptr backstageCompHbmp uintptr
@@ -282,10 +284,17 @@ type backstageWinCacheEntry struct {
hbmp uintptr hbmp uintptr
bits unsafe.Pointer bits unsafe.Pointer
w, h int w, h int
bytes int64
usedAt uint64
lastOK bool lastOK bool
age int age int
} }
const (
backstageMaxWindowCacheEntries = 64
backstageMaxWindowCacheBytes = int64(256 << 20)
)
type keybdInput struct { type keybdInput struct {
wVk uint16 wVk uint16
wScan uint16 wScan uint16
@@ -375,16 +384,16 @@ func InitializebackstageDesktop() error {
func CleanupbackstageDesktop() { func CleanupbackstageDesktop() {
backstageDesktopMu.Lock() backstageDesktopMu.Lock()
defer backstageDesktopMu.Unlock() defer backstageDesktopMu.Unlock()
ResetPrevbackstage()
resetH264D3D11TextureEncoder("backstage")
backstageCleanupFrameReaders() backstageCleanupFrameReaders()
captureMu.Lock()
backstageCleanupDWMThumbnails() backstageCleanupDWMThumbnails()
backstageFreeCapCache() backstageFreeCapCache()
for _, entry := range backstageWinCache { backstageClearWindowCache()
backstageFreeCacheEntry(entry)
}
backstageWinCache = nil
backstageWinCachePrev = nil backstageWinCachePrev = nil
if backstageCompHbmp != 0 { if backstageCompHbmp != 0 {
@@ -398,6 +407,7 @@ func CleanupbackstageDesktop() {
backstageCompBits = nil backstageCompBits = nil
backstageCompW = 0 backstageCompW = 0
backstageCompH = 0 backstageCompH = 0
captureMu.Unlock()
backstageInputMu.Lock() backstageInputMu.Lock()
backstageShiftDown = false backstageShiftDown = false
@@ -1705,10 +1715,9 @@ func drawbackstageWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, tar
} }
// Evict cache entries for windows that no longer exist // Evict cache entries for windows that no longer exist
for h, entry := range backstageWinCache { for h := range backstageWinCache {
if !alive[h] { if !alive[h] {
backstageFreeCacheEntry(entry) backstageRemoveWindowCacheEntry(h)
delete(backstageWinCache, h)
} }
} }
@@ -1716,13 +1725,35 @@ func drawbackstageWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, tar
} }
func backstageGetOrCreateCache(hdcScreen uintptr, hwnd uintptr, w, h int) *backstageWinCacheEntry { 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] entry, ok := backstageWinCache[hwnd]
if ok && entry.w == w && entry.h == h && entry.hdcMem != 0 && entry.hbmp != 0 { if ok && entry.w == w && entry.h == h && entry.hdcMem != 0 && entry.hbmp != 0 {
entry.age = 0 entry.age = 0
entry.usedAt = backstageWinCacheSeq
return entry return entry
} }
if ok { 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) hdcMem := createCompatibleDC(hdcScreen)
if hdcMem == 0 { if hdcMem == 0 {
@@ -1752,11 +1783,42 @@ func backstageGetOrCreateCache(hdcScreen uintptr, hwnd uintptr, w, h int) *backs
bits: bits, bits: bits,
w: w, w: w,
h: h, h: h,
bytes: entryBytes,
usedAt: backstageWinCacheSeq,
} }
backstageWinCache[hwnd] = entry backstageWinCache[hwnd] = entry
backstageWinCacheBytes += entryBytes
return entry 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) { func backstageFreeCacheEntry(entry *backstageWinCacheEntry) {
if entry.hbmp != 0 { if entry.hbmp != 0 {
deleteObject(entry.hbmp) deleteObject(entry.hbmp)
+28 -5
View File
@@ -445,18 +445,30 @@ func captureAllDisplaysAndSend(ctx context.Context, env *rt.Env) error {
} }
quality := desktopJPEGQuality() quality := desktopJPEGQuality()
frame, _, err := buildFrame(canvas, 0, quality) frame, err := buildInitialDesktopThumbnail(canvas, quality)
canvasW, canvasH := canvas.Rect.Dx(), canvas.Rect.Dy() canvasW, canvasH := canvas.Rect.Dx(), canvas.Rect.Dy()
PutRGBA(canvas) PutRGBA(canvas)
canvas = nil canvas = nil
if err != nil { if err != nil {
return err return err
} }
frame.Header.FPS = 1
log.Printf("capture: all-displays initial frame %dx%d (%d monitors)", canvasW, canvasH, n) log.Printf("capture: all-displays initial frame %dx%d (%d monitors)", canvasW, canvasH, n)
return wire.WriteMsg(ctx, env.Conn, frame) 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 { func supportsCapture() bool {
return safeActiveDisplays() > 0 return safeActiveDisplays() > 0
} }
@@ -665,14 +677,25 @@ func consumeFullFrame() bool {
} }
func ResetPrev() { func ResetPrev() {
prevMu.Lock() clearDesktopPreviousFrame()
prevFrame = nil
prevMu.Unlock()
lastKeyframe.Store(0) lastKeyframe.Store(0)
requestFullFrames(2) requestFullFrames(2)
resetH264Encoder() resetH264Encoder()
} }
func clearDesktopPreviousFrame() {
prevMu.Lock()
prevFrame = nil
prevMu.Unlock()
}
func CleanupDesktopStream() {
ResetDesktopCapture()
ResetPrev()
resetHEVCEncoder()
ResetFrameSlots()
}
func ResetPrevbackstage() { func ResetPrevbackstage() {
backstagePrevMu.Lock() backstagePrevMu.Lock()
backstagePrevFrame = nil backstagePrevFrame = nil
@@ -23,6 +23,53 @@ type recordingWriter struct {
types []websocket.MessageType 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) { func TestEncodeBlocksBackstageAdvancesOnlyBackstageHistory(t *testing.T) {
img := image.NewRGBA(image.Rect(0, 0, 64, 64)) img := image.NewRGBA(image.Rect(0, 0, 64, 64))
for i := range img.Pix { for i := range img.Pix {
@@ -2,7 +2,9 @@
package capture package capture
func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities { import "context"
func ProbeDesktopEncoderCapabilities(_ context.Context, display int) DesktopEncoderCapabilities {
return completeDesktopEncoderCapabilities(DesktopEncoderCapabilities{ return completeDesktopEncoderCapabilities(DesktopEncoderCapabilities{
Display: display, Display: display,
Profiles: []DesktopEncoderProfile{ Profiles: []DesktopEncoderProfile{
@@ -3,6 +3,7 @@
package capture package capture
import ( import (
"context"
"fmt" "fmt"
"image" "image"
"sort" "sort"
@@ -35,7 +36,7 @@ type desktopProfileProbe struct {
run func(h264D3D11TextureRequest) (time.Duration, error) run func(h264D3D11TextureRequest) (time.Duration, error)
} }
func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities { func ProbeDesktopEncoderCapabilities(ctx context.Context, display int) DesktopEncoderCapabilities {
bounds := DisplayBounds(display) bounds := DisplayBounds(display)
cacheKey := fmt.Sprintf("%d:%d:%d:%d:%d:%.3f", display, bounds.Min.X, bounds.Min.Y, bounds.Dx(), bounds.Dy(), captureScale()) cacheKey := fmt.Sprintf("%d:%d:%d:%d:%d:%.3f", display, bounds.Min.X, bounds.Min.Y, bounds.Dx(), bounds.Dy(), captureScale())
desktopCapabilityCache.Lock() desktopCapabilityCache.Lock()
@@ -43,8 +44,12 @@ func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities {
desktopCapabilityCache.Unlock() desktopCapabilityCache.Unlock()
return cached.caps return cached.caps
} }
caps := probeDesktopEncoderCapabilities(display, bounds) caps := probeDesktopEncoderCapabilities(ctx, display, bounds)
caps = completeDesktopEncoderCapabilities(caps) caps = completeDesktopEncoderCapabilities(caps)
if ctx.Err() != nil {
desktopCapabilityCache.Unlock()
return caps
}
desktopCapabilityCache.entries[cacheKey] = desktopCapabilityCacheEntry{at: time.Now(), caps: caps} desktopCapabilityCache.entries[cacheKey] = desktopCapabilityCacheEntry{at: time.Now(), caps: caps}
for key, entry := range desktopCapabilityCache.entries { for key, entry := range desktopCapabilityCache.entries {
if time.Since(entry.at) >= desktopCapabilityCacheTTL { if time.Since(entry.at) >= desktopCapabilityCacheTTL {
@@ -55,8 +60,13 @@ func ProbeDesktopEncoderCapabilities(display int) DesktopEncoderCapabilities {
return caps 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} 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 { if bounds.Dx() <= 0 || bounds.Dy() <= 0 {
caps.Profiles = safeDesktopProfiles(1280, 720) caps.Profiles = safeDesktopProfiles(1280, 720)
caps.Detail = "Could not determine the selected display size; safe software profiles are shown." 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} fpsOptions := []int{240, 165, 144, 120, 90, 60, 30}
sizes := desktopProbeSizes(nativeWidth, nativeHeight) sizes := desktopProbeSizes(nativeWidth, nativeHeight)
providersByProfile := make(map[string][]string) providersByProfile := make(map[string][]string)
probeLoop:
for _, size := range sizes { for _, size := range sizes {
for _, probe := range probes { for _, probe := range probes {
maxFPS := 0 maxFPS := 0
for _, fps := range fpsOptions { for _, fps := range fpsOptions {
if ctx.Err() != nil {
break probeLoop
}
if fps > desktopProfileFPSCeiling(size.height) { if fps > desktopProfileFPSCeiling(size.height) {
continue continue
} }
@@ -175,11 +189,13 @@ func probeDesktopEncoderCapabilities(display int, bounds image.Rectangle) Deskto
} }
return caps.Profiles[i].FPS > caps.Profiles[j].FPS return caps.Profiles[i].FPS > caps.Profiles[j].FPS
}) })
caps.Probed = true caps.Probed = ctx.Err() == nil
if len(caps.Profiles) == 0 { if len(caps.Profiles) == 0 {
caps.Profiles = safeDesktopProfiles(nativeWidth, nativeHeight) 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." caps.Detail = "No hardware H.264 profile initialized successfully; safe software profiles are shown."
} else { } else {
caps.Detail = "Profiles were tested on the selected display adapter." caps.Detail = "Profiles were tested on the selected display adapter."
@@ -1,6 +1,10 @@
package capture package capture
import "testing" import (
"context"
"image"
"testing"
)
func TestDesktopProfileFPSCeiling(t *testing.T) { func TestDesktopProfileFPSCeiling(t *testing.T) {
tests := []struct { 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)
}
}
@@ -124,6 +124,36 @@ type fakeH264FrameEncoder struct {
closed bool 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) { func (e *fakeH264FrameEncoder) Encode(*image.RGBA) ([]byte, error) {
e.encodeCall++ e.encodeCall++
return e.output, e.err return e.output, e.err
@@ -137,6 +137,13 @@ func CleanupVirtualMode() {
virtualDupState.mu.Lock() virtualDupState.mu.Lock()
virtualDupState.close() virtualDupState.close()
virtualDupState.mu.Unlock() virtualDupState.mu.Unlock()
ResetPrevbackstage()
resetH264D3D11TextureEncoder("backstage")
captureMu.Lock()
backstageFreeCapCache()
backstageClearWindowCache()
virtualCapImg = nil
captureMu.Unlock()
virtualInputMu.Lock() virtualInputMu.Lock()
virtualShiftDown = false virtualShiftDown = false
@@ -2495,10 +2502,9 @@ func drawVirtualWindowsToBuffer(hdcScreen uintptr, bounds image.Rectangle, targe
hwnd = getWindow(hwnd, GW_HWNDPREV) hwnd = getWindow(hwnd, GW_HWNDPREV)
} }
for h, entry := range backstageWinCache { for h := range backstageWinCache {
if !alive[h] { if !alive[h] {
backstageFreeCacheEntry(entry) backstageRemoveWindowCacheEntry(h)
delete(backstageWinCache, h)
} }
} }
@@ -899,7 +899,7 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
display = 0 display = 0
} }
goSafe("desktop encoder capability probe", env.Cancel, func() { 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)) profiles := make([]wire.DesktopEncoderProfile, 0, len(caps.Profiles))
for _, profile := range caps.Profiles { for _, profile := range caps.Profiles {
profiles = append(profiles, wire.DesktopEncoderProfile{ profiles = append(profiles, wire.DesktopEncoderProfile{
@@ -30,6 +30,7 @@ func GetPersistedDisplay() int {
func DesktopStart(ctx context.Context, env *rt.Env) error { func DesktopStart(ctx context.Context, env *rt.Env) error {
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2 //garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
defer capture.CleanupDesktopStream()
fps := activeDesktopTargetFPS() fps := activeDesktopTargetFPS()
interval := time.Second / time.Duration(fps) interval := time.Second / time.Duration(fps)
capture.SetDesktopH264TargetFPS(fps) capture.SetDesktopH264TargetFPS(fps)