From dbaba29838271bbde2de2b5ad724b59a360d3e46 Mon Sep 17 00:00:00 2001
From: Oleh Kuznetsov <okuznetsov@google.com>
Date: Wed, 18 Dec 2024 16:30:18 +0000
Subject: [PATCH] Add waitForStartOrCancel helper

---
 gapis/perfetto/android/trace.go | 37 ++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/gapis/perfetto/android/trace.go b/gapis/perfetto/android/trace.go
index 8c8bb9a0c..7979642c5 100644
--- a/gapis/perfetto/android/trace.go
+++ b/gapis/perfetto/android/trace.go
@@ -262,14 +262,8 @@ func (p *Process) Capture(ctx context.Context, start task.Signal, stop task.Sign
 	// Signal that we are ready to start.
 	atomic.StoreInt64(written, 1)
 
-	if p.deferred {
-		select {
-		case <-start:
-			break
-		case <-stop:
-		case <-task.ShouldStop(ctx):
-			return 0, log.Err(ctx, nil, "Cancelled")
-		}
+	if !p.waitForStartOrCancel(ctx, start, stop) {
+		return 0, log.Err(ctx, nil, "Cancelled")
 	}
 
 	if err := p.device.StartPerfettoTrace(ctx, p.config, perfettoTraceFile, stop, ready); err != nil {
@@ -307,6 +301,21 @@ func (w *trackingWriter) Write(p []byte) (int, error) {
 	return n, err
 }
 
+// Waits for start capture signal or cancel signal.
+// Returns true if start signal received or false if need to cancel.
+func (p *Process) waitForStartOrCancel(ctx context.Context, start task.Signal, stop task.Signal) bool {
+	if p.deferred {
+		select {
+		case <-start:
+			return true
+		case <-stop:
+		case <-task.ShouldStop(ctx):
+			return false
+		}
+	}
+	return true
+}
+
 func (p *Process) captureWithClientApi(ctx context.Context, start task.Signal, stop task.Signal, ready task.Task, w io.Writer, written *int64) (int64, error) {
 	defer p.perfettoClient.Close(ctx)
 
@@ -384,15 +393,9 @@ func (p *Process) captureWithClientApi(ctx context.Context, start task.Signal, s
 	delayedReady := task.Delay(ready, 250*time.Millisecond)
 
 	atomic.StoreInt64(written, 1)
-	if p.deferred {
-		select {
-		case <-start:
-			break
-		case <-stop:
-		case <-task.ShouldStop(ctx):
-			ts.Stop(ctx)
-			return 0, errors.New("Cancelled")
-		}
+	if !p.waitForStartOrCancel(ctx, start, stop) {
+		ts.Stop(ctx)
+		return 0, errors.New("Cancelled")
 	}
 	delayedReady(ctx)
 	ts.Start(ctx)