Update WatchDmChannels signature (#5447)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/dataservice/binlog_helper.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/golang/protobuf/proto"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
 )
 
 // binlog helper functions persisting binlog paths into kv storage.
@@ -199,3 +200,61 @@ func (s *Server) prepareSegmentPos(segInfo *datapb.SegmentInfo, dmlPos, ddlPos *
 
 	return map[string]string{}, nil
 }
+
+// GetVChanPositions get vchannel latest postitions with provided dml channel names
+func (s *Server) GetVChanPositions(vchans []vchannel) ([]*datapb.VchannelPair, error) {
+	if s.kvClient == nil {
+		return nil, errNilKvClient
+	}
+	pairs := make([]*datapb.VchannelPair, 0, len(vchans))
+
+	for _, vchan := range vchans {
+
+		dmlKey := path.Join(Params.DmlChannelPosSubPath, vchan.DmlChannel)
+		ddlKey := path.Join(Params.DdlChannelPosSubPath, vchan.DmlChannel)
+
+		dmlPos := &datapb.PositionPair{}
+		ddlPos := &datapb.PositionPair{}
+
+		dmlVal, err := s.kvClient.Load(dmlKey)
+		zp := zeroPos(vchan.DmlChannel)
+		if err != nil {
+			dmlPos.StartPosition = &zp
+			dmlPos.EndPosition = &zp
+		} else {
+			err = proto.UnmarshalText(dmlVal, dmlPos)
+			if err != nil {
+				dmlPos.StartPosition = &zp
+				dmlPos.EndPosition = &zp
+			}
+		}
+
+		ddlVal, err := s.kvClient.Load(ddlKey)
+		zp = zeroPos(vchan.DdlChannel)
+		if err != nil {
+			ddlPos.StartPosition = &zp
+			ddlPos.EndPosition = &zp
+		} else {
+			err = proto.UnmarshalText(ddlVal, ddlPos)
+			if err != nil {
+				ddlPos.StartPosition = &zp
+				ddlPos.EndPosition = &zp
+			}
+		}
+
+		pairs = append(pairs, &datapb.VchannelPair{
+			CollectionID:    vchan.CollectionID,
+			DmlVchannelName: vchan.DmlChannel,
+			DdlVchannelName: vchan.DmlChannel,
+			DdlPosition:     ddlPos,
+			DmlPosition:     dmlPos,
+		})
+	}
+	return pairs, nil
+}
+
+func zeroPos(name string) internalpb.MsgPosition {
+	return internalpb.MsgPosition{
+		ChannelName: name,
+	}
+}

internal/dataservice/cluster.go

@@ -25,6 +25,7 @@ type cluster struct {
 	ctx            context.Context
 	dataManager    *clusterNodeManager
 	sessionManager sessionManager
+	posProvider    positionProvider
 
 	startupPolicy    clusterStartupPolicy
 	registerPolicy   dataNodeRegisterPolicy
@@ -76,11 +77,12 @@ func defaultAssignPolicy() channelAssignPolicy {
 	return newAllAssignPolicy()
 }
 
-func newCluster(ctx context.Context, dataManager *clusterNodeManager, sessionManager sessionManager, opts ...clusterOption) *cluster {
+func newCluster(ctx context.Context, dataManager *clusterNodeManager, sessionManager sessionManager, posProvider positionProvider, opts ...clusterOption) *cluster {
 	c := &cluster{
 		ctx:              ctx,
 		sessionManager:   sessionManager,
 		dataManager:      dataManager,
+		posProvider:      posProvider,
 		startupPolicy:    defaultStartupPolicy(),
 		registerPolicy:   defaultRegisterPolicy(),
 		unregisterPolicy: defaultUnregisterPolicy(),
@@ -105,12 +107,21 @@ func (c *cluster) startup(dataNodes []*datapb.DataNodeInfo) error {
 
 func (c *cluster) watch(nodes []*datapb.DataNodeInfo) []*datapb.DataNodeInfo {
 	for _, n := range nodes {
-		uncompletes := make([]string, 0)
+		uncompletes := make([]vchannel, 0, len(nodes))
 		for _, ch := range n.Channels {
 			if ch.State == datapb.ChannelWatchState_Uncomplete {
-				uncompletes = append(uncompletes, ch.Name)
+				uncompletes = append(uncompletes, vchannel{
+					CollectionID: ch.CollectionID,
+					DmlChannel:   ch.Name,
+					DdlChannel:   c.posProvider.GetDdlChannel(),
+				})
 			}
 		}
+		pairs, err := c.posProvider.GetVChanPositions(uncompletes)
+		if err != nil {
+			log.Warn("get vchannel position failed", zap.Error(err))
+			continue
+		}
 		cli, err := c.sessionManager.getOrCreateSession(n.Address)
 		if err != nil {
 			log.Warn("get session failed", zap.String("addr", n.Address), zap.Error(err))
@@ -120,7 +131,7 @@ func (c *cluster) watch(nodes []*datapb.DataNodeInfo) []*datapb.DataNodeInfo {
 			Base: &commonpb.MsgBase{
 				SourceID: Params.NodeID,
 			},
-			//ChannelNames: uncompletes,
+			Vchannels: pairs,
 		}
 		resp, err := cli.WatchDmChannels(c.ctx, req)
 		if err != nil {
@@ -162,11 +173,11 @@ func (c *cluster) unregister(n *datapb.DataNodeInfo) {
 	c.dataManager.updateDataNodes(rets)
 }
 
-func (c *cluster) watchIfNeeded(channel string) {
+func (c *cluster) watchIfNeeded(channel string, collectionID UniqueID) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	cNodes := c.dataManager.getDataNodes(true)
-	rets := c.assginPolicy.apply(cNodes, channel)
+	rets := c.assginPolicy.apply(cNodes, channel, collectionID)
 	c.dataManager.updateDataNodes(rets)
 	rets = c.watch(rets)
 	c.dataManager.updateDataNodes(rets)

internal/dataservice/cluster_test.go

@@ -101,11 +101,11 @@ func TestWatchIfNeeded(t *testing.T) {
 	assert.EqualValues(t, "localhost:8080", dataNodes[addr].Address)
 
 	chName := "ch1"
-	cluster.watchIfNeeded(chName)
+	cluster.watchIfNeeded(chName, 0)
 	dataNodes = cluster.dataManager.getDataNodes(true)
 	assert.EqualValues(t, 1, len(dataNodes[addr].Channels))
 	assert.EqualValues(t, chName, dataNodes[addr].Channels[0].Name)
-	cluster.watchIfNeeded(chName)
+	cluster.watchIfNeeded(chName, 0)
 	assert.EqualValues(t, 1, len(dataNodes[addr].Channels))
 	assert.EqualValues(t, chName, dataNodes[addr].Channels[0].Name)
 }
@@ -139,5 +139,5 @@ func createCluster(t *testing.T, ch chan interface{}, options ...clusterOption)
 	sessionManager := newMockSessionManager(ch)
 	dataManager, err := newClusterNodeManager(kv)
 	assert.Nil(t, err)
-	return newCluster(context.TODO(), dataManager, sessionManager, options...)
+	return newCluster(context.TODO(), dataManager, sessionManager, dummyPosProvider{}, options...)
 }

internal/dataservice/datanode_helper.go

+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+package dataservice
+
+import (
+	"github.com/milvus-io/milvus/internal/log"
+	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"go.uber.org/zap"
+)
+
+type vchannel struct {
+	CollectionID UniqueID
+	DmlChannel   string
+	DdlChannel   string
+}
+
+// positionProvider provides vchannel pair related position pairs
+type positionProvider interface {
+	GetVChanPositions(vchans []vchannel) ([]*datapb.VchannelPair, error)
+	GetDdlChannel() string
+}
+
+type dummyPosProvider struct{}
+
+//GetVChanPositions implements positionProvider
+func (dp dummyPosProvider) GetVChanPositions(vchans []vchannel) ([]*datapb.VchannelPair, error) {
+	pairs := make([]*datapb.VchannelPair, len(vchans))
+	for _, vchan := range vchans {
+		dmlPos := &datapb.PositionPair{}
+		ddlPos := &datapb.PositionPair{}
+		pairs = append(pairs, &datapb.VchannelPair{
+			CollectionID:    vchan.CollectionID,
+			DmlVchannelName: vchan.DmlChannel,
+			DdlVchannelName: vchan.DmlChannel,
+			DdlPosition:     ddlPos,
+			DmlPosition:     dmlPos,
+		})
+	}
+	return pairs, nil
+}
+
+//GetDdlChannel implements positionProvider
+func (dp dummyPosProvider) GetDdlChannel() string {
+	return "dummy_ddl"
+}
+
+//GetDdlChannel implements positionProvider
+func (s *Server) GetDdlChannel() string {
+	s.getDDChannel()
+	return s.ddChannelMu.name
+}
+
+// getAllActiveVChannels get all vchannels with unflushed segments
+func (s *Server) getAllActiveVChannels() []vchannel {
+	s.getDDChannel()
+	segments := s.meta.GetUnFlushedSegments()
+
+	mChanCol := make(map[string]UniqueID)
+	for _, segment := range segments {
+		ocid, has := mChanCol[segment.InsertChannel]
+		if has && ocid != segment.CollectionID {
+			log.Error("col:vchan not 1:N",
+				zap.Int64("colid 1", ocid),
+				zap.Int64("colid 2", segment.CollectionID),
+				zap.String("channel", segment.InsertChannel))
+		}
+		mChanCol[segment.InsertChannel] = segment.CollectionID
+	}
+
+	vchans := make([]vchannel, 0, len(mChanCol))
+	for dmChan, colID := range mChanCol {
+		vchans = append(vchans, vchannel{
+			CollectionID: colID,
+			DmlChannel:   dmChan,
+			DdlChannel:   s.ddChannelMu.name,
+		})
+	}
+	return vchans
+}

internal/dataservice/param.go

@@ -106,6 +106,8 @@ func (p *ParamTable) Init() {
 		p.initStatsStreamPosSubPath()
 		p.initSegmentDmlPosSubPath()
 		p.initSegmentDdlPosSubPath()
+		p.initDmlChannelPosSubPath()
+		p.initDdlChannelPosSubPath()
 	})
 }
 

internal/dataservice/policy.go

@@ -11,9 +11,13 @@
 package dataservice
 
 import (
+	"crypto/rand"
 	"fmt"
+	"math/big"
 
+	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"go.uber.org/zap"
 )
 
 type clusterDeltaChange struct {
@@ -74,8 +78,43 @@ func (p *doNothingUnregisterPolicy) apply(cluster map[string]*datapb.DataNodeInf
 	return nil
 }
 
+type reassignRandomUnregisterPolicy struct{}
+
+func (p *reassignRandomUnregisterPolicy) apply(cluster map[string]*datapb.DataNodeInfo, session *datapb.DataNodeInfo) []*datapb.DataNodeInfo {
+	if len(cluster) == 0 || // no available node
+		len(session.Channels) == 0 { // lost node not watching any channels
+		return []*datapb.DataNodeInfo{}
+	}
+
+	mChan := make(map[string]struct{}, len(session.Channels))
+	for _, chanSt := range session.Channels {
+		mChan[chanSt.Name] = struct{}{}
+	}
+
+	bIdx, err := rand.Int(rand.Reader, big.NewInt(int64(len(cluster))))
+	if err != nil {
+		log.Error("error generated rand idx", zap.Error(err))
+		return []*datapb.DataNodeInfo{}
+	}
+	idx := bIdx.Int64()
+	if int(idx) >= len(cluster) {
+		return []*datapb.DataNodeInfo{}
+	}
+	i := 0
+	for _, node := range cluster {
+		if i == int(idx) {
+			//TODO add channel to node
+			return []*datapb.DataNodeInfo{
+				node,
+			}
+		}
+		i++
+	}
+	return []*datapb.DataNodeInfo{}
+}
+
 type channelAssignPolicy interface {
-	apply(cluster map[string]*datapb.DataNodeInfo, channel string) []*datapb.DataNodeInfo
+	apply(cluster map[string]*datapb.DataNodeInfo, channel string, collectionID UniqueID) []*datapb.DataNodeInfo
 }
 
 type allAssignPolicy struct {
@@ -85,7 +124,7 @@ func newAllAssignPolicy() channelAssignPolicy {
 	return &allAssignPolicy{}
 }
 
-func (p *allAssignPolicy) apply(cluster map[string]*datapb.DataNodeInfo, channel string) []*datapb.DataNodeInfo {
+func (p *allAssignPolicy) apply(cluster map[string]*datapb.DataNodeInfo, channel string, collectionID UniqueID) []*datapb.DataNodeInfo {
 	ret := make([]*datapb.DataNodeInfo, 0)
 	for _, node := range cluster {
 		fmt.Printf("xxxxnode: %v\n", node.Address)
@@ -100,8 +139,9 @@ func (p *allAssignPolicy) apply(cluster map[string]*datapb.DataNodeInfo, channel
 			continue
 		}
 		node.Channels = append(node.Channels, &datapb.ChannelStatus{
-			Name:  channel,
-			State: datapb.ChannelWatchState_Uncomplete,
+			Name:         channel,
+			State:        datapb.ChannelWatchState_Uncomplete,
+			CollectionID: collectionID,
 		})
 		fmt.Printf("channelxxxx: %v\n", node.Channels)
 		ret = append(ret, node)

internal/dataservice/server.go

@@ -171,7 +171,7 @@ func (s *Server) initCluster() error {
 		return err
 	}
 	sManager := newClusterSessionManager(s.dataClientCreator)
-	s.cluster = newCluster(s.ctx, dManager, sManager)
+	s.cluster = newCluster(s.ctx, dManager, sManager, s)
 	return nil
 }
 

internal/proto/data_service.proto

@@ -298,6 +298,7 @@ enum ChannelWatchState {
 message ChannelStatus {
   string name = 1;
   ChannelWatchState state=2;
+  int64 collectionID = 3;
 }
 
 message DataNodeInfo {