Add inrement method for monotonous counter

onflow · Apr 30, 2024 · c6b01e6 · c6b01e6
1 parent 542529f
commit c6b01e6
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 9 deletions.
diff --git a/access/handler.go b/access/handler.go
@@ -1119,9 +1119,7 @@ func (h *Handler) SendAndSubscribeTransactionStatuses(
 	return subscription.HandleSubscription(sub, func(txResults []*TransactionResult) error {
 		for i := range txResults {
 			value := messageIndex.Value()
-			if ok := messageIndex.Set(value + 1); !ok {
-				return status.Errorf(codes.Internal, "the message index has already been incremented to %d", messageIndex.Value())
-			}
+			messageIndex.Increment()
 
 			err = stream.Send(&access.SendAndSubscribeTransactionStatusesResponse{
 				TransactionResults: TransactionResultToMessage(txResults[i]),

diff --git a/engine/access/state_stream/backend/handler.go b/engine/access/state_stream/backend/handler.go
@@ -298,9 +298,7 @@ func (h *Handler) handleEventsResponse(send sendSubscribeEventsResponseFunc, hea
 		}
 
 		index := messageIndex.Value()
-		if ok := messageIndex.Set(index + 1); !ok {
-			return status.Errorf(codes.Internal, "message index already incremented to %d", messageIndex.Value())
-		}
+		messageIndex.Increment()
 
 		err = send(&executiondata.SubscribeEventsResponse{
 			BlockHeight:    resp.Height,
@@ -415,9 +413,7 @@ func (h *Handler) handleAccountStatusesResponse(
 		}
 
 		index := messageIndex.Value()
-		if ok := messageIndex.Set(index + 1); !ok {
-			return status.Errorf(codes.Internal, "message index already incremented to %d", messageIndex.Value())
-		}
+		messageIndex.Increment()
 
 		err = send(&executiondata.SubscribeAccountStatusesResponse{
 			BlockId:      convert.IdentifierToMessage(resp.BlockID),

diff --git a/module/counters/monotonous_counter.go b/module/counters/monotonous_counter.go
@@ -35,3 +35,13 @@ func (c *StrictMonotonousCounter) Set(newValue uint64) bool {
 func (c *StrictMonotonousCounter) Value() uint64 {
 	return atomic.LoadUint64(&c.atomicCounter)
 }
+
+func (c *StrictMonotonousCounter) Increment() uint64 {
+	for {
+		oldValue := c.Value()
+		newValue := oldValue + 1
+		if atomic.CompareAndSwapUint64(&c.atomicCounter, oldValue, newValue) {
+			return newValue
+		}
+	}
+}
diff --git a/module/counters/monotonous_counter_test.go b/module/counters/monotonous_counter_test.go
@@ -16,6 +16,12 @@ func TestSet(t *testing.T) {
 	require.Equal(t, uint64(4), counter.Value())
 }
 
+func TestIncrement(t *testing.T) {
+	counter := NewMonotonousCounter(1)
+	require.Equal(t, uint64(2), counter.Increment())
+	require.Equal(t, uint64(3), counter.Increment())
+}
+
 func TestFuzzy(t *testing.T) {
 	counter := NewMonotonousCounter(3)
 	require.True(t, counter.Set(4))