From 9502d77b29edca4799ec5fdba2b54a5239bdf775 Mon Sep 17 00:00:00 2001
From: PJ <peterjan.brone@gmail.com>
Date: Fri, 5 Apr 2024 12:10:17 +0200
Subject: [PATCH] stores: fix retryTransaction

---
 stores/sql.go      | 25 ++++++++++++++------
 stores/sql_test.go | 59 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 7 deletions(-)

diff --git a/stores/sql.go b/stores/sql.go
index 34a6d78ab7..320ba1cf6b 100644
--- a/stores/sql.go
+++ b/stores/sql.go
@@ -525,7 +525,7 @@ func (ss *SQLStore) applyUpdates(force bool) error {
 	return nil
 }
 
-func (s *SQLStore) retryTransaction(ctx context.Context, fc func(tx *gorm.DB) error) error {
+func (s *SQLStore) retryTransaction(ctx context.Context, fc func(tx *gorm.DB) error) (err error) {
 	abortRetry := func(err error) bool {
 		if err == nil ||
 			errors.Is(err, context.Canceled) ||
@@ -549,14 +549,25 @@ func (s *SQLStore) retryTransaction(ctx context.Context, fc func(tx *gorm.DB) er
 		}
 		return false
 	}
-	var err error
-	for i := 0; i < len(s.retryTransactionIntervals); i++ {
+
+	attempts := len(s.retryTransactionIntervals) + 1
+	for i := 1; i <= attempts; i++ {
+		// execute the transaction
 		err = s.db.WithContext(ctx).Transaction(fc)
-		if abortRetry(err) {
-			return err
+		if err == nil || abortRetry(err) {
+			return
 		}
-		s.logger.Warn(fmt.Sprintf("transaction attempt %d/%d failed, retry in %v,  err: %v", i+1, len(s.retryTransactionIntervals), s.retryTransactionIntervals[i], err))
-		time.Sleep(s.retryTransactionIntervals[i])
+
+		// if this was the last attempt, return the error
+		if i-1 == len(s.retryTransactionIntervals) {
+			s.logger.Warn(fmt.Sprintf("transaction attempt %d/%d failed, err: %v", i, attempts, err))
+			return
+		}
+
+		// log the failed attempt and sleep before retrying
+		interval := s.retryTransactionIntervals[i-1]
+		s.logger.Warn(fmt.Sprintf("transaction attempt %d/%d failed, retry in %v,  err: %v", i, attempts, interval, err))
+		time.Sleep(interval)
 	}
 	return fmt.Errorf("retryTransaction failed: %w", err)
 }
diff --git a/stores/sql_test.go b/stores/sql_test.go
index 2d29763bba..cae87b1fec 100644
--- a/stores/sql_test.go
+++ b/stores/sql_test.go
@@ -4,13 +4,16 @@ import (
 	"bytes"
 	"context"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	"go.sia.tech/core/types"
 	"go.sia.tech/renterd/alerts"
 	"go.sia.tech/renterd/api"
@@ -18,6 +21,7 @@ import (
 	"go.sia.tech/siad/modules"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
+	"go.uber.org/zap/zaptest/observer"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
 	"lukechampine.com/frand"
@@ -417,3 +421,58 @@ func TestApplyUpdatesErr(t *testing.T) {
 		t.Fatal("lastSave should not have changed")
 	}
 }
+
+func TestRetryTransaction(t *testing.T) {
+	ss := newTestSQLStore(t, defaultTestSQLStoreConfig)
+	defer ss.Close()
+
+	// create custom logger to capture logs
+	observedZapCore, observedLogs := observer.New(zap.InfoLevel)
+	ss.logger = zap.New(observedZapCore).Sugar()
+
+	// collectLogs returns all logs
+	collectLogs := func() (logs []string) {
+		t.Helper()
+		for _, entry := range observedLogs.All() {
+			logs = append(logs, entry.Message)
+		}
+		return
+	}
+
+	// disable retries and retry a transaction that fails
+	ss.retryTransactionIntervals = nil
+	ss.retryTransaction(context.Background(), func(tx *gorm.DB) error { return errors.New("database locked") })
+
+	// assert transaction is attempted once and not retried
+	got := collectLogs()
+	want := []string{"transaction attempt 1/1 failed, err: database locked"}
+	if !reflect.DeepEqual(got, want) {
+		t.Fatal("unexpected logs", cmp.Diff(got, want))
+	}
+
+	// enable retries and retry the same transaction
+	ss.retryTransactionIntervals = []time.Duration{
+		5 * time.Millisecond,
+		10 * time.Millisecond,
+		15 * time.Millisecond,
+	}
+	ss.retryTransaction(context.Background(), func(tx *gorm.DB) error { return errors.New("database locked") })
+
+	// assert transaction is retried 4 times in total
+	got = collectLogs()
+	want = append(want, []string{
+		"transaction attempt 1/4 failed, retry in 5ms,  err: database locked",
+		"transaction attempt 2/4 failed, retry in 10ms,  err: database locked",
+		"transaction attempt 3/4 failed, retry in 15ms,  err: database locked",
+		"transaction attempt 4/4 failed, err: database locked",
+	}...)
+	if !reflect.DeepEqual(got, want) {
+		t.Fatal("unexpected logs", cmp.Diff(got, want))
+	}
+
+	// retry transaction that aborts, assert no logs were added
+	ss.retryTransaction(context.Background(), func(tx *gorm.DB) error { return context.Canceled })
+	if len(observedLogs.All()) != 5 {
+		t.Fatal("expected no logs")
+	}
+}