From b1c25febcb2d69f3ec40c8595941145b720bd275 Mon Sep 17 00:00:00 2001
From: Ramy Medhat <abdelaal@google.com>
Date: Thu, 24 Aug 2023 09:41:19 -0400
Subject: [PATCH] Respect dial_timeout flag in rewrapper.

This CL addresses a regression introduced in
[]
that makes rewrapper no longer respect the dial timeout. This likely
explains the elevated rate of build failures due to socket
unavailability in Android starting at 0.111.

Bug: b/297167920
Test: unit tests, manual tests with varying dial timeouts.
Change-Id: If13e000693b66b87f09146d6d9062da26b4bd672
GitOrigin-RevId: c03f4286d021c219b1da1a6f50ac1b3e570df629
---
 cmd/rewrapper/main.go                    |  6 ++----
 internal/pkg/rewrapper/BUILD.bazel       |  2 ++
 internal/pkg/rewrapper/rewrapper.go      | 13 ++++++++-----
 internal/pkg/rewrapper/rewrapper_test.go | 11 ++++++++++-
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/cmd/rewrapper/main.go b/cmd/rewrapper/main.go
index cde0b5ae..41c217ed 100644
--- a/cmd/rewrapper/main.go
+++ b/cmd/rewrapper/main.go
@@ -120,8 +120,7 @@ func main() {
 		log.Fatalf("No exec_strategy provided, must be one of %v", execStrategies)
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), *dialTimeout)
-	defer cancel()
+	ctx := context.Background()
 	conn, err := ipc.DialContext(ctx, serverAddr)
 	if err != nil {
 		log.Fatalf("Fail to dial %s: %v", serverAddr, err)
@@ -129,7 +128,6 @@ func main() {
 	defer conn.Close()
 
 	proxy := pb.NewCommandsClient(conn)
-	ctx = context.Background()
 	wd, err := os.Getwd()
 	if err != nil {
 		log.Fatalf("Failed to get current working directory: %v", err)
@@ -162,7 +160,7 @@ func main() {
 
 	// TODO (b/296409009): Add support for preserve true and download outputs false for downloading stubs.
 
-	resp, err := rewrapper.RunCommand(ctx, proxy, cmd, cOpts)
+	resp, err := rewrapper.RunCommand(ctx, *dialTimeout, proxy, cmd, cOpts)
 	if err != nil {
 		// Don't use log.Fatalf to avoid printing a stack trace.
 		log.Exitf("Command failed: %v", err)
diff --git a/internal/pkg/rewrapper/BUILD.bazel b/internal/pkg/rewrapper/BUILD.bazel
index f5315a1f..03cdb0c2 100644
--- a/internal/pkg/rewrapper/BUILD.bazel
+++ b/internal/pkg/rewrapper/BUILD.bazel
@@ -28,6 +28,8 @@ go_test(
         "@com_github_bazelbuild_remote_apis_sdks//go/pkg/command",
         "@com_github_google_go_cmp//cmp",
         "@org_golang_google_grpc//:go_default_library",
+        "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//status",
         "@org_golang_google_protobuf//testing/protocmp",
     ],
 )
diff --git a/internal/pkg/rewrapper/rewrapper.go b/internal/pkg/rewrapper/rewrapper.go
index f660f037..81eb31da 100644
--- a/internal/pkg/rewrapper/rewrapper.go
+++ b/internal/pkg/rewrapper/rewrapper.go
@@ -18,6 +18,7 @@ package rewrapper
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"strings"
 	"time"
@@ -42,11 +43,9 @@ const (
 )
 
 var (
-	backoff     = retry.ExponentialBackoff(1*time.Second, 15*time.Second, retry.Attempts(10))
+	// backoff has unlimited attempts because retry overall time is controlled by dialTimeout.
+	backoff     = retry.ExponentialBackoff(1*time.Second, 15*time.Second, retry.UnlimitedAttempts)
 	shouldRetry = func(err error) bool {
-		if err == context.DeadlineExceeded {
-			return true
-		}
 		s, ok := status.FromError(err)
 		if !ok {
 			return false
@@ -103,13 +102,17 @@ type CommandOptions struct {
 }
 
 // RunCommand runs a command through the RE proxy.
-func RunCommand(ctx context.Context, proxy Proxy, cmd []string, opts *CommandOptions) (*ppb.RunResponse, error) {
+func RunCommand(ctx context.Context, dialTimeout time.Duration, proxy Proxy, cmd []string, opts *CommandOptions) (*ppb.RunResponse, error) {
 	req, err := createRequest(cmd, opts)
 	if err != nil {
 		return nil, err
 	}
 	var resp *ppb.RunResponse
+	st := time.Now()
 	err = retry.WithPolicy(ctx, shouldRetry, backoff, func() error {
+		if time.Since(st) > dialTimeout {
+			return fmt.Errorf("dial_timeout of %v expired before being able to connect to reproxy", dialTimeout)
+		}
 		resp, err = proxy.RunCommand(ctx, req)
 		return err
 	})
diff --git a/internal/pkg/rewrapper/rewrapper_test.go b/internal/pkg/rewrapper/rewrapper_test.go
index d0a20fd8..702a4f0f 100644
--- a/internal/pkg/rewrapper/rewrapper_test.go
+++ b/internal/pkg/rewrapper/rewrapper_test.go
@@ -24,6 +24,8 @@ import (
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/testing/protocmp"
 
 	ppb "team/foundry-x/re-client/api/proxy"
@@ -115,7 +117,7 @@ func TestRunCommand(t *testing.T) {
 		},
 	}
 	p := &proxyStub{}
-	if _, err := RunCommand(context.Background(), p, cmd, opts); err != nil {
+	if _, err := RunCommand(context.Background(), time.Hour, p, cmd, opts); err != nil {
 		t.Errorf("RunCommand(%v,%v) returned error: %v", cmd, opts, err)
 	}
 	strSliceCmp := protocmp.SortRepeated(func(a, b string) bool { return a < b })
@@ -124,6 +126,13 @@ func TestRunCommand(t *testing.T) {
 	}
 }
 
+func TestRunCommandTimeout(t *testing.T) {
+	p := &proxyStub{err: status.Error(codes.Unavailable, "error")}
+	if _, err := RunCommand(context.Background(), time.Second, p, []string{}, &CommandOptions{}); err == nil {
+		t.Fatalf("RunCommand returned no error, expecting dial timeout error")
+	}
+}
+
 func createRspFile(t *testing.T, contents []string) string {
 	tmpFile, err := os.CreateTemp(os.TempDir(), "")
 	if err != nil {