fix(concurrency): use atomic specific types instead of atomic helpers (…

…#682)

block/manager.go

@@ -53,13 +53,14 @@ type Manager struct {
 
 	// Synchronization
 	syncTargetDiode diodes.Diode
-	syncTarget      uint64
-	isSyncedCond    sync.Cond
+
+	syncTarget   atomic.Uint64
+	isSyncedCond sync.Cond
 
 	// Block production
 	shouldProduceBlocksCh chan bool
 	produceEmptyBlockCh   chan bool
-	lastSubmissionTime    int64
+	lastSubmissionTime    atomic.Int64
 	batchInProcess        atomic.Value
 	produceBlockMutex     sync.Mutex
 	applyCachedBlockMutex sync.Mutex
@@ -181,27 +182,27 @@ func (m *Manager) syncBlockManager(ctx context.Context) error {
 			// Since we requested the latest batch and got batch not found it means
 			// the SL still hasn't got any batches for this chain.
 			m.logger.Info("No batches for chain found in SL. Start writing first batch")
-			atomic.StoreUint64(&m.syncTarget, uint64(m.genesis.InitialHeight-1))
+			m.syncTarget.Store(uint64(m.genesis.InitialHeight - 1))
 			return nil
 		}
 		return err
 	}
-	atomic.StoreUint64(&m.syncTarget, resultRetrieveBatch.EndHeight)
+	m.syncTarget.Store(resultRetrieveBatch.EndHeight)
 	err = m.syncUntilTarget(ctx, resultRetrieveBatch.EndHeight)
 	if err != nil {
 		return err
 	}
 
-	m.logger.Info("Synced", "current height", m.store.Height(), "syncTarget", atomic.LoadUint64(&m.syncTarget))
+	m.logger.Info("Synced", "current height", m.store.Height(), "syncTarget", m.syncTarget.Load())
 	return nil
 }
 
 // updateSyncParams updates the sync target and state index if necessary
 func (m *Manager) updateSyncParams(endHeight uint64) {
 	types.RollappHubHeightGauge.Set(float64(endHeight))
 	m.logger.Info("Received new syncTarget", "syncTarget", endHeight)
-	atomic.StoreUint64(&m.syncTarget, endHeight)
-	atomic.StoreInt64(&m.lastSubmissionTime, time.Now().UnixNano())
+	m.syncTarget.Store(endHeight)
+	m.lastSubmissionTime.Store(time.Now().UnixNano())
 }
 
 func getAddress(key crypto.PrivKey) ([]byte, error) {

block/manager_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"crypto/rand"
 	"errors"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -118,7 +117,8 @@ func TestProduceOnlyAfterSynced(t *testing.T) {
 	require.NotNil(t, manager)
 
 	t.Log("Taking the manager out of sync by submitting a batch")
-	syncTarget := atomic.LoadUint64(&manager.syncTarget)
+
+	syncTarget := manager.syncTarget.Load()
 	numBatchesToAdd := 2
 	nextBatchStartHeight := syncTarget + 1
 	var batch *types.Batch
@@ -135,7 +135,7 @@ func TestProduceOnlyAfterSynced(t *testing.T) {
 	}
 
 	// Initially sync target is 0
-	assert.True(t, manager.syncTarget == 0)
+	assert.Zero(t, manager.syncTarget.Load())
 	assert.True(t, manager.store.Height() == 0)
 
 	// enough time to sync and produce blocks
@@ -150,7 +150,7 @@ func TestProduceOnlyAfterSynced(t *testing.T) {
 		assert.NoError(t, err, "Manager start should not produce an error")
 	}()
 	<-ctx.Done()
-	assert.True(t, manager.syncTarget == batch.EndHeight)
+	assert.Equal(t, batch.EndHeight, manager.syncTarget.Load())
 	// validate that we produced blocks
 	assert.Greater(t, manager.store.Height(), batch.EndHeight)
 }
@@ -440,7 +440,7 @@ func TestCreateNextDABatchWithBytesLimit(t *testing.T) {
 			}
 
 			// Call createNextDABatch function
-			startHeight := atomic.LoadUint64(&manager.syncTarget) + 1
+			startHeight := manager.syncTarget.Load() + 1
 			endHeight := startHeight + uint64(tc.blocksToProduce) - 1
 			batch, err := manager.createNextDABatch(startHeight, endHeight)
 			assert.NoError(err)

block/pruning.go

@@ -2,11 +2,10 @@ package block
 
 import (
 	"fmt"
-	"sync/atomic"
 )
 
 func (m *Manager) pruneBlocks(retainHeight int64) (uint64, error) {
-	syncTarget := atomic.LoadUint64(&m.syncTarget)
+	syncTarget := m.syncTarget.Load()
 
 	if retainHeight > int64(syncTarget) {
 		return 0, fmt.Errorf("cannot prune uncommitted blocks")

block/retriever.go

@@ -28,7 +28,7 @@ func (m *Manager) RetriveLoop(ctx context.Context) {
 			}
 			// Check if after we sync we are synced or a new syncTarget was already set.
 			// If we are synced then signal all goroutines waiting on isSyncedCond.
-			if m.store.Height() >= atomic.LoadUint64(&m.syncTarget) {
+			if m.store.Height() >= m.syncTarget.Load() {
 				m.logger.Info("Synced at height", "height", m.store.Height())
 				m.isSyncedCond.L.Lock()
 				m.isSyncedCond.Signal()

block/submit.go

@@ -3,7 +3,6 @@ package block
 import (
 	"context"
 	"fmt"
-	"sync/atomic"
 	"time"
 
 	"github.com/dymensionxyz/dymint/da"
@@ -29,7 +28,7 @@ func (m *Manager) SubmitLoop(ctx context.Context) {
 
 func (m *Manager) handleSubmissionTrigger(ctx context.Context) {
 	// SyncTarget is the height of the last block in the last batch as seen by this node.
-	syncTarget := atomic.LoadUint64(&m.syncTarget)
+	syncTarget := m.syncTarget.Load()
 	height := m.store.Height()
 	// no new blocks produced yet
 	if height <= syncTarget {
@@ -65,7 +64,7 @@ func (m *Manager) handleSubmissionTrigger(ctx context.Context) {
 
 func (m *Manager) submitNextBatch() (uint64, error) {
 	// Get the batch start and end height
-	startHeight := atomic.LoadUint64(&m.syncTarget) + 1
+	startHeight := m.syncTarget.Load() + 1
 	endHeight := uint64(m.store.Height())
 
 	// Create the batch
@@ -112,7 +111,7 @@ func (m *Manager) submitNextBatch() (uint64, error) {
 }
 
 func (m *Manager) validateBatch(batch *types.Batch) error {
-	syncTarget := atomic.LoadUint64(&m.syncTarget)
+	syncTarget := m.syncTarget.Load()
 	if batch.StartHeight != syncTarget+1 {
 		return fmt.Errorf("batch start height != syncTarget + 1. StartHeight %d, m.syncTarget %d", batch.StartHeight, syncTarget)
 	}

block/submit_test.go

@@ -41,17 +41,17 @@ func TestBatchSubmissionHappyFlow(t *testing.T) {
 	initialHeight := uint64(0)
 	require.Zero(manager.store.Height())
 	require.True(manager.batchInProcess.Load() == false)
-	require.Zero(manager.syncTarget)
+	require.Zero(manager.syncTarget.Load())
 
 	// Produce block and validate that we produced blocks
 	err = manager.produceBlock(ctx, true)
 	require.NoError(err)
 	assert.Greater(t, manager.store.Height(), initialHeight)
-	assert.Zero(t, manager.syncTarget)
+	assert.Zero(t, manager.syncTarget.Load())
 
 	// submit and validate sync target
 	manager.handleSubmissionTrigger(ctx)
-	assert.EqualValues(t, 1, manager.syncTarget)
+	assert.EqualValues(t, 1, manager.syncTarget.Load())
 }
 
 func TestBatchSubmissionFailedSubmission(t *testing.T) {
@@ -88,23 +88,23 @@ func TestBatchSubmissionFailedSubmission(t *testing.T) {
 	initialHeight := uint64(0)
 	require.Zero(manager.store.Height())
 	require.True(manager.batchInProcess.Load() == false)
-	require.Zero(manager.syncTarget)
+	require.Zero(manager.syncTarget.Load())
 
 	// Produce block and validate that we produced blocks
 	err = manager.produceBlock(ctx, true)
 	require.NoError(err)
 	assert.Greater(t, manager.store.Height(), initialHeight)
-	assert.Zero(t, manager.syncTarget)
+	assert.Zero(t, manager.syncTarget.Load())
 
 	// try to submit, we expect failure
 	mockLayerI.On("SubmitBatch", mock.Anything, mock.Anything, mock.Anything).Return(fmt.Errorf("submit batch")).Once()
 	manager.handleSubmissionTrigger(ctx)
-	assert.EqualValues(t, 0, manager.syncTarget)
+	assert.EqualValues(t, 0, manager.syncTarget.Load())
 
 	// try to submit again, we expect success
 	mockLayerI.On("SubmitBatch", mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
 	manager.handleSubmissionTrigger(ctx)
-	assert.EqualValues(t, 1, manager.syncTarget)
+	assert.EqualValues(t, 1, manager.syncTarget.Load())
 }
 
 func TestBatchSubmissionAfterTimeout(t *testing.T) {
@@ -142,7 +142,7 @@ func TestBatchSubmissionAfterTimeout(t *testing.T) {
 	require.Equal(initialHeight, manager.store.Height())
 	require.True(manager.batchInProcess.Load() == false)
 
-	require.True(manager.syncTarget == 0)
+	require.Zero(manager.syncTarget.Load())
 
 	var wg sync.WaitGroup
 	mCtx, cancel := context.WithTimeout(context.Background(), runTime)
@@ -162,5 +162,5 @@ func TestBatchSubmissionAfterTimeout(t *testing.T) {
 
 	<-mCtx.Done()
 	wg.Wait() // Wait for all goroutines to finish
-	require.True(manager.syncTarget > 0)
+	require.True(manager.syncTarget.Load() > 0)
 }

da/local/local.go

@@ -18,7 +18,7 @@ import (
 type DataAvailabilityLayerClient struct {
 	logger   types.Logger
 	dalcKV   store.KVStore
-	daHeight uint64
+	daHeight atomic.Uint64
 	config   config
 }
 
@@ -37,7 +37,7 @@ var (
 func (m *DataAvailabilityLayerClient) Init(config []byte, _ *pubsub.Server, dalcKV store.KVStore, logger types.Logger, options ...da.Option) error {
 	m.logger = logger
 	m.dalcKV = dalcKV
-	m.daHeight = 1
+	m.daHeight.Store(1)
 	if len(config) > 0 {
 		var err error
 		m.config.BlockTime, err = time.ParseDuration(string(config))
@@ -77,7 +77,8 @@ func (m *DataAvailabilityLayerClient) GetClientType() da.Client {
 // This should create a transaction which (potentially)
 // triggers a state transition in the DA layer.
 func (m *DataAvailabilityLayerClient) SubmitBatch(batch *types.Batch) da.ResultSubmitBatch {
-	daHeight := atomic.LoadUint64(&m.daHeight)
+	daHeight := m.daHeight.Load()
+
 	m.logger.Debug("Submitting batch to DA layer", "start height", batch.StartHeight, "end height", batch.EndHeight, "da height", daHeight)
 
 	blob, err := batch.MarshalBinary()
@@ -94,7 +95,7 @@ func (m *DataAvailabilityLayerClient) SubmitBatch(batch *types.Batch) da.ResultS
 		return da.ResultSubmitBatch{BaseResult: da.BaseResult{Code: da.StatusError, Message: err.Error(), Error: err}}
 	}
 
-	atomic.StoreUint64(&m.daHeight, daHeight+1)
+	m.daHeight.Store(daHeight + 1) // guaranteed no ABA problem as submit batch is only called when the object is locked
 
 	return da.ResultSubmitBatch{
 		BaseResult: da.BaseResult{
@@ -116,7 +117,7 @@ func (m *DataAvailabilityLayerClient) CheckBatchAvailability(daMetaData *da.DASu
 
 // RetrieveBatches returns block at given height from data availability layer.
 func (m *DataAvailabilityLayerClient) RetrieveBatches(daMetaData *da.DASubmitMetaData) da.ResultRetrieveBatch {
-	if daMetaData.Height >= atomic.LoadUint64(&m.daHeight) {
+	if daMetaData.Height >= m.daHeight.Load() {
 		return da.ResultRetrieveBatch{BaseResult: da.BaseResult{Code: da.StatusError, Message: "batch not found", Error: da.ErrBlobNotFound}}
 	}
 
@@ -160,5 +161,5 @@ func getKey(daHeight uint64, height uint64) []byte {
 
 func (m *DataAvailabilityLayerClient) updateDAHeight() {
 	blockStep := rand.Uint64()%10 + 1 //#nosec
-	atomic.AddUint64(&m.daHeight, blockStep)
+	m.daHeight.Add(blockStep)
 }

settlement/grpc/grpc.go

@@ -58,22 +58,21 @@ func (m *LayerClient) Init(config settlement.Config, pubsub *pubsub.Server, logg
 	return nil
 }
 
-// HubClient implements The HubClient interface
+var _ settlement.HubClient = (*HubGrpcClient)(nil)
+
 type HubGrpcClient struct {
 	ctx            context.Context
 	ProposerPubKey string
 	slStateIndex   uint64
 	logger         types.Logger
 	pubsub         *pubsub.Server
-	latestHeight   uint64
+	latestHeight   atomic.Uint64
 	conn           *grpc.ClientConn
 	sl             slmock.MockSLClient
 	stopchan       chan struct{}
 	refreshTime    int
 }
 
-var _ settlement.HubClient = &HubGrpcClient{}
-
 func newHubClient(config settlement.Config, pubsub *pubsub.Server, logger types.Logger) (*HubGrpcClient, error) {
 	ctx := context.Background()
 
@@ -113,18 +112,19 @@ func newHubClient(config settlement.Config, pubsub *pubsub.Server, logger types.
 	}
 	logger.Debug("Starting grpc SL ", "index", slStateIndex)
 
-	return &HubGrpcClient{
+	ret := &HubGrpcClient{
 		ctx:            ctx,
 		ProposerPubKey: proposer,
 		logger:         logger,
 		pubsub:         pubsub,
-		latestHeight:   latestHeight,
 		slStateIndex:   slStateIndex,
 		conn:           conn,
 		sl:             client,
 		stopchan:       stopchan,
 		refreshTime:    config.SLGrpc.RefreshTime,
-	}, nil
+	}
+	ret.latestHeight.Store(latestHeight)
+	return ret, nil
 }
 
 func initConfig(conf settlement.Config) (proposer string, err error) {
@@ -273,7 +273,7 @@ func (c *HubGrpcClient) saveBatch(batch *settlement.Batch) {
 	}
 	c.logger.Debug("Setting grpc SL Index to ", "index", setIndexReply.GetIndex())
 	// Save latest height in memory and in store
-	atomic.StoreUint64(&c.latestHeight, batch.EndHeight)
+	c.latestHeight.Store(batch.EndHeight)
 }
 
 func (c *HubGrpcClient) convertBatchtoSettlementBatch(batch *types.Batch, daResult *da.ResultSubmitBatch) *settlement.Batch {

store/store.go

@@ -48,7 +48,7 @@ func (s *DefaultStore) NewBatch() Batch {
 
 // SetHeight sets the height saved in the Store if it is higher than the existing height
 func (s *DefaultStore) SetHeight(height uint64) {
-	storeHeight := atomic.LoadUint64(&s.height)
+	storeHeight := s.Height()
 	if height > storeHeight {
 		_ = atomic.CompareAndSwapUint64(&s.height, storeHeight, height)
 	}
@@ -61,7 +61,7 @@ func (s *DefaultStore) Height() uint64 {
 
 // SetBase sets the height saved in the Store of the earliest block
 func (s *DefaultStore) SetBase(height uint64) {
-	baseHeight := atomic.LoadUint64(&s.baseHeight)
+	baseHeight := s.Base()
 	if height > baseHeight {
 		_ = atomic.CompareAndSwapUint64(&s.baseHeight, baseHeight, height)
 	}