fix: ChannelManager double assignment

This pr fixs double assign in channels, recurr of dropped channel,
amend extra channel meta, and refresh incorrect state with DN by the
following edits:
1. Loose the lock in advanceStandbys to avoid concurrent assignment.
2. Trasfer ToWatch to ToRelease if DN returns ErrChannelRedupulicate.
3. Remove dup channel when recovering

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

Author: XuanYang-cn

internal/datacoord/channel.go

@@ -226,11 +226,9 @@ func (c *StateChannel) TransitionOnFailure(opID int64) {
 		return
 	}
 	switch c.currentState {
-	case Watching:
-		c.setState(Standby)
-	case Releasing:
+	case Watching, Releasing, ToWatch:
 		c.setState(Standby)
-	case Standby, ToWatch, Watched, ToRelease:
+	case Standby, Watched, ToRelease:
 		// Stay original state
 	}
 }

internal/datacoord/channel_manager.go

@@ -156,14 +156,14 @@ 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()
 
 	for _, info := range nodeChannels {
 		m.finishRemoveChannel(info.NodeID, lo.Values(info.Channels)...)
 	}
+	m.mu.Unlock()
 
 	if m.balanceCheckLoop != nil {
 		log.Ctx(ctx).Info("starting channel balance loop")
@@ -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)
 
@@ -453,9 +452,8 @@ func (m *ChannelManagerImpl) AdvanceChannelState(ctx context.Context) {
 	}
 }
 
+// inner method need lock
 func (m *ChannelManagerImpl) finishRemoveChannel(nodeID int64, channels ...RWChannel) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
 	for _, ch := range channels {
 		if err := m.removeChannel(nodeID, ch); err != nil {
 			log.Warn("Failed to remove channel", zap.Any("channel", ch), zap.Error(err))
@@ -469,6 +467,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 +561,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 +640,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
@@ -712,6 +720,7 @@ func (m *ChannelManagerImpl) Check(ctx context.Context, nodeID int64, info *data
 	return false, false
 }
 
+// inner method need lock
 func (m *ChannelManagerImpl) execute(updates *ChannelOpSet) error {
 	for _, op := range updates.ops {
 		if op.Type != Delete {

internal/datacoord/channel_manager_test.go

@@ -378,6 +378,31 @@ func (s *ChannelManagerSuite) TestFindWatcher() {
 func (s *ChannelManagerSuite) TestAdvanceChannelState() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
+	s.Run("advance towatch dn watched to torelease", func() {
+		chNodes := map[string]int64{
+			"ch1": 1,
+			"ch2": 1,
+		}
+		s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
+		s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).
+			RunAndReturn(func(ctx context.Context, nodeID int64, req *datapb.ChannelOperationsRequest) error {
+				s.Require().Equal(1, len(req.GetInfos()))
+				switch req.GetInfos()[0].GetVchan().GetChannelName() {
+				case "ch2":
+					return merr.WrapErrChannelReduplicate("ch2")
+				default:
+					return nil
+				}
+			}).Twice()
+		m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
+		s.Require().NoError(err)
+		s.checkAssignment(m, 1, "ch1", ToWatch)
+		s.checkAssignment(m, 1, "ch2", ToWatch)
+
+		m.AdvanceChannelState(ctx)
+		s.checkAssignment(m, 1, "ch1", Watching)
+		s.checkAssignment(m, 1, "ch2", ToRelease)
+	})
 	s.Run("advance statndby with no available nodes", func() {
 		chNodes := map[string]int64{
 			"ch1": bufferID,

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)
 
@@ -339,6 +339,8 @@ func (c *StateChannelStore) Reload() error {
 	if err != nil {
 		return err
 	}
+
+	dupChannel := []*StateChannel{}
 	for i := 0; i < len(keys); i++ {
 		k := keys[i]
 		v := values[i]
@@ -353,14 +355,36 @@ func (c *StateChannelStore) Reload() error {
 		}
 		reviseVChannelInfo(info.GetVchan())
 
-		c.AddNode(nodeID)
-
+		channelName := info.GetVchan().GetChannelName()
 		channel := NewStateChannelByWatchInfo(nodeID, info)
+
+		if c.HasChannel(channelName) {
+			dupChannel = append(dupChannel, channel)
+			log.Warn("channel store detects duplicated channel, skip recovering it",
+				zap.Int64("nodeID", nodeID),
+				zap.String("channel", channelName))
+			continue
+		}
+
+		c.AddNode(nodeID)
 		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)))
 	}
+
+	for _, channel := range dupChannel {
+		log.Warn("channel store clearing duplicated channel",
+			zap.String("channel", channel.GetName()), zap.Int64("nodeID", channel.assignedNode))
+		chOp := NewChannelOpSet(NewChannelOp(channel.assignedNode, Delete, channel))
+		if err := c.Update(chOp); err != nil {
+			log.Warn("channel store failed to remove duplicated channel, will retry later",
+				zap.String("channel", channel.GetName()),
+				zap.Int64("nodeID", channel.assignedNode),
+				zap.Error(err))
+		}
+	}
 	log.Info("channel store reload done", zap.Duration("duration", record.ElapseSpan()))
 	return nil
 }
@@ -375,14 +399,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,12 +449,45 @@ func (s *StateChannelStoreSuite) TestUpdateState() {
 				},
 			}
 
-			store.UpdateState(test.inSuccess, bufferID, channel, 0)
+			store.UpdateState(test.inAction, bufferID, channel, 0)
 			s.Equal(test.outChannelState, channel.currentState)
 		})
 	}
 }
 
+func (s *StateChannelStoreSuite) TestReloadDupCh() {
+	s.mockTxn.ExpectedCalls = nil
+
+	tests := []struct {
+		channelName string
+		nodeID      int64
+	}{
+		{"ch1", 1},
+		{"ch1", bufferID},
+		{"ch1", 2},
+	}
+
+	var keys, values []string
+	for _, test := range tests {
+		keys = append(keys, fmt.Sprintf("channel_store/%d/%s", test.nodeID, test.channelName))
+		info := generateWatchInfo(test.channelName, datapb.ChannelWatchState_WatchSuccess)
+		bs, err := proto.Marshal(info)
+		s.Require().NoError(err)
+		values = append(values, string(bs))
+	}
+	s.mockTxn.EXPECT().LoadWithPrefix(mock.Anything, mock.AnythingOfType("string")).Return(keys, values, nil)
+	s.mockTxn.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(2)
+
+	store := NewStateChannelStore(s.mockTxn)
+	err := store.Reload()
+	s.Require().NoError(err)
+
+	s.True(store.HasChannel("ch1"))
+	s.ElementsMatch([]int64{1}, store.GetNodes())
+	s.EqualValues(1, store.GetNodeChannelCount(1))
+	s.EqualValues(0, store.GetNodeChannelCount(2))
+}
+
 func (s *StateChannelStoreSuite) TestReload() {
 	type item struct {
 		nodeID      int64