Remove taskIDBlock data race

Author: stephanos

service/matching/backlog_manager.go

@@ -223,14 +223,14 @@ func (c *backlogManagerImpl) TotalApproximateBacklogCount() int64 {
 }
 
 func (c *backlogManagerImpl) InternalStatus() []*taskqueuespb.InternalTaskQueueStatus {
+	currentTaskIDBlock := c.taskWriter.getTaskIDBlock()
 	return []*taskqueuespb.InternalTaskQueueStatus{
 		&taskqueuespb.InternalTaskQueueStatus{
 			ReadLevel: c.taskAckManager.getReadLevel(),
 			AckLevel:  c.taskAckManager.getAckLevel(),
 			TaskIdBlock: &taskqueuepb.TaskIdBlock{
-				// TODO(pri): this is a data race, it should only be read by taskWriterLoop
-				StartId: c.taskWriter.taskIDBlock.start,
-				EndId:   c.taskWriter.taskIDBlock.end,
+				StartId: currentTaskIDBlock.start,
+				EndId:   currentTaskIDBlock.end,
 			},
 			LoadedTasks:  c.taskAckManager.getBacklogCountHint(),
 			MaxReadLevel: c.db.GetMaxReadLevel(subqueueZero),

service/matching/pri_backlog_manager.go

@@ -323,11 +323,7 @@ func (c *priBacklogManagerImpl) TotalApproximateBacklogCount() int64 {
 }
 
 func (c *priBacklogManagerImpl) InternalStatus() []*taskqueuespb.InternalTaskQueueStatus {
-	// TODO(pri): this is a data race, it should only be read by taskWriterLoop
-	idBlock := &taskqueuepb.TaskIdBlock{
-		StartId: c.taskWriter.taskIDBlock.start,
-		EndId:   c.taskWriter.taskIDBlock.end,
-	}
+	currentTaskIDBlock := c.taskWriter.currentTaskIDBlock()
 
 	c.subqueueLock.Lock()
 	defer c.subqueueLock.Unlock()
@@ -336,9 +332,12 @@ func (c *priBacklogManagerImpl) InternalStatus() []*taskqueuespb.InternalTaskQue
 	for i, r := range c.subqueues {
 		readLevel, ackLevel := r.getLevels()
 		status[i] = &taskqueuespb.InternalTaskQueueStatus{
-			ReadLevel:               readLevel,
-			AckLevel:                ackLevel,
-			TaskIdBlock:             idBlock,
+			ReadLevel: readLevel,
+			AckLevel:  ackLevel,
+			TaskIdBlock: &taskqueuepb.TaskIdBlock{
+				StartId: currentTaskIDBlock.start,
+				EndId:   currentTaskIDBlock.end,
+			},
 			LoadedTasks:             int64(r.getLoadedTasks()),
 			MaxReadLevel:            c.db.GetMaxReadLevel(i),
 			ApproximateBacklogCount: c.db.getApproximateBacklogCount(i),

service/matching/pri_task_writer.go

@@ -26,6 +26,7 @@ package matching
 
 import (
 	"context"
+	"sync/atomic"
 	"time"
 
 	enumspb "go.temporal.io/api/enums/v1"
@@ -52,12 +53,13 @@ type (
 
 	// priTaskWriter writes tasks persistence split among subqueues
 	priTaskWriter struct {
-		backlogMgr  *priBacklogManagerImpl
-		config      *taskQueueConfig
-		db          *taskQueueDB
-		logger      log.Logger
-		appendCh    chan *writeTaskRequest
-		taskIDBlock taskIDBlock
+		backlogMgr      *priBacklogManagerImpl
+		config          *taskQueueConfig
+		db              *taskQueueDB
+		logger          log.Logger
+		appendCh        chan *writeTaskRequest
+		taskIDBlock     taskIDBlock
+		lastTaskIDBlock taskIDBlock // copy of the last taskIDBlock for safe concurrent access via currentTaskIDBlock()
 	}
 )
 
@@ -91,7 +93,6 @@ func (w *priTaskWriter) appendTask(
 	subqueue int,
 	taskInfo *persistencespb.TaskInfo,
 ) error {
-
 	select {
 	case <-w.backlogMgr.tqCtx.Done():
 		return errShutdown
@@ -174,6 +175,7 @@ func (w *priTaskWriter) initState() error {
 		return err
 	}
 	w.taskIDBlock = rangeIDToTaskIDBlock(state.rangeID, w.config.RangeSize)
+	w.lastTaskIDBlock = w.lastTaskIDBlock
 	w.backlogMgr.initState(state, nil)
 	return nil
 }
@@ -185,6 +187,9 @@ func (w *priTaskWriter) taskWriterLoop() {
 
 	var reqs []*writeTaskRequest
 	for {
+		atomic.StoreInt64(&w.lastTaskIDBlock.start, w.taskIDBlock.start)
+		atomic.StoreInt64(&w.lastTaskIDBlock.end, w.taskIDBlock.end)
+
 		select {
 		case request := <-w.appendCh:
 			// read a batch of requests from the channel
@@ -249,3 +254,10 @@ func (w *priTaskWriter) allocTaskIDBlock(prevBlockEnd int64) (taskIDBlock, error
 	}
 	return rangeIDToTaskIDBlock(state.rangeID, w.config.RangeSize), nil
 }
+
+func (w *priTaskWriter) currentTaskIDBlock() taskIDBlock {
+	return taskIDBlock{
+		start: atomic.LoadInt64(&w.lastTaskIDBlock.start),
+		end:   atomic.LoadInt64(&w.lastTaskIDBlock.end),
+	}
+}

service/matching/task_writer.go

@@ -27,6 +27,7 @@ package matching
 import (
 	"context"
 	"errors"
+	"sync/atomic"
 	"time"
 
 	enumspb "go.temporal.io/api/enums/v1"
@@ -54,12 +55,13 @@ type (
 
 	// taskWriter writes tasks sequentially to persistence
 	taskWriter struct {
-		backlogMgr  *backlogManagerImpl
-		config      *taskQueueConfig
-		db          *taskQueueDB
-		logger      log.Logger
-		appendCh    chan *writeTaskRequest
-		taskIDBlock taskIDBlock
+		backlogMgr      *backlogManagerImpl
+		config          *taskQueueConfig
+		db              *taskQueueDB
+		logger          log.Logger
+		appendCh        chan *writeTaskRequest
+		taskIDBlock     taskIDBlock
+		lastTaskIDBlock taskIDBlock // copy of the last taskIDBlock for safe concurrent access via getTaskIDBlock()
 	}
 )
 
@@ -99,6 +101,7 @@ func (w *taskWriter) initReadWriteState() error {
 		return err
 	}
 	w.taskIDBlock = rangeIDToTaskIDBlock(state.rangeID, w.config.RangeSize)
+	w.lastTaskIDBlock = w.taskIDBlock
 	w.backlogMgr.taskAckManager.setAckLevel(state.ackLevel)
 
 	return nil
@@ -182,6 +185,9 @@ func (w *taskWriter) taskWriterLoop() {
 	w.backlogMgr.SetInitializedError(err)
 
 	for {
+		atomic.StoreInt64(&w.lastTaskIDBlock.start, w.taskIDBlock.start)
+		atomic.StoreInt64(&w.lastTaskIDBlock.end, w.taskIDBlock.end)
+
 		select {
 		case request := <-w.appendCh:
 			// read a batch of requests from the channel
@@ -248,3 +254,11 @@ func (w *taskWriter) allocTaskIDBlock(prevBlockEnd int64) (taskIDBlock, error) {
 	}
 	return rangeIDToTaskIDBlock(state.rangeID, w.config.RangeSize), nil
 }
+
+// getTaskIDBlock returns the latest taskIDBlock. Safe to be called concurrently.
+func (w *taskWriter) getTaskIDBlock() taskIDBlock {
+	return taskIDBlock{
+		start: atomic.LoadInt64(&w.lastTaskIDBlock.start),
+		end:   atomic.LoadInt64(&w.lastTaskIDBlock.end),
+	}
+}