From ce4ca86ca6be1ce6747cb0d563d4595a9ec9a5df Mon Sep 17 00:00:00 2001
From: Andras Banki-Horvath <bhandras@gmail.com>
Date: Fri, 8 Jan 2021 15:27:47 +0100
Subject: [PATCH] kvdb: fix kvdb.Batch to use state reset when using etcd
 backend

---
 channeldb/kvdb/etcd/db.go   | 12 ------------
 channeldb/kvdb/interface.go | 13 +++++++++++--
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/channeldb/kvdb/etcd/db.go b/channeldb/kvdb/etcd/db.go
index 81835ef0..576591aa 100644
--- a/channeldb/kvdb/etcd/db.go
+++ b/channeldb/kvdb/etcd/db.go
@@ -304,15 +304,3 @@ func (db *db) Copy(w io.Writer) error {
 func (db *db) Close() error {
 	return db.cli.Close()
 }
-
-// Batch opens a database read/write transaction and executes the function f
-// with the transaction passed as a parameter.  After f exits, if f did not
-// error, the transaction is committed.  Otherwise, if f did error, the
-// transaction is rolled back.  If the rollback fails, the original error
-// returned by f is still returned.  If the commit fails, the commit error is
-// returned.
-//
-// Batch is only useful when there are multiple goroutines calling it.
-func (db *db) Batch(apply func(tx walletdb.ReadWriteTx) error) error {
-	return db.Update(apply, func() {})
-}
diff --git a/channeldb/kvdb/interface.go b/channeldb/kvdb/interface.go
index 9ea4ccdf..616e6317 100644
--- a/channeldb/kvdb/interface.go
+++ b/channeldb/kvdb/interface.go
@@ -44,8 +44,17 @@ func View(db Backend, f func(tx RTx) error, reset func()) error {
 // Batch is identical to the Update call, but it attempts to combine several
 // individual Update transactions into a single write database transaction on
 // an optimistic basis. This only has benefits if multiple goroutines call
-// Batch.
-var Batch = walletdb.Batch
+// Batch. For etcd Batch simply does an Update since combination is more complex
+// in that case due to STM retries.
+func Batch(db Backend, f func(tx RwTx) error) error {
+	if extendedDB, ok := db.(ExtendedBackend); ok {
+		// Since Batch calls handle external state reset, we can safely
+		// pass in an empty reset closure.
+		return extendedDB.Update(f, func() {})
+	}
+
+	return walletdb.Batch(db, f)
+}
 
 // Create initializes and opens a database for the specified type. The
 // arguments are specific to the database type driver. See the documentation