optimize streaming

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
 	}

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)
+}

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)
+}

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 {

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

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 {

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 {

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"`

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

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);