From 08f145a9ffccb461b4d869d14deaceee6101a95e Mon Sep 17 00:00:00 2001
From: Baha Aiman <bahaaiman@google.com>
Date: Tue, 21 Jan 2025 08:21:44 +0000
Subject: [PATCH] fix(bigtable): Retry RST stream errors

---
 bigtable/bigtable.go   |  3 +++
 bigtable/retry_test.go | 26 ++++++++++++++------------
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/bigtable/bigtable.go b/bigtable/bigtable.go
index 569c7f329870..d6b72a60820e 100644
--- a/bigtable/bigtable.go
+++ b/bigtable/bigtable.go
@@ -182,6 +182,9 @@ var (
 	}
 	retryableInternalErrMsgs = []string{
 		"stream terminated by RST_STREAM", // Retry similar to spanner client. Special case due to https://github.com/googleapis/google-cloud-go/issues/6476
+		"Received Rst stream",
+		"RST_STREAM closed stream",
+		"Received RST_STREAM",
 	}
 )
 
diff --git a/bigtable/retry_test.go b/bigtable/retry_test.go
index e4c12aa44dbf..a67316189db5 100644
--- a/bigtable/retry_test.go
+++ b/bigtable/retry_test.go
@@ -124,18 +124,20 @@ func TestRetryApply(t *testing.T) {
 		t.Errorf("conditionally mutating row with no retries: no error")
 	}
 
-	errCount = 0
-	code = codes.Internal // Will be retried
-	errMsg = "stream terminated by RST_STREAM"
-	if err := tbl.Apply(ctx, "row", mut); err != nil {
-		t.Errorf("applying single mutation with retries: %v", err)
-	}
-	row, err = tbl.ReadRow(ctx, "row")
-	if err != nil {
-		t.Errorf("reading single value with retries: %v", err)
-	}
-	if row == nil {
-		t.Errorf("applying single mutation with retries: could not read back row")
+	for _, msg := range retryableInternalErrMsgs {
+		errCount = 0
+		code = codes.Internal // Will be retried
+		errMsg = msg
+		if err := tbl.Apply(ctx, "row", mut); err != nil {
+			t.Errorf("applying single mutation with retries: %v, errMsg: %v", err, errMsg)
+		}
+		row, err = tbl.ReadRow(ctx, "row")
+		if err != nil {
+			t.Errorf("reading single value with retries: %v, errMsg: %v", err, errMsg)
+		}
+		if row == nil {
+			t.Errorf("applying single mutation with retries: could not read back row. errMsg: %v", errMsg)
+		}
 	}
 
 	errCount = 0