Make DataNode release rather than delete when reassign (#17293)

1. Reassgin now will assign to the original Node if no other nodes
   avaliable
2. Make AddNode balance async: ToRealse + Reassign

See also: #16114, #17270
Signed-off-by: Nyangxuan <xuan.yang@zilliz.com>

internal/datacoord/channel_manager.go

@@ -172,12 +172,17 @@ func (c *ChannelManager) Startup(ctx context.Context, nodes []int64) error {
 // ReleaseSuccess  remove
 // ReleaseFail     clean up and remove
 func (c *ChannelManager) checkOldNodes(nodes []UniqueID) error {
+	// Load all the watch infos before processing
+	nodeWatchInfos := make(map[UniqueID][]*datapb.ChannelWatchInfo)
 	for _, nodeID := range nodes {
 		watchInfos, err := c.stateTimer.loadAllChannels(nodeID)
 		if err != nil {
 			return err
 		}
+		nodeWatchInfos[nodeID] = watchInfos
+	}
 
+	for nodeID, watchInfos := range nodeWatchInfos {
 		for _, info := range watchInfos {
 			channelName := info.GetVchan().GetChannelName()
 
@@ -198,12 +203,12 @@ func (c *ChannelManager) checkOldNodes(nodes []UniqueID) error {
 				c.stateTimer.startOne(datapb.ChannelWatchState_ToRelease, channelName, nodeID, info.GetTimeoutTs())
 
 			case datapb.ChannelWatchState_ReleaseSuccess:
-				if err := c.toDelete(nodeID, channelName); err != nil {
+				if err := c.Reassign(nodeID, channelName); err != nil {
 					return err
 				}
 
 			case datapb.ChannelWatchState_ReleaseFailure:
-				if err := c.cleanUpAndDelete(nodeID, channelName); err != nil {
+				if err := c.CleanupAndReassign(nodeID, channelName); err != nil {
 					return err
 				}
 			}
@@ -318,9 +323,9 @@ func (c *ChannelManager) AddNode(nodeID int64) error {
 
 	c.store.Add(nodeID)
 
-	// the default registerPolicy doesn't reassgin channels already there
 	updates := c.registerPolicy(c.store, nodeID)
 	if len(updates) <= 0 {
+		log.Info("register node with no reassignment", zap.Int64("registered node", nodeID))
 		return nil
 	}
 
@@ -328,7 +333,7 @@ func (c *ChannelManager) AddNode(nodeID int64) error {
 		zap.Int64("registered node", nodeID),
 		zap.Array("updates", updates))
 
-	return c.updateWithTimer(updates, datapb.ChannelWatchState_ToWatch)
+	return c.updateWithTimer(updates, datapb.ChannelWatchState_ToRelease)
 }
 
 // DeleteNode deletes the node from the cluster.
@@ -348,10 +353,30 @@ func (c *ChannelManager) DeleteNode(nodeID int64) error {
 	log.Warn("deregister node",
 		zap.Int64("unregistered node", nodeID),
 		zap.Array("updates", updates))
+	if len(updates) <= 0 {
+		return nil
+	}
+
+	var channels []*channel
+	for _, op := range updates {
+		if op.Type == Delete {
+			channels = op.Channels
+		}
+	}
+
+	chNames := make([]string, 0, len(channels))
+	for _, ch := range channels {
+		chNames = append(chNames, ch.Name)
+	}
+	log.Debug("remove timers for channel of the deregistered node",
+		zap.Any("channels", chNames), zap.Int64("nodeID", nodeID))
+	c.stateTimer.removeTimers(chNames)
 
 	if err := c.updateWithTimer(updates, datapb.ChannelWatchState_ToWatch); err != nil {
 		return err
 	}
+
+	// No channels will be return
 	_, err := c.store.Delete(nodeID)
 	return err
 }
@@ -579,14 +604,16 @@ func (c *ChannelManager) processAck(e *ackEvent) {
 		}
 
 	case releaseFailAck, releaseTimeoutAck: // failure acks from toRelease
-		err := c.cleanUpAndDelete(e.nodeID, e.channelName)
+		// Cleanup, Delete and Reassign
+		err := c.CleanupAndReassign(e.nodeID, e.channelName)
 		if err != nil {
-			log.Warn("fail to clean and delete channels for release failure ACKs",
+			log.Warn("fail to clean and reassign channels for release failure ACKs",
 				zap.Int64("nodeID", e.nodeID), zap.String("channel name", e.channelName), zap.Error(err))
 		}
 
 	case releaseSuccessAck:
-		err := c.toDelete(e.nodeID, e.channelName)
+		// Delete and Reassign
+		err := c.Reassign(e.nodeID, e.channelName)
 		if err != nil {
 			log.Warn("fail to response to release success ACK",
 				zap.Int64("nodeID", e.nodeID), zap.String("channel name", e.channelName), zap.Error(err))
@@ -594,8 +621,8 @@ func (c *ChannelManager) processAck(e *ackEvent) {
 	}
 }
 
-// cleanUpAndDelete tries to clean up datanode's subscription, and then delete channel watch info.
-func (c *ChannelManager) cleanUpAndDelete(nodeID UniqueID, channelName string) error {
+// CleanupAndReassign tries to clean up datanode's subscription, and then delete channel watch info.
+func (c *ChannelManager) CleanupAndReassign(nodeID UniqueID, channelName string) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -612,35 +639,31 @@ func (c *ChannelManager) cleanUpAndDelete(nodeID UniqueID, channelName string) e
 		msgstream.UnsubscribeChannels(c.ctx, c.msgstreamFactory, subName, []string{pchannelName})
 	}
 
-	if !c.isMarkedDrop(channelName) {
-		reallocates := &NodeChannelInfo{nodeID, []*channel{chToCleanUp}}
-
-		// reassign policy won't choose the same Node for a ressignment of a channel
-		updates := c.reassignPolicy(c.store, []*NodeChannelInfo{reallocates})
-		if len(updates) <= 0 {
-			log.Warn("fail to reassign channel to other nodes, add channel to buffer", zap.String("channel name", channelName))
-			updates.Add(bufferID, []*channel{chToCleanUp})
-		}
-
-		err := c.remove(nodeID, chToCleanUp)
-		if err != nil {
-			return fmt.Errorf("failed to remove watch info: %v,%s", chToCleanUp, err.Error())
-		}
-
-		log.Info("channel manager reassign channels", zap.Int64("old node ID", nodeID), zap.Array("updates", updates))
-
-		return c.updateWithTimer(updates, datapb.ChannelWatchState_ToWatch)
-	}
-
 	err := c.remove(nodeID, chToCleanUp)
 	if err != nil {
 		return fmt.Errorf("failed to remove watch info: %v,%s", chToCleanUp, err.Error())
 	}
 
-	log.Debug("try to cleanup removal flag ", zap.String("channel name", channelName))
-	c.h.FinishDropChannel(channelName)
+	if c.isMarkedDrop(channelName) {
+		log.Debug("try to cleanup removal flag ", zap.String("channel name", channelName))
+		c.h.FinishDropChannel(channelName)
+		return nil
+	}
 
-	return nil
+	reallocates := &NodeChannelInfo{nodeID, []*channel{chToCleanUp}}
+
+	// reassign policy won't choose the same Node for a ressignment of a channel
+	updates := c.reassignPolicy(c.store, []*NodeChannelInfo{reallocates})
+	if len(updates) <= 0 {
+		log.Warn("fail to reassign channel to other nodes, add channel to the original node",
+			zap.Int64("nodeID", nodeID),
+			zap.String("channel name", channelName))
+		updates.Add(nodeID, []*channel{chToCleanUp})
+	}
+
+	log.Info("channel manager reassign channels", zap.Int64("old nodeID", nodeID), zap.Array("updates", updates))
+
+	return c.updateWithTimer(updates, datapb.ChannelWatchState_ToWatch)
 }
 
 type channelStateChecker func(context.Context)
@@ -724,8 +747,8 @@ func (c *ChannelManager) Release(nodeID UniqueID, channelName string) error {
 	return err
 }
 
-// toDelete removes channel assignment from a datanode
-func (c *ChannelManager) toDelete(nodeID UniqueID, channelName string) error {
+// Reassign removes channel assignment from a datanode
+func (c *ChannelManager) Reassign(nodeID UniqueID, channelName string) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -734,36 +757,33 @@ func (c *ChannelManager) toDelete(nodeID UniqueID, channelName string) error {
 		return fmt.Errorf("fail to find matching nodeID: %d with channelName: %s", nodeID, channelName)
 	}
 
-	if !c.isMarkedDrop(channelName) {
-		reallocates := &NodeChannelInfo{nodeID, []*channel{ch}}
+	if err := c.remove(nodeID, ch); err != nil {
+		return fmt.Errorf("failed to remove watch info: %v,%s", ch, err.Error())
+	}
 
-		// reassign policy won't choose the same Node for a ressignment of a channel
-		updates := c.reassignPolicy(c.store, []*NodeChannelInfo{reallocates})
-		if len(updates) <= 0 {
-			log.Warn("fail to reassign channel to other nodes, add to the buffer", zap.String("channel name", channelName))
-			updates.Add(bufferID, []*channel{ch})
-		}
+	if c.isMarkedDrop(channelName) {
+		log.Debug("try to cleanup removal flag ", zap.String("channel name", channelName))
+		c.h.FinishDropChannel(channelName)
 
-		err := c.remove(nodeID, ch)
-		if err != nil {
-			return fmt.Errorf("failed to remove watch info: %v,%s", ch, err.Error())
-		}
+		log.Info("removed channel assignment", zap.Any("channel", ch))
+		return nil
+	}
 
-		log.Info("channel manager reassign channels", zap.Int64("old node ID", nodeID), zap.Array("updates", updates))
+	reallocates := &NodeChannelInfo{nodeID, []*channel{ch}}
 
-		return c.updateWithTimer(updates, datapb.ChannelWatchState_ToWatch)
+	// reassign policy won't choose the same Node for a ressignment of a channel
+	updates := c.reassignPolicy(c.store, []*NodeChannelInfo{reallocates})
+	if len(updates) <= 0 {
+		log.Warn("fail to reassign channel to other nodes, assign to the original Node",
+			zap.Int64("nodeID", nodeID),
+			zap.String("channel name", channelName))
+		updates.Add(nodeID, []*channel{ch})
 	}
 
-	err := c.remove(nodeID, ch)
-	if err != nil {
-		return err
-	}
+	log.Info("channel manager reassign channels", zap.Int64("old node ID", nodeID), zap.Array("updates", updates))
 
-	log.Debug("try to cleanup removal flag ", zap.String("channel name", channelName))
-	c.h.FinishDropChannel(channelName)
+	return c.updateWithTimer(updates, datapb.ChannelWatchState_ToWatch)
 
-	log.Info("removed channel assignment", zap.Any("channel", ch))
-	return nil
 }
 
 func (c *ChannelManager) getChannelByNodeAndName(nodeID UniqueID, channelName string) *channel {

internal/datacoord/policy.go

@@ -51,6 +51,7 @@ func BufferChannelAssignPolicy(store ROChannelStore, nodeID int64) ChannelOpSet
 }
 
 // AvgAssignRegisterPolicy assigns channels with average to new registered node
+// Register will not directly delete the node-channel pair, channel manager handles the release itself
 func AvgAssignRegisterPolicy(store ROChannelStore, nodeID int64) ChannelOpSet {
 	opSet := BufferChannelAssignPolicy(store, nodeID)
 	if len(opSet) != 0 {
@@ -74,24 +75,19 @@ func AvgAssignRegisterPolicy(store ROChannelStore, nodeID int64) ChannelOpSet {
 		return len(infos[i].Channels) > len(infos[j].Channels)
 	})
 
-	deletes := make(map[int64][]*channel)
-	adds := make(map[int64][]*channel)
+	releases := make(map[int64][]*channel)
 	for i := 0; i < avg; {
 		t := infos[i%len(infos)]
 		idx := i / len(infos)
 		if idx >= len(t.Channels) {
 			continue
 		}
-		deletes[t.NodeID] = append(deletes[t.NodeID], t.Channels[idx])
-		adds[nodeID] = append(adds[nodeID], t.Channels[idx])
+		releases[t.NodeID] = append(releases[t.NodeID], t.Channels[idx])
 		i++
 	}
 
 	opSet = ChannelOpSet{}
-	for k, v := range deletes {
-		opSet.Delete(k, v)
-	}
-	for k, v := range adds {
+	for k, v := range releases {
 		opSet.Add(k, v)
 	}
 	return opSet
@@ -114,8 +110,7 @@ func ConsistentHashRegisterPolicy(hashRing *consistent.Consistent) RegisterPolic
 		elems := formatNodeIDs(store.GetNodes())
 		hashRing.Set(elems)
 
-		removes := make(map[int64][]*channel)
-		adds := make(map[int64][]*channel)
+		releases := make(map[int64][]*channel)
 
 		// If there are buffer channels, then nodeID is the first node.
 		opSet := BufferChannelAssignPolicy(store, nodeID)
@@ -141,16 +136,12 @@ func ConsistentHashRegisterPolicy(hashRing *consistent.Consistent) RegisterPolic
 					return nil
 				}
 				if did != c.NodeID {
-					removes[c.NodeID] = append(removes[c.NodeID], ch)
-					adds[did] = append(adds[did], ch)
+					releases[c.NodeID] = append(releases[c.NodeID], ch)
 				}
 			}
 		}
 
-		for id, channels := range removes {
-			opSet.Delete(id, channels)
-		}
-		for id, channels := range adds {
+		for id, channels := range releases {
 			opSet.Add(id, channels)
 		}
 		return opSet

internal/datacoord/policy_test.go

@@ -97,11 +97,10 @@ func TestConsistentHashRegisterPolicy(t *testing.T) {
 
 		updates := policy(store, 2)
 
-		// chan1 will be hash to 2, chan2 will be hash to 1
 		assert.NotNil(t, updates)
-		assert.Equal(t, 2, len(updates))
-		assert.EqualValues(t, &ChannelOp{Type: Delete, NodeID: 1, Channels: []*channel{channels[0]}}, updates[0])
-		assert.EqualValues(t, &ChannelOp{Type: Add, NodeID: 2, Channels: []*channel{channels[0]}}, updates[1])
+		assert.Equal(t, 1, len(updates))
+		// No Delete operation will be generated
+		assert.EqualValues(t, &ChannelOp{Type: Add, NodeID: 1, Channels: []*channel{channels[0]}}, updates[0])
 	})
 }
 
@@ -486,14 +485,9 @@ func TestAvgAssignRegisterPolicy(t *testing.T) {
 				3,
 			},
 			[]*ChannelOp{
-				{
-					Type:     Delete,
-					NodeID:   1,
-					Channels: []*channel{{"ch1", 1}},
-				},
 				{
 					Type:     Add,
-					NodeID:   3,
+					NodeID:   1,
 					Channels: []*channel{{"ch1", 1}},
 				},
 			},
@@ -525,14 +519,9 @@ func TestAvgAssignRegisterPolicy(t *testing.T) {
 				3,
 			},
 			[]*ChannelOp{
-				{
-					Type:     Delete,
-					NodeID:   1,
-					Channels: []*channel{{"ch1", 1}},
-				},
 				{
 					Type:     Add,
-					NodeID:   3,
+					NodeID:   1,
 					Channels: []*channel{{"ch1", 1}},
 				},
 			},