From a2d9f9918019f4d0296eb1b9314c663a5822bbf1 Mon Sep 17 00:00:00 2001
From: Mark Sirek <sirek@cockroachlabs.com>
Date: Sun, 29 Jan 2023 15:35:51 -0800
Subject: [PATCH] logictest: retry support for "statement" and "query" commands

This adds a retry command to logic tests which may be issued on a
separate line preceding either a "statement" or "query" command.
It causes the statement to be retried with exponential backoff up
to some maximum duration, e.g.
```
retry
statement error column "non_exist" does not exist
ALTER TABLE created_as_global SET LOCALITY REGIONAL BY ROW AS "non_exist"
```
This has the same effect as the retry option of the query command, but
now also supports "statement ok", "statement error" and "query error"
commands.

Retry of a query command may be specified by the standalone retry
command, the retry option of the query command, or both.

Fixes: #95668

Epic: CRDB-20535

Release note: None
---
 pkg/sql/logictest/logic.go | 43 +++++++++++++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/pkg/sql/logictest/logic.go b/pkg/sql/logictest/logic.go
index 5bff2cbf264a..f10d2b0a3324 100644
--- a/pkg/sql/logictest/logic.go
+++ b/pkg/sql/logictest/logic.go
@@ -383,6 +383,12 @@ import (
 //    that occur after this command until the end of file or the next subtest
 //    command.
 //
+//  - retry
+//    Specifies that the next occurrence of a statement or query directive
+//    (including those which expect errors) will be retried for a fixed
+//    duration until the test passes, or the alloted time has elapsed.
+//    This is similar to the retry option of the query directive.
+//
 // The overall architecture of TestLogic is as follows:
 //
 // - TestLogic() selects the input files and instantiates
@@ -834,9 +840,6 @@ type logicQuery struct {
 	colTypes string
 	// colNames controls the inclusion of column names in the query result.
 	colNames bool
-	// retry indicates if the query should be retried in case of failure with
-	// exponential backoff up to some maximum duration.
-	retry bool
 	// some tests require the output to match modulo sorting.
 	sorter logicSorter
 	// expectedErr and expectedErrCode are as in logicStatement.
@@ -1016,6 +1019,12 @@ type logicTest struct {
 	// declarativeCorpusCollector used to save declarative schema changer state
 	// to disk.
 	declarativeCorpusCollector *corpus.Collector
+
+	// retry indicates if the statement or query should be retried in case of
+	// failure with exponential backoff up to some maximum duration. It is reset
+	// to false after every successful statement or query test point, including
+	// those which are supposed to error out.
+	retry bool
 }
 
 func (t *logicTest) t() *testing.T {
@@ -2323,6 +2332,8 @@ func (t *logicTest) processSubtest(
 	t.lastProgress = timeutil.Now()
 
 	repeat := 1
+	t.retry = false
+
 	for s.Scan() {
 		t.curPath, t.curLineNo = path, s.Line+subtest.lineLineIndexIntoFile
 		if *maxErrs > 0 && t.failures >= *maxErrs {
@@ -2411,6 +2422,11 @@ func (t *logicTest) processSubtest(
 
 			t.success(path)
 
+		case "retry":
+			// retry is a standalone command that may precede a "statement" or "query"
+			// command. It has the same retry effect as the retry option of the query
+			// command.
+			t.retry = true
 		case "statement":
 			stmt := logicStatement{
 				pos:         fmt.Sprintf("\n%s:%d", path, s.Line+subtest.lineLineIndexIntoFile),
@@ -2441,7 +2457,19 @@ func (t *logicTest) processSubtest(
 			}
 			if !s.Skip {
 				for i := 0; i < repeat; i++ {
-					if cont, err := t.execStatement(stmt); err != nil {
+					var cont bool
+					var err error
+					if t.retry {
+						err = testutils.SucceedsSoonError(func() error {
+							t.purgeZoneConfig()
+							var tempErr error
+							cont, tempErr = t.execStatement(stmt)
+							return tempErr
+						})
+					} else {
+						cont, err = t.execStatement(stmt)
+					}
+					if err != nil {
 						if !cont {
 							return err
 						}
@@ -2582,7 +2610,7 @@ func (t *logicTest) processSubtest(
 							query.colNames = true
 
 						case "retry":
-							query.retry = true
+							t.retry = true
 
 						case "kvtrace":
 							// kvtrace without any arguments doesn't perform any additional
@@ -2765,7 +2793,7 @@ func (t *logicTest) processSubtest(
 				}
 
 				for i := 0; i < repeat; i++ {
-					if query.retry && !*rewriteResultsInTestfiles {
+					if t.retry && !*rewriteResultsInTestfiles {
 						if err := testutils.SucceedsSoonError(func() error {
 							t.purgeZoneConfig()
 							return t.execQuery(query)
@@ -2773,7 +2801,7 @@ func (t *logicTest) processSubtest(
 							t.Error(err)
 						}
 					} else {
-						if query.retry && *rewriteResultsInTestfiles {
+						if t.retry && *rewriteResultsInTestfiles {
 							t.purgeZoneConfig()
 							// The presence of the retry flag indicates that we expect this
 							// query may need some time to succeed. If we are rewriting, wait
@@ -3679,6 +3707,7 @@ func (t *logicTest) formatValues(vals []string, valsPerLine int) []string {
 }
 
 func (t *logicTest) success(file string) {
+	t.retry = false
 	t.progress++
 	now := timeutil.Now()
 	if now.Sub(t.lastProgress) >= 2*time.Second {