fix: ChannelManager double assignment

Signed-off-by: yangxuan <[email protected]>

Author: XuanYang-cn

internal/datacoord/channel_manager.go

@@ -156,7 +156,7 @@ func (m *ChannelManagerImpl) Startup(ctx context.Context, legacyNodes, allNodes
 	m.mu.Lock()
 	nodeChannels := m.store.GetNodeChannelsBy(
 		WithAllNodes(),
-		func(ch *StateChannel) bool {
+		func(ch *StateChannel) bool { // Channel with drop-mark
 			return m.h.CheckShouldDropChannel(ch.GetName())
 		})
 	m.mu.Unlock()
@@ -238,6 +238,7 @@ func (m *ChannelManagerImpl) Watch(ctx context.Context, ch RWChannel) error {
 			zap.Array("updates", updates), zap.Error(err))
 	}
 
+	// Speed up channel assignment
 	// channel already written into meta, try to assign it to the cluster
 	// not error is returned if failed, the assignment will retry later
 	updates = m.assignPolicy(m.store.GetNodesChannels(), m.store.GetBufferChannelInfo(), m.legacyNodes.Collect())
@@ -286,11 +287,8 @@ func (m *ChannelManagerImpl) DeleteNode(nodeID UniqueID) error {
 	return nil
 }
 
-// reassign reassigns a channel to another DataNode.
+// inner method, lock before using it, reassign reassigns a channel to another DataNode.
 func (m *ChannelManagerImpl) reassign(original *NodeChannelInfo) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
 	updates := m.assignPolicy(m.store.GetNodesChannels(), original, m.legacyNodes.Collect())
 	if updates != nil {
 		return m.execute(updates)
@@ -436,15 +434,16 @@ func (m *ChannelManagerImpl) CheckLoop(ctx context.Context) {
 }
 
 func (m *ChannelManagerImpl) AdvanceChannelState(ctx context.Context) {
-	m.mu.RLock()
+	m.mu.Lock()
 	standbys := m.store.GetNodeChannelsBy(WithAllNodes(), WithChannelStates(Standby))
 	toNotifies := m.store.GetNodeChannelsBy(WithoutBufferNode(), WithChannelStates(ToWatch, ToRelease))
 	toChecks := m.store.GetNodeChannelsBy(WithoutBufferNode(), WithChannelStates(Watching, Releasing))
-	m.mu.RUnlock()
 
-	// Processing standby channels
-	updatedStandbys := false
-	updatedStandbys = m.advanceStandbys(ctx, standbys)
+	// Reassigning standby channels in locks to avoid concurrent assignment with Watch, Remove, AddNode, DeleteNode
+	updatedStandbys := m.advanceStandbys(ctx, standbys)
+	m.mu.Unlock()
+
+	// RPCs stays out of locks
 	updatedToCheckes := m.advanceToChecks(ctx, toChecks)
 	updatedToNotifies := m.advanceToNotifies(ctx, toNotifies)
 
@@ -469,6 +468,7 @@ func (m *ChannelManagerImpl) finishRemoveChannel(nodeID int64, channels ...RWCha
 	}
 }
 
+// inner method need locks
 func (m *ChannelManagerImpl) advanceStandbys(ctx context.Context, standbys []*NodeChannelInfo) bool {
 	var advanced bool = false
 	for _, nodeAssign := range standbys {
@@ -562,21 +562,26 @@ func (m *ChannelManagerImpl) advanceToNotifies(ctx context.Context, toNotifies [
 			got, err := f.Await()
 			res := got.(poolResult)
 
+			action := OnSuccess
 			if err != nil {
 				log.Ctx(ctx).Warn("Failed to notify channel operations to datanode",
 					zap.Int64("assignment", nodeAssign.NodeID),
 					zap.Int("operation count", channelCount),
 					zap.String("channel name", res.ch.GetName()),
 					zap.Error(err),
 				)
+				action = OnFailure
+				if err == merr.ErrChannelReduplicate {
+					action = OnNotifyDuplicate
+				}
 				failedChannels++
 			} else {
 				succeededChannels++
 				advanced = true
 			}
 
 			m.mu.Lock()
-			m.store.UpdateState(err == nil, nodeID, res.ch, res.opID)
+			m.store.UpdateState(action, nodeID, res.ch, res.opID)
 			m.mu.Unlock()
 		}
 
@@ -636,7 +641,11 @@ func (m *ChannelManagerImpl) advanceToChecks(ctx context.Context, toChecks []*No
 			if err == nil {
 				m.mu.Lock()
 				result := got.(poolResult)
-				m.store.UpdateState(result.successful, nodeID, result.ch, result.opID)
+				action := OnSuccess
+				if !result.successful {
+					action = OnFailure
+				}
+				m.store.UpdateState(action, nodeID, result.ch, result.opID)
 				m.mu.Unlock()
 
 				advanced = true

internal/datacoord/channel_store.go

@@ -72,7 +72,7 @@ type RWChannelStore interface {
 	Update(op *ChannelOpSet) error
 
 	// UpdateState is used by StateChannelStore only
-	UpdateState(isSuccessful bool, nodeID int64, channel RWChannel, opID int64)
+	UpdateState(action Action, nodeID int64, channel RWChannel, opID int64)
 	// SegLegacyChannelByNode is used by StateChannelStore only
 	SetLegacyChannelByNode(nodeIDs ...int64)
 
@@ -355,10 +355,18 @@ func (c *StateChannelStore) Reload() error {
 
 		c.AddNode(nodeID)
 
+		channelName := info.GetVchan().GetChannelName()
+		if c.HasChannel(channelName) {
+			log.Warn("channel store detects duplicated channel, skip recovering it",
+				zap.Int64("nodeID", nodeID),
+				zap.String("channel", channelName))
+		}
+
 		channel := NewStateChannelByWatchInfo(nodeID, info)
 		c.channelsInfo[nodeID].AddChannel(channel)
-		log.Info("channel store reload channel",
-			zap.Int64("nodeID", nodeID), zap.String("channel", channel.Name))
+		log.Info("channel store reloads channel from meta",
+			zap.Int64("nodeID", nodeID),
+			zap.String("channel", channelName))
 		metrics.DataCoordDmlChannelNum.WithLabelValues(strconv.FormatInt(nodeID, 10)).Set(float64(len(c.channelsInfo[nodeID].Channels)))
 	}
 	log.Info("channel store reload done", zap.Duration("duration", record.ElapseSpan()))
@@ -375,14 +383,39 @@ func (c *StateChannelStore) AddNode(nodeID int64) {
 	}
 }
 
-func (c *StateChannelStore) UpdateState(isSuccessful bool, nodeID int64, channel RWChannel, opID int64) {
+func (c *StateChannelStore) SetState(targetState ChannelState, nodeID int64, channel RWChannel, opID int64) {
+	channelName := channel.GetName()
+	if cInfo, ok := c.channelsInfo[nodeID]; ok {
+		if stateChannel, ok := cInfo.Channels[channelName]; ok {
+			stateChannel.(*StateChannel).setState(targetState)
+		}
+	}
+}
+
+type Action string
+
+const (
+	OnSuccess         Action = "OnSuccess"
+	OnFailure         Action = "OnFailure"
+	OnNotifyDuplicate Action = "OnNotifyDuplicate" // notify ToWatch to DataNode already subscribed to the channel
+)
+
+func (c *StateChannelStore) UpdateState(action Action, nodeID int64, channel RWChannel, opID int64) {
 	channelName := channel.GetName()
 	if cInfo, ok := c.channelsInfo[nodeID]; ok {
 		if stateChannel, ok := cInfo.Channels[channelName]; ok {
-			if isSuccessful {
+			switch action {
+			case OnSuccess:
 				stateChannel.(*StateChannel).TransitionOnSuccess(opID)
-			} else {
+			case OnFailure:
 				stateChannel.(*StateChannel).TransitionOnFailure(opID)
+			case OnNotifyDuplicate:
+				stateChannel.(*StateChannel).setState(ToRelease)
+			default:
+				log.Warn("unknown action", zap.Any("action", action),
+					zap.Int64("nodeID", nodeID),
+					zap.String("channel", channelName),
+					zap.Int64("opID", opID))
 			}
 		}
 	}

internal/datacoord/channel_store_test.go

@@ -421,12 +421,18 @@ func (s *StateChannelStoreSuite) TestUpdateState() {
 	tests := []struct {
 		description string
 
-		inSuccess       bool
+		inAction        Action
 		inChannelState  ChannelState
 		outChannelState ChannelState
 	}{
-		{"input standby, fail", false, Standby, Standby},
-		{"input standby, success", true, Standby, ToWatch},
+		{"input standby, fail", OnFailure, Standby, Standby},
+		{"input standby, success", OnSuccess, Standby, ToWatch},
+		{"input towatch, duplicate", OnNotifyDuplicate, ToWatch, ToRelease},
+		{"input towatch, fail", OnFailure, ToWatch, Standby},
+		{"input towatch, success", OnSuccess, ToWatch, Watching},
+		{"input torelease, duplicate", OnNotifyDuplicate, ToRelease, ToRelease},
+		{"input torelease, fail", OnFailure, ToRelease, ToRelease},
+		{"input torelease, success", OnSuccess, ToRelease, Releasing},
 	}
 
 	for _, test := range tests {
@@ -443,7 +449,7 @@ func (s *StateChannelStoreSuite) TestUpdateState() {
 				},
 			}
 
-			store.UpdateState(test.inSuccess, bufferID, channel, 0)
+			store.UpdateState(test.inAction, bufferID, channel, 0)
 			s.Equal(test.outChannelState, channel.currentState)
 		})
 	}

internal/datacoord/mock_channel_store.go

@@ -105,6 +105,12 @@ func (m *ChannelManagerImpl) Submit(info *datapb.ChannelWatchInfo) error {
 		return nil
 	}
 
+	// DataNode already watched this channel of other OpID
+	if info.GetState() == datapb.ChannelWatchState_ToWatch &&
+		m.fgManager.HasFlowgraph(channel) {
+		return merr.WrapErrChannelReduplicate(channel)
+	}
+
 	if info.GetState() == datapb.ChannelWatchState_ToRelease &&
 		!m.fgManager.HasFlowgraph(channel) {
 		log.Warn("Release op already finished, skip", zap.Int64("opID", info.GetOpID()), zap.String("channel", channel))