From 90dff730ce294a171bfa9f52a4f336554049166b Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Mon, 15 Jul 2024 15:00:10 +0200
Subject: [PATCH] graph: updated builder to use atomic ints

Instead of relying on devs to remember that they must only be accessed
atomically.
---
 graph/builder.go       | 18 +++++++++---------
 graph/builder_test.go  |  7 ++-----
 graph/notifications.go |  3 +--
 3 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/graph/builder.go b/graph/builder.go
index 264a2aacd..2436f3176 100644
--- a/graph/builder.go
+++ b/graph/builder.go
@@ -116,8 +116,8 @@ type Builder struct {
 	started atomic.Bool
 	stopped atomic.Bool
 
-	ntfnClientCounter uint64 // To be used atomically.
-	bestHeight        uint32 // To be used atomically.
+	ntfnClientCounter atomic.Uint64
+	bestHeight        atomic.Uint32
 
 	cfg *Config
 
@@ -278,7 +278,7 @@ func (b *Builder) Start() error {
 		if err != nil {
 			return err
 		}
-		b.bestHeight = uint32(bestHeight)
+		b.bestHeight.Store(uint32(bestHeight))
 
 		// Before we begin normal operation of the router, we first need
 		// to synchronize the channel graph to the latest state of the
@@ -340,7 +340,7 @@ func (b *Builder) syncGraphWithChain() error {
 	if err != nil {
 		return err
 	}
-	b.bestHeight = uint32(bestHeight)
+	b.bestHeight.Store(uint32(bestHeight))
 
 	pruneHash, pruneHeight, err := b.cfg.Graph.PruneTip()
 	if err != nil {
@@ -806,7 +806,7 @@ func (b *Builder) networkHandler() {
 			// Since this block is stale, we update our best height
 			// to the previous block.
 			blockHeight := chainUpdate.Height
-			atomic.StoreUint32(&b.bestHeight, blockHeight-1)
+			b.bestHeight.Store(blockHeight - 1)
 
 			// Update the channel graph to reflect that this block
 			// was disconnected.
@@ -834,7 +834,7 @@ func (b *Builder) networkHandler() {
 			// directly to the end of our main chain. If not, then
 			// we've somehow missed some blocks. Here we'll catch
 			// up the chain with the latest blocks.
-			currentHeight := atomic.LoadUint32(&b.bestHeight)
+			currentHeight := b.bestHeight.Load()
 			switch {
 			case chainUpdate.Height == currentHeight+1:
 				err := b.updateGraphWithClosedChannels(
@@ -991,7 +991,7 @@ func (b *Builder) updateGraphWithClosedChannels(
 	// of the chain tip.
 	blockHeight := chainUpdate.Height
 
-	atomic.StoreUint32(&b.bestHeight, blockHeight)
+	b.bestHeight.Store(blockHeight)
 	log.Infof("Pruning channel graph using block %v (height=%v)",
 		chainUpdate.Hash, blockHeight)
 
@@ -1342,7 +1342,7 @@ func (b *Builder) processUpdate(msg interface{},
 			},
 		}
 		err = b.cfg.ChainView.UpdateFilter(
-			filterUpdate, atomic.LoadUint32(&b.bestHeight),
+			filterUpdate, b.bestHeight.Load(),
 		)
 		if err != nil {
 			return errors.Errorf("unable to update chain "+
@@ -1658,7 +1658,7 @@ func (b *Builder) CurrentBlockHeight() (uint32, error) {
 // is synced to. This can differ from the above chain height if the goroutine
 // responsible for processing the blocks isn't yet up to speed.
 func (b *Builder) SyncedHeight() uint32 {
-	return atomic.LoadUint32(&b.bestHeight)
+	return b.bestHeight.Load()
 }
 
 // GetChannelByID return the channel by the channel id.
diff --git a/graph/builder_test.go b/graph/builder_test.go
index e3a3bd091..600bd8634 100644
--- a/graph/builder_test.go
+++ b/graph/builder_test.go
@@ -11,7 +11,6 @@ import (
 	"net"
 	"os"
 	"strings"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -1392,12 +1391,10 @@ func TestBlockDifferenceFix(t *testing.T) {
 
 	err := wait.NoError(func() error {
 		// Then router height should be updated to the latest block.
-		if atomic.LoadUint32(&ctx.builder.bestHeight) !=
-			newBlockHeight {
-
+		if ctx.builder.bestHeight.Load() != newBlockHeight {
 			return fmt.Errorf("height should have been updated "+
 				"to %v, instead got %v", newBlockHeight,
-				ctx.builder.bestHeight)
+				ctx.builder.bestHeight.Load())
 		}
 
 		return nil
diff --git a/graph/notifications.go b/graph/notifications.go
index 90748b05a..14ea3d127 100644
--- a/graph/notifications.go
+++ b/graph/notifications.go
@@ -5,7 +5,6 @@ import (
 	"image/color"
 	"net"
 	"sync"
-	"sync/atomic"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcutil"
@@ -65,7 +64,7 @@ func (b *Builder) SubscribeTopology() (*TopologyClient, error) {
 
 	// We'll first atomically obtain the next ID for this client from the
 	// incrementing client ID counter.
-	clientID := atomic.AddUint64(&b.ntfnClientCounter, 1)
+	clientID := b.ntfnClientCounter.Add(1)
 
 	log.Debugf("New graph topology client subscription, client %v",
 		clientID)