From 9820426c89a56aa275b24da0634e7e93bdca705f Mon Sep 17 00:00:00 2001
From: PJ <peterjan.brone@gmail.com>
Date: Fri, 19 Apr 2024 11:23:42 +0200
Subject: [PATCH] stores: improve perf in UpdateStateElements

---
 chain/subscriber.go               |  4 ++--
 internal/test/e2e/cluster.go      |  2 +-
 internal/test/e2e/cluster_test.go |  2 +-
 stores/chain.go                   | 29 ++++++++++++++++++++---------
 4 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/chain/subscriber.go b/chain/subscriber.go
index 2e52e9bef..be6f897a9 100644
--- a/chain/subscriber.go
+++ b/chain/subscriber.go
@@ -162,8 +162,8 @@ func (s *Subscriber) Run() (func(), error) {
 				s.logger.Errorf("failed to get chain index: %v", err)
 				continue
 			}
-			err = s.sync(ci)
-			if err != nil && !errors.Is(err, errClosed) {
+
+			if err := s.sync(ci); err != nil && !errors.Is(err, errClosed) {
 				s.logger.Errorf("failed to sync: %v", err)
 			}
 		}
diff --git a/internal/test/e2e/cluster.go b/internal/test/e2e/cluster.go
index 1a4fc14c1..33c555bb8 100644
--- a/internal/test/e2e/cluster.go
+++ b/internal/test/e2e/cluster.go
@@ -742,9 +742,9 @@ func (c *TestCluster) AddHost(h *Host) {
 	c.tt.Retry(10, time.Second, func() error {
 		c.tt.Helper()
 
-		c.MineBlocks(1)
 		_, err = c.Bus.Host(context.Background(), h.PublicKey())
 		if err != nil {
+			c.MineBlocks(1)
 			return err
 		}
 		return nil
diff --git a/internal/test/e2e/cluster_test.go b/internal/test/e2e/cluster_test.go
index f537a415f..64a595013 100644
--- a/internal/test/e2e/cluster_test.go
+++ b/internal/test/e2e/cluster_test.go
@@ -34,7 +34,7 @@ import (
 // TestNewTestCluster is a test for creating a cluster of Nodes for testing,
 // making sure that it forms contracts, renews contracts and shuts down.
 func TestNewTestCluster(t *testing.T) {
-	cluster := newTestCluster(t, testClusterOptions{logger: newTestLoggerCustom(zapcore.WarnLevel)})
+	cluster := newTestCluster(t, testClusterOptions{logger: newTestLoggerCustom(zapcore.ErrorLevel)})
 	defer cluster.Shutdown()
 	b := cluster.Bus
 	tt := cluster.tt
diff --git a/stores/chain.go b/stores/chain.go
index ea5e3350e..c07885664 100644
--- a/stores/chain.go
+++ b/stores/chain.go
@@ -329,18 +329,29 @@ func (u *chainUpdateTx) UpdateHost(hk types.PublicKey, ha chain.HostAnnouncement
 // UpdateStateElements updates the proofs of all state elements affected by the
 // update.
 func (u *chainUpdateTx) UpdateStateElements(elements []types.StateElement) error {
+	if len(elements) == 0 {
+		return nil
+	}
+
+	var entities []dbWalletOutput
+	if err := u.tx.Model(&dbWalletOutput{}).Find(&entities).Error; err != nil {
+		return err
+	} else if len(entities) == 0 {
+		return nil
+	}
+
+	indices := make(map[types.Hash256]int)
+	for i, e := range entities {
+		indices[types.Hash256(e.OutputID)] = i
+	}
+
 	for _, se := range elements {
-		if err := u.tx.
-			Model(&dbWalletOutput{}).
-			Where("output_id", hash256(se.ID)).
-			Updates(map[string]interface{}{
-				"merkle_proof": merkleProof{proof: se.MerkleProof},
-				"leaf_index":   se.LeafIndex,
-			}).Error; err != nil {
-			return err
+		if index, ok := indices[se.ID]; ok {
+			entities[index].LeafIndex = se.LeafIndex
+			entities[index].MerkleProof = merkleProof{proof: se.MerkleProof}
 		}
 	}
-	return nil
+	return u.tx.Save(&entities).Error
 }
 
 // WalletStateElements implements the ChainStore interface and returns all state