From edccd8b5966b05c938cc0277c1cd29f230f369b8 Mon Sep 17 00:00:00 2001
From: Doug Fawley <dfawley@google.com>
Date: Fri, 12 Jun 2020 14:47:00 -0700
Subject: [PATCH] Revert "retry: prevent per-RPC creds error from being
 transparently retried (#3677)"

This reverts commit eb11ffdf9bd6e525ab49264f15b2b8b3841aaad2.
---
 balancer/grpclb/grpclb_test.go | 12 ++++++------
 stream.go                      | 17 ++++++++++++-----
 test/creds_test.go             |  6 +-----
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/balancer/grpclb/grpclb_test.go b/balancer/grpclb/grpclb_test.go
index d55fe3eb8be5..45d2df64567f 100644
--- a/balancer/grpclb/grpclb_test.go
+++ b/balancer/grpclb/grpclb_test.go
@@ -1351,9 +1351,9 @@ func (s) TestGRPCLBStatsUnaryFailedToSend(t *testing.T) {
 			cc.Invoke(context.Background(), failtosendURI, &testpb.Empty{}, nil)
 		}
 	}, &rpcStats{
-		numCallsStarted:                        int64(countRPC),
-		numCallsFinished:                       int64(countRPC),
-		numCallsFinishedWithClientFailedToSend: int64(countRPC) - 1,
+		numCallsStarted:                        int64(countRPC)*2 - 1,
+		numCallsFinished:                       int64(countRPC)*2 - 1,
+		numCallsFinishedWithClientFailedToSend: int64(countRPC-1) * 2,
 		numCallsFinishedKnownReceived:          1,
 	}); err != nil {
 		t.Fatal(err)
@@ -1444,9 +1444,9 @@ func (s) TestGRPCLBStatsStreamingFailedToSend(t *testing.T) {
 			cc.NewStream(context.Background(), &grpc.StreamDesc{}, failtosendURI)
 		}
 	}, &rpcStats{
-		numCallsStarted:                        int64(countRPC),
-		numCallsFinished:                       int64(countRPC),
-		numCallsFinishedWithClientFailedToSend: int64(countRPC) - 1,
+		numCallsStarted:                        int64(countRPC)*2 - 1,
+		numCallsFinished:                       int64(countRPC)*2 - 1,
+		numCallsFinishedWithClientFailedToSend: int64(countRPC-1) * 2,
 		numCallsFinishedKnownReceived:          1,
 	}); err != nil {
 		t.Fatal(err)
diff --git a/stream.go b/stream.go
index 68d5c14506d8..62d51334488d 100644
--- a/stream.go
+++ b/stream.go
@@ -459,6 +459,12 @@ func (cs *clientStream) commitAttempt() {
 // shouldRetry returns nil if the RPC should be retried; otherwise it returns
 // the error that should be returned by the operation.
 func (cs *clientStream) shouldRetry(err error) error {
+	if cs.attempt.s == nil && !cs.callInfo.failFast {
+		// In the event of any error from NewStream (attempt.s == nil), we
+		// never attempted to write anything to the wire, so we can retry
+		// indefinitely for non-fail-fast RPCs.
+		return nil
+	}
 	if cs.finished || cs.committed {
 		// RPC is finished or committed; cannot retry.
 		return err
@@ -466,11 +472,13 @@ func (cs *clientStream) shouldRetry(err error) error {
 	// Wait for the trailers.
 	if cs.attempt.s != nil {
 		<-cs.attempt.s.Done()
-		if cs.firstAttempt && cs.attempt.s.Unprocessed() {
-			// First attempt, stream unprocessed: transparently retry.
-			return nil
-		}
 	}
+	if cs.firstAttempt && (cs.attempt.s == nil || cs.attempt.s.Unprocessed()) {
+		// First attempt, stream unprocessed: transparently retry.
+		cs.firstAttempt = false
+		return nil
+	}
+	cs.firstAttempt = false
 	if cs.cc.dopts.disableRetry {
 		return err
 	}
@@ -556,7 +564,6 @@ func (cs *clientStream) retryLocked(lastErr error) error {
 			cs.commitAttemptLocked()
 			return err
 		}
-		cs.firstAttempt = false
 		if err := cs.newAttemptLocked(nil, nil); err != nil {
 			return err
 		}
diff --git a/test/creds_test.go b/test/creds_test.go
index 8f87af125ec3..7b607ef03716 100644
--- a/test/creds_test.go
+++ b/test/creds_test.go
@@ -189,15 +189,11 @@ func (s) TestGRPCMethodAccessibleToCredsViaContextRequestInfo(t *testing.T) {
 	cc := te.clientConn(grpc.WithPerRPCCredentials(&methodTestCreds{}))
 	tc := testpb.NewTestServiceClient(cc)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()
 	if _, err := tc.EmptyCall(ctx, &testpb.Empty{}); status.Convert(err).Message() != wantMethod {
 		t.Fatalf("ss.client.EmptyCall(_, _) = _, %v; want _, _.Message()=%q", err, wantMethod)
 	}
-
-	if _, err := tc.EmptyCall(ctx, &testpb.Empty{}, grpc.WaitForReady(true)); status.Convert(err).Message() != wantMethod {
-		t.Fatalf("ss.client.EmptyCall(_, _) = _, %v; want _, _.Message()=%q", err, wantMethod)
-	}
 }
 
 const clientAlwaysFailCredErrorMsg = "clientAlwaysFailCred always fails"