From c29fb81d1b7e46af0a5e8c2b37b8221cf1a91f94 Mon Sep 17 00:00:00 2001
From: Alex Akselrod <aakselrod@lightspark.com>
Date: Fri, 1 Nov 2024 19:50:05 -0700
Subject: [PATCH] batch: handle serialization errors correctly

---
 batch/batch.go      | 17 +++++++++--
 batch/batch_test.go | 74 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 batch/batch_test.go

diff --git a/batch/batch.go b/batch/batch.go
index fcc691582a..9f4842c655 100644
--- a/batch/batch.go
+++ b/batch/batch.go
@@ -5,6 +5,7 @@ import (
 	"sync"
 
 	"github.com/lightningnetwork/lnd/kvdb"
+	"github.com/lightningnetwork/lnd/sqldb"
 )
 
 // errSolo is a sentinel error indicating that the requester should re-run the
@@ -55,8 +56,20 @@ func (b *batch) run() {
 			for i, req := range b.reqs {
 				err := req.Update(tx)
 				if err != nil {
-					failIdx = i
-					return err
+					// If we get a serialization error, we
+					// want the underlying SQL retry
+					// mechanism to retry the entire batch.
+					// Otherwise, we can succeed in an
+					// sqldb retry and still re-execute the
+					// failing request individually.
+					dbErr := sqldb.MapSQLError(err)
+					if !sqldb.IsSerializationError(dbErr) {
+						failIdx = i
+
+						return err
+					}
+
+					return dbErr
 				}
 			}
 			return nil
diff --git a/batch/batch_test.go b/batch/batch_test.go
new file mode 100644
index 0000000000..fef2c55979
--- /dev/null
+++ b/batch/batch_test.go
@@ -0,0 +1,74 @@
+package batch
+
+import (
+	"errors"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/lightningnetwork/lnd/kvdb"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRetry(t *testing.T) {
+	dbDir := t.TempDir()
+
+	dbName := filepath.Join(dbDir, "weks.db")
+	db, err := walletdb.Create(
+		"bdb", dbName, true, kvdb.DefaultDBTimeout,
+	)
+	if err != nil {
+		t.Fatalf("unable to create walletdb: %v", err)
+	}
+	t.Cleanup(func() {
+		db.Close()
+	})
+
+	var (
+		mu     sync.Mutex
+		called int
+	)
+	sched := NewTimeScheduler(db, &mu, time.Second)
+
+	// First, we construct a request that should retry individually and
+	// execute it non-lazily. It should still return the error the second
+	// time.
+	req := &Request{
+		Update: func(tx kvdb.RwTx) error {
+			called++
+
+			return errors.New("test")
+		},
+	}
+	err = sched.Execute(req)
+
+	// Check and reset the called counter.
+	mu.Lock()
+	require.Equal(t, 2, called)
+	called = 0
+	mu.Unlock()
+
+	require.ErrorContains(t, err, "test")
+
+	// Now, we construct a request that should NOT retry because it returns
+	// a serialization error, which should cause the underlying postgres
+	// transaction to retry. Since we aren't using postgres, this will
+	// cause the transaction to not be retried at all.
+	req = &Request{
+		Update: func(tx kvdb.RwTx) error {
+			called++
+
+			return errors.New("could not serialize access")
+		},
+	}
+	err = sched.Execute(req)
+
+	// Check the called counter.
+	mu.Lock()
+	require.Equal(t, 1, called)
+	mu.Unlock()
+
+	require.ErrorContains(t, err, "could not serialize access")
+}