From b711efe6bde9ae42d1b74a1f8b7f186648418e37 Mon Sep 17 00:00:00 2001 From: KingKDot Date: Wed, 8 Apr 2026 14:14:09 -0500 Subject: [PATCH] optimize streaming --- Overlord-Client/cmd/agent/capture/capture.go | 12 ++++- .../cmd/agent/capture/flowcontrol.go | 49 +++++++++++++++++++ Overlord-Client/cmd/agent/capture/rgbapool.go | 33 +++++++++++++ .../cmd/agent/capture/win_bitblt.go | 4 +- .../cmd/agent/capture/win_duplication.go | 12 ++--- Overlord-Client/cmd/agent/handlers/command.go | 1 + .../cmd/agent/handlers/dispatcher.go | 4 ++ Overlord-Client/cmd/agent/wire/protocol.go | 4 ++ Overlord-Server/src/protocol.ts | 2 + .../routes/websocket-lifecycle-routes.ts | 1 + 10 files changed, 113 insertions(+), 9 deletions(-) create mode 100644 Overlord-Client/cmd/agent/capture/flowcontrol.go create mode 100644 Overlord-Client/cmd/agent/capture/rgbapool.go diff --git a/Overlord-Client/cmd/agent/capture/capture.go b/Overlord-Client/cmd/agent/capture/capture.go index 1403898..36c5b18 100644 --- a/Overlord-Client/cmd/agent/capture/capture.go +++ b/Overlord-Client/cmd/agent/capture/capture.go @@ -130,6 +130,8 @@ func CaptureAndSend(ctx context.Context, env *rt.Env) error { quality := jpegQuality() frame, encodeDur, err := buildFrame(img, display, quality) + PutRGBA(img) + img = nil if err != nil { return err } @@ -142,6 +144,9 @@ func CaptureAndSend(ctx context.Context, env *rt.Env) error { if ctx.Err() != nil { return nil } + if !AcquireFrameSlot() { + return nil // backpressure: server hasn't acked previous frames + } sendStart := time.Now() err = wire.WriteMsg(ctx, env.Conn, frame) sendDur := time.Since(sendStart) @@ -263,14 +268,19 @@ func captureAllDisplaysAndSend(ctx context.Context, env *rt.Env) error { cH = ry } } - canvas := image.NewRGBA(image.Rect(0, 0, cW, cH)) + canvas := GetRGBA(cW, cH) for _, sp := range scaled { dst := image.Rect(sp.offX, sp.offY, sp.offX+sp.img.Rect.Dx(), sp.offY+sp.img.Rect.Dy()) draw.Draw(canvas, dst, sp.img, sp.img.Rect.Min, draw.Src) } + for _, part := range parts { + PutRGBA(part.img) + } quality := jpegQuality() frame, _, err := buildFrame(canvas, 0, quality) + PutRGBA(canvas) + canvas = nil if err != nil { return err } diff --git a/Overlord-Client/cmd/agent/capture/flowcontrol.go b/Overlord-Client/cmd/agent/capture/flowcontrol.go new file mode 100644 index 0000000..4a58737 --- /dev/null +++ b/Overlord-Client/cmd/agent/capture/flowcontrol.go @@ -0,0 +1,49 @@ +package capture + +import ( + "sync/atomic" + "time" +) + +const maxInFlightFrames int64 = 2 + +var ( + inFlightFrames atomic.Int64 + frameAckSeen atomic.Bool + lastAckNano atomic.Int64 +) + +func AcquireFrameSlot() bool { + if !frameAckSeen.Load() { + return true + } + + for { + cur := inFlightFrames.Load() + if cur >= maxInFlightFrames { + lastAck := lastAckNano.Load() + if lastAck > 0 && time.Since(time.Unix(0, lastAck)) > time.Second { + inFlightFrames.Store(0) + continue + } + return false + } + if inFlightFrames.CompareAndSwap(cur, cur+1) { + return true + } + } +} + +func ReleaseFrameSlot() { + frameAckSeen.Store(true) + lastAckNano.Store(time.Now().UnixNano()) + if inFlightFrames.Add(-1) < 0 { + inFlightFrames.Store(0) + } +} + +func ResetFrameSlots() { + inFlightFrames.Store(0) + frameAckSeen.Store(false) + lastAckNano.Store(0) +} diff --git a/Overlord-Client/cmd/agent/capture/rgbapool.go b/Overlord-Client/cmd/agent/capture/rgbapool.go new file mode 100644 index 0000000..9bada99 --- /dev/null +++ b/Overlord-Client/cmd/agent/capture/rgbapool.go @@ -0,0 +1,33 @@ +package capture + +import ( + "image" + "sync" +) + +var rgbaPool sync.Pool + +func GetRGBA(w, h int) *image.RGBA { + need := w * h * 4 + if need <= 0 { + return image.NewRGBA(image.Rect(0, 0, w, h)) + } + if v := rgbaPool.Get(); v != nil { + buf := v.([]byte) + if cap(buf) >= need { + return &image.RGBA{ + Pix: buf[:need], + Stride: w * 4, + Rect: image.Rect(0, 0, w, h), + } + } + } + return image.NewRGBA(image.Rect(0, 0, w, h)) +} + +func PutRGBA(img *image.RGBA) { + if img == nil || len(img.Pix) == 0 { + return + } + rgbaPool.Put(img.Pix) +} diff --git a/Overlord-Client/cmd/agent/capture/win_bitblt.go b/Overlord-Client/cmd/agent/capture/win_bitblt.go index 617c086..4b58c7e 100644 --- a/Overlord-Client/cmd/agent/capture/win_bitblt.go +++ b/Overlord-Client/cmd/agent/capture/win_bitblt.go @@ -320,7 +320,7 @@ func captureDisplayBitBlt(display int) (*image.RGBA, error) { img = resizeNearest(img, dstW, dstH) } else { // Detach from the reusable DIB buffer to avoid concurrent mutation during encode. - stable := image.NewRGBA(img.Rect) + stable := GetRGBA(img.Rect.Dx(), img.Rect.Dy()) copy(stable.Pix, img.Pix) img = stable } @@ -422,7 +422,7 @@ func logCaptureTimings(bitDur, dibDur, convDur time.Duration) { } func resizeNearest(src *image.RGBA, w, h int) *image.RGBA { - dst := image.NewRGBA(image.Rect(0, 0, w, h)) + dst := GetRGBA(w, h) srcW := src.Bounds().Dx() srcH := src.Bounds().Dy() if srcW <= 0 || srcH <= 0 || w <= 0 || h <= 0 { diff --git a/Overlord-Client/cmd/agent/capture/win_duplication.go b/Overlord-Client/cmd/agent/capture/win_duplication.go index 38c7f5e..86c5ec4 100644 --- a/Overlord-Client/cmd/agent/capture/win_duplication.go +++ b/Overlord-Client/cmd/agent/capture/win_duplication.go @@ -828,7 +828,7 @@ func cloneRGBA(src *image.RGBA) *image.RGBA { if src == nil { return nil } - dst := image.NewRGBA(src.Rect) + dst := GetRGBA(src.Rect.Dx(), src.Rect.Dy()) copy(dst.Pix, src.Pix) return dst } @@ -1096,7 +1096,7 @@ func createD3DDevice(adapter *idxgiAdapter1) (*d3d11Device, *d3d11DeviceContext, func convertBGRA(src []byte, pitch, width, height int, rotation uint32) *image.RGBA { switch rotation { case dxgiModeRotationRotate90: - dst := image.NewRGBA(image.Rect(0, 0, height, width)) + dst := GetRGBA(height, width) for y := 0; y < height; y++ { sRow := src[y*pitch : y*pitch+width*4] for x := 0; x < width; x++ { @@ -1112,7 +1112,7 @@ func convertBGRA(src []byte, pitch, width, height int, rotation uint32) *image.R } return dst case dxgiModeRotationRotate180: - dst := image.NewRGBA(image.Rect(0, 0, width, height)) + dst := GetRGBA(width, height) for y := 0; y < height; y++ { sRow := src[y*pitch : y*pitch+width*4] dy := height - 1 - y @@ -1129,7 +1129,7 @@ func convertBGRA(src []byte, pitch, width, height int, rotation uint32) *image.R } return dst case dxgiModeRotationRotate270: - dst := image.NewRGBA(image.Rect(0, 0, height, width)) + dst := GetRGBA(height, width) for y := 0; y < height; y++ { sRow := src[y*pitch : y*pitch+width*4] for x := 0; x < width; x++ { @@ -1145,7 +1145,7 @@ func convertBGRA(src []byte, pitch, width, height int, rotation uint32) *image.R } return dst default: - dst := image.NewRGBA(image.Rect(0, 0, width, height)) + dst := GetRGBA(width, height) for y := 0; y < height; y++ { sRow := src[y*pitch : y*pitch+width*4] dRow := dst.Pix[y*dst.Stride : y*dst.Stride+width*4] @@ -1166,7 +1166,7 @@ func convertBGRAScaled(src []byte, pitch, srcW, srcH, dstW, dstH int, rotation u if rotation != dxgiModeRotationIdentity && rotation != 0 { return nil } - dst := image.NewRGBA(image.Rect(0, 0, dstW, dstH)) + dst := GetRGBA(dstW, dstH) xOff := make([]int, dstW) for x := 0; x < dstW; x++ { xOff[x] = (x * srcW / dstW) * 4 diff --git a/Overlord-Client/cmd/agent/handlers/command.go b/Overlord-Client/cmd/agent/handlers/command.go index 45f3412..6b35713 100644 --- a/Overlord-Client/cmd/agent/handlers/command.go +++ b/Overlord-Client/cmd/agent/handlers/command.go @@ -81,6 +81,7 @@ func resetForReconnect(env *runtime.Env) { } cancelAllCommands() + capture.ResetFrameSlots() env.DesktopMu.Lock() if env.DesktopCancel != nil { diff --git a/Overlord-Client/cmd/agent/handlers/dispatcher.go b/Overlord-Client/cmd/agent/handlers/dispatcher.go index 77289ed..8353405 100644 --- a/Overlord-Client/cmd/agent/handlers/dispatcher.go +++ b/Overlord-Client/cmd/agent/handlers/dispatcher.go @@ -5,6 +5,7 @@ import ( "fmt" "log" + "overlord-client/cmd/agent/capture" "overlord-client/cmd/agent/runtime" ) @@ -37,6 +38,9 @@ func (d *Dispatcher) Dispatch(ctx context.Context, envelope map[string]interface return HandlePing(ctx, d.Env, envelope) case "pong": return HandlePong(ctx, d.Env, envelope) + case "frame_ack": + capture.ReleaseFrameSlot() + return nil case "command": cmdType, _ := envelope["commandType"].(string) switch cmdType { diff --git a/Overlord-Client/cmd/agent/wire/protocol.go b/Overlord-Client/cmd/agent/wire/protocol.go index 2d2cd2c..c2c62bc 100644 --- a/Overlord-Client/cmd/agent/wire/protocol.go +++ b/Overlord-Client/cmd/agent/wire/protocol.go @@ -75,6 +75,10 @@ type Frame struct { Data []byte `msgpack:"data"` } +type FrameAck struct { + Type string `msgpack:"type"` +} + type ScreenshotResult struct { Type string `msgpack:"type"` CommandID string `msgpack:"commandId,omitempty"` diff --git a/Overlord-Server/src/protocol.ts b/Overlord-Server/src/protocol.ts index af43e03..a70e665 100644 --- a/Overlord-Server/src/protocol.ts +++ b/Overlord-Server/src/protocol.ts @@ -138,6 +138,7 @@ export type FrameHeader = { }; export type Frame = { type: "frame"; header: FrameHeader; data: Uint8Array }; +export type FrameAck = { type: "frame_ack" }; export type Status = { type: "status"; state: "idle" | "streaming" | "error"; @@ -304,6 +305,7 @@ export type WireMessage = | CommandResult | ScreenshotResult | Frame + | FrameAck | Status | ConsoleOutput | FileListResult diff --git a/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts b/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts index fca0296..0924372 100644 --- a/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts +++ b/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts @@ -454,6 +454,7 @@ export async function handleWebSocketMessage( } } handleFrame(client, payload); + try { ws.send(encodeMessage({ type: "frame_ack" })); } catch {} break; case "screenshot_result": deps.handleNotificationScreenshotResult(client.id, payload);