enhance: add proxy and datanode checker when wal balance startup (milvus-io#40877)

issue: milvus-io#40532

- balance should enable only when there's no proxy and datanode which
version is lower than 2.6.0

---------

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

Author: chyezh

internal/mocks/util/streamingutil/service/mock_discoverer/mock_Discoverer.go

@@ -68,7 +68,7 @@ type Client interface {
 func NewClient(etcdCli *clientv3.Client) Client {
 	// StreamingCoord is deployed on DataCoord node.
 	role := sessionutil.GetSessionPrefixByRole(typeutil.RootCoordRole)
-	rb := resolver.NewSessionExclusiveBuilder(etcdCli, role)
+	rb := resolver.NewSessionExclusiveBuilder(etcdCli, role, ">=2.6.0-dev")
 	dialTimeout := paramtable.Get().StreamingCoordGrpcClientCfg.DialTimeout.GetAsDuration(time.Millisecond)
 	dialOptions := getDialOptions(rb)
 	conn := lazygrpc.NewConn(func(ctx context.Context) (*grpc.ClientConn, error) {

internal/mocks/util/streamingutil/service/mock_resolver/mock_Resolver.go

@@ -10,6 +10,8 @@ import (
 
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/resource"
+	"github.com/milvus-io/milvus/internal/util/sessionutil"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
 	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
 	"github.com/milvus-io/milvus/pkg/v2/log"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
@@ -122,6 +124,11 @@ func (b *balancerImpl) execute() {
 		b.logger.Info("balancer execute finished")
 	}()
 
+	if err := b.blockUntilAllNodeIsGreaterThan260(b.ctx); err != nil {
+		b.logger.Warn("fail to block until all node is greater than 2.6.0", zap.Error(err))
+		return
+	}
+
 	balanceTimer := typeutil.NewBackoffTimer(&backoffConfigFetcher{})
 	nodeChanged, err := resource.Resource().StreamingNodeManagerClient().WatchNodeChanged(b.backgroundTaskNotifier.Context())
 	if err != nil {
@@ -164,6 +171,35 @@ func (b *balancerImpl) execute() {
 	}
 }
 
+// blockUntilAllNodeIsGreaterThan260 block until all node is greater than 2.6.0.
+// It's just a protection, but didn't promised that there will never be a node with version < 2.6.0 join the cluster.
+// These promise can only be achieved by the cluster dev-ops.
+func (b *balancerImpl) blockUntilAllNodeIsGreaterThan260(ctx context.Context) error {
+	doneErr := errors.New("done")
+	expectedRoles := []string{typeutil.ProxyRole, typeutil.DataNodeRole}
+	for _, role := range expectedRoles {
+		logger := b.logger.With(zap.String("role", role))
+		logger.Info("start to wait that the nodes is greater than 2.6.0")
+		// Check if there's any proxy or data node with version < 2.6.0.
+		proxyResolver := resolver.NewSessionBuilder(resource.Resource().ETCD(), sessionutil.GetSessionPrefixByRole(role), "<2.6.0")
+		r := proxyResolver.Resolver()
+		err := r.Watch(ctx, func(vs resolver.VersionedState) error {
+			if len(vs.Sessions()) == 0 {
+				return doneErr
+			}
+			logger.Info("session changes", zap.Int("sessionCount", len(vs.Sessions())))
+			return nil
+		})
+		if err != nil && !errors.Is(err, doneErr) {
+			logger.Info("fail to wait that the nodes is greater than 2.6.0", zap.Error(err))
+			return err
+		}
+		logger.Info("all nodes is greater than 2.6.0")
+		proxyResolver.Close()
+	}
+	return nil
+}
+
 // applyAllRequest apply all request in the request channel.
 func (b *balancerImpl) applyAllRequest() {
 	for {

internal/streamingcoord/client/client.go

@@ -2,6 +2,8 @@ package balancer_test
 
 import (
 	"context"
+	"encoding/json"
+	"path"
 	"testing"
 	"time"
 
@@ -14,15 +16,23 @@ import (
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
 	_ "github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/policy"
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/resource"
+	"github.com/milvus-io/milvus/internal/util/sessionutil"
 	"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/v2/util/etcd"
 	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
 	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
 )
 
 func TestBalancer(t *testing.T) {
 	paramtable.Init()
+	err := etcd.InitEtcdServer(true, "", t.TempDir(), "stdout", "info")
+	assert.NoError(t, err)
+	defer etcd.StopEtcdServer()
+
+	etcdClient, err := etcd.GetEmbedEtcdClient()
+	assert.NoError(t, err)
 
 	streamingNodeManager := mock_manager.NewMockManagerClient(t)
 	streamingNodeManager.EXPECT().WatchNodeChanged(mock.Anything).Return(make(chan struct{}), nil)
@@ -57,7 +67,7 @@ func TestBalancer(t *testing.T) {
 	}, nil)
 
 	catalog := mock_metastore.NewMockStreamingCoordCataLog(t)
-	resource.InitForTest(resource.OptStreamingCatalog(catalog), resource.OptStreamingManagerClient(streamingNodeManager))
+	resource.InitForTest(resource.OptETCD(etcdClient), resource.OptStreamingCatalog(catalog), resource.OptStreamingManagerClient(streamingNodeManager))
 	catalog.EXPECT().ListPChannel(mock.Anything).Unset()
 	catalog.EXPECT().ListPChannel(mock.Anything).RunAndReturn(func(ctx context.Context) ([]*streamingpb.PChannelMeta, error) {
 		return []*streamingpb.PChannelMeta{
@@ -89,11 +99,36 @@ func TestBalancer(t *testing.T) {
 	})
 	catalog.EXPECT().SavePChannels(mock.Anything, mock.Anything).Return(nil).Maybe()
 
+	// Test for lower datanode and proxy version protection.
+	metaRoot := paramtable.Get().EtcdCfg.MetaRootPath.GetValue()
+	proxyPath1 := path.Join(metaRoot, sessionutil.DefaultServiceRoot, typeutil.ProxyRole+"-1")
+	r := sessionutil.SessionRaw{Version: "2.5.11", ServerID: 1}
+	data, _ := json.Marshal(r)
+	resource.Resource().ETCD().Put(context.Background(), proxyPath1, string(data))
+	proxyPath2 := path.Join(metaRoot, sessionutil.DefaultServiceRoot, typeutil.ProxyRole+"-2")
+	r = sessionutil.SessionRaw{Version: "2.5.11", ServerID: 2}
+	data, _ = json.Marshal(r)
+	resource.Resource().ETCD().Put(context.Background(), proxyPath2, string(data))
+	metaRoot = paramtable.Get().EtcdCfg.MetaRootPath.GetValue()
+	dataNodePath := path.Join(metaRoot, sessionutil.DefaultServiceRoot, typeutil.DataNodeRole)
+	resource.Resource().ETCD().Put(context.Background(), dataNodePath, string(data))
+
 	ctx := context.Background()
 	b, err := balancer.RecoverBalancer(ctx, "pchannel_count_fair")
 	assert.NoError(t, err)
 	assert.NotNil(t, b)
 
+	ctx1, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+	err = b.WatchChannelAssignments(ctx1, func(version typeutil.VersionInt64Pair, relations []types.PChannelInfoAssigned) error {
+		assert.Len(t, relations, 2)
+		return nil
+	})
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+	resource.Resource().ETCD().Delete(context.Background(), proxyPath1)
+	resource.Resource().ETCD().Delete(context.Background(), proxyPath2)
+	resource.Resource().ETCD().Delete(context.Background(), dataNodePath)
+
 	b.MarkAsUnavailable(ctx, []types.PChannelInfo{{
 		Name: "test-channel-1",
 		Term: 1,

internal/streamingcoord/server/balancer/balancer_impl.go

@@ -47,7 +47,7 @@ type ManagerClient interface {
 // NewManagerClient creates a new manager client.
 func NewManagerClient(etcdCli *clientv3.Client) ManagerClient {
 	role := sessionutil.GetSessionPrefixByRole(typeutil.StreamingNodeRole)
-	rb := resolver.NewSessionBuilder(etcdCli, role)
+	rb := resolver.NewSessionBuilder(etcdCli, role, ">=2.6.0-dev")
 	dialTimeout := paramtable.Get().StreamingNodeGrpcClientCfg.DialTimeout.GetAsDuration(time.Millisecond)
 	dialOptions := getDialOptions(rb)
 	conn := lazygrpc.NewConn(func(ctx context.Context) (*grpc.ClientConn, error) {

internal/streamingcoord/server/balancer/balancer_test.go

@@ -61,7 +61,10 @@ func (c *managerClientImpl) CollectAllStatus(ctx context.Context) (map[int64]*ty
 	defer c.lifetime.Done()
 
 	// Get all discovered streamingnode.
-	state := c.rb.Resolver().GetLatestState()
+	state, err := c.rb.Resolver().GetLatestState(ctx)
+	if err != nil {
+		return nil, err
+	}
 	if len(state.State.Addresses) == 0 {
 		return make(map[int64]*types.StreamingNodeStatus), nil
 	}
@@ -73,7 +76,10 @@ func (c *managerClientImpl) CollectAllStatus(ctx context.Context) (map[int64]*ty
 	}
 
 	// Collect status may cost some time, so we need to check the lifetime again.
-	newState := c.rb.Resolver().GetLatestState()
+	newState, err := c.rb.Resolver().GetLatestState(ctx)
+	if err != nil {
+		return nil, err
+	}
 	if newState.Version.GT(state.Version) {
 		newSession := newState.Sessions()
 		for serverID := range result {

internal/streamingnode/client/manager/manager_client.go

@@ -50,11 +50,11 @@ func TestManager(t *testing.T) {
 	assert.False(t, ok)
 
 	// Test CollectAllStatus
-	r.EXPECT().GetLatestState().RunAndReturn(func() discoverer.VersionedState {
+	r.EXPECT().GetLatestState(mock.Anything).RunAndReturn(func(ctx context.Context) (discoverer.VersionedState, error) {
 		return discoverer.VersionedState{
 			Version: typeutil.VersionInt64(1),
 			State:   resolver.State{},
-		}
+		}, nil
 	})
 	// Not address here.
 	nodes, err := m.CollectAllStatus(context.Background())
@@ -75,11 +75,11 @@ func TestManager(t *testing.T) {
 		{1: false, 2: false, 3: true},
 		{1: true, 2: false},
 	}
-	r.EXPECT().GetLatestState().Unset()
-	r.EXPECT().GetLatestState().RunAndReturn(func() discoverer.VersionedState {
+	r.EXPECT().GetLatestState(mock.Anything).Unset()
+	r.EXPECT().GetLatestState(mock.Anything).RunAndReturn(func(ctx context.Context) (discoverer.VersionedState, error) {
 		s := newVersionedState(int64(i), states[i])
 		i++
-		return s
+		return s, nil
 	})
 
 	nodes, err = m.CollectAllStatus(context.Background())

internal/streamingnode/client/manager/manager_client_impl.go

@@ -36,10 +36,12 @@ func (d *channelAssignmentDiscoverer) NewVersionedState() VersionedState {
 
 // channelAssignmentDiscoverer implements the resolver.Discoverer interface.
 func (d *channelAssignmentDiscoverer) Discover(ctx context.Context, cb func(VersionedState) error) error {
-	// Always send the current state first.
-	// Outside logic may lost the last state before retry Discover function.
-	if err := cb(d.parseState()); err != nil {
-		return err
+	if d.lastDiscovery != nil {
+		// Always send the current state first if there's.
+		// Outside logic may lost the last state before retry Discover function.
+		if err := cb(d.parseState()); err != nil {
+			return err
+		}
 	}
 	return d.assignmentWatcher.AssignmentDiscover(ctx, func(assignments *types.VersionedStreamingNodeAssignments) error {
 		d.lastDiscovery = assignments
@@ -50,10 +52,6 @@ func (d *channelAssignmentDiscoverer) Discover(ctx context.Context, cb func(Vers
 // parseState parses the addresses from the discovery response.
 // Always perform a copy here.
 func (d *channelAssignmentDiscoverer) parseState() VersionedState {
-	if d.lastDiscovery == nil {
-		return d.NewVersionedState()
-	}
-
 	addrs := make([]resolver.Address, 0, len(d.lastDiscovery.Assignments))
 	for _, assignment := range d.lastDiscovery.Assignments {
 		assignment := assignment

internal/streamingnode/client/manager/manager_test.go

@@ -36,14 +36,8 @@ func TestChannelAssignmentDiscoverer(t *testing.T) {
 		})
 
 	d := NewChannelAssignmentDiscoverer(w)
-	s := d.NewVersionedState()
-	assert.True(t, s.Version.EQ(typeutil.VersionInt64Pair{Global: -1, Local: -1}))
 
 	expected := []*types.VersionedStreamingNodeAssignments{
-		{
-			Version:     typeutil.VersionInt64Pair{Global: -1, Local: -1},
-			Assignments: map[int64]types.StreamingNodeAssignment{},
-		},
 		{
 			Version: typeutil.VersionInt64Pair{
 				Global: 1,
@@ -81,6 +75,7 @@ func TestChannelAssignmentDiscoverer(t *testing.T) {
 		},
 	}
 
+	ch <- expected[0]
 	idx := 0
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

internal/util/streamingutil/service/discoverer/channel_assignment_discoverer.go

@@ -13,9 +13,6 @@ import (
 // 1. concurrent safe.
 // 2. the version of discovery may be repeated or decreasing. So user should check the version in callback.
 type Discoverer interface {
-	// NewVersionedState returns a lowest versioned state.
-	NewVersionedState() VersionedState
-
 	// Discover watches the service discovery on these goroutine.
 	// 1. Call the callback when the discovery is changed, and block until the discovery is canceled or break down.
 	// 2. Discover should always send the current state first and then block.