Split watch operations to many transactions (#10655)

issue: #10633
Signed-off-by: sunby <bingyi.sun@zilliz.com>

internal/datacoord/channel_store.go

@@ -30,8 +30,9 @@ import (
 )
 
 const (
-	bufferID  = math.MinInt64
-	delimeter = "/"
+	bufferID            = math.MinInt64
+	delimeter           = "/"
+	maxOperationsPerTxn = 128
 )
 
 var errUnknownOpType error = errors.New("unknown operation type")
@@ -172,6 +173,50 @@ func (c *ChannelStore) Add(nodeID int64) {
 
 // Update applies the operations in opSet
 func (c *ChannelStore) Update(opSet ChannelOpSet) error {
+	totalChannelNum := 0
+	for _, op := range opSet {
+		totalChannelNum += len(op.Channels)
+	}
+	if totalChannelNum <= maxOperationsPerTxn {
+		return c.update(opSet)
+	}
+	// split opset to many txn; same channel's operations should be executed in one txn.
+	channelsOpSet := make(map[string]ChannelOpSet)
+	for _, op := range opSet {
+		for i, ch := range op.Channels {
+			chOp := &ChannelOp{
+				Type:     op.Type,
+				NodeID:   op.NodeID,
+				Channels: []*channel{ch},
+			}
+			if op.Type == Add {
+				chOp.ChannelWatchInfos = []*datapb.ChannelWatchInfo{op.ChannelWatchInfos[i]}
+			}
+			channelsOpSet[ch.name] = append(channelsOpSet[ch.name], chOp)
+		}
+	}
+
+	// execute a txn per 128 operations
+	count := 0
+	operations := make([]*ChannelOp, 0, maxOperationsPerTxn)
+	for _, opset := range channelsOpSet {
+		if count+len(opset) > maxOperationsPerTxn {
+			if err := c.update(operations); err != nil {
+				return err
+			}
+			count = 0
+			operations = make([]*ChannelOp, 0, maxOperationsPerTxn)
+		}
+		count += len(opset)
+		operations = append(operations, opset...)
+	}
+	if count == 0 {
+		return nil
+	}
+	return c.update(operations)
+}
+
+func (c *ChannelStore) update(opSet ChannelOpSet) error {
 	if err := c.txn(opSet); err != nil {
 		return err
 	}

internal/datacoord/channel_store_test.go

@@ -0,0 +1,159 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 datacoord
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/milvus-io/milvus/internal/kv"
+	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/stretchr/testify/assert"
+)
+
+type mockTxnKv struct{}
+
+func (m *mockTxnKv) Load(key string) (string, error) {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) MultiLoad(keys []string) ([]string, error) {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) LoadWithPrefix(key string) ([]string, []string, error) {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) Save(key string, value string) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) MultiSave(kvs map[string]string) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) Remove(key string) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) MultiRemove(keys []string) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) RemoveWithPrefix(key string) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) Close() {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) MultiSaveAndRemove(saves map[string]string, removals []string) error {
+	if len(saves)+len(removals) > 128 {
+		return errors.New("too many operations")
+	}
+	return nil
+}
+
+func (m *mockTxnKv) MultiRemoveWithPrefix(keys []string) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockTxnKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func genNodeChannelInfos(id int64, num int) *NodeChannelInfo {
+	channels := make([]*channel, 0, num)
+	for i := 0; i < num; i++ {
+		name := fmt.Sprintf("ch%d", i)
+		channels = append(channels, &channel{name, 1})
+	}
+	return &NodeChannelInfo{
+		NodeID:   id,
+		Channels: channels,
+	}
+}
+
+func genChannelOperations(from, to int64, num int) ChannelOpSet {
+	ops := make([]*ChannelOp, 0, 2)
+	channels := make([]*channel, 0, num)
+	channelWatchInfos := make([]*datapb.ChannelWatchInfo, 0, num)
+	for i := 0; i < num; i++ {
+		name := fmt.Sprintf("ch%d", i)
+		channels = append(channels, &channel{name, 1})
+		channelWatchInfos = append(channelWatchInfos, &datapb.ChannelWatchInfo{})
+	}
+
+	ops = append(ops, &ChannelOp{
+		Type:     Delete,
+		NodeID:   from,
+		Channels: channels,
+	})
+
+	ops = append(ops, &ChannelOp{
+		Type:              Add,
+		NodeID:            to,
+		Channels:          channels,
+		ChannelWatchInfos: channelWatchInfos,
+	})
+
+	return ops
+}
+
+func TestChannelStore_Update(t *testing.T) {
+	type fields struct {
+		store        kv.TxnKV
+		channelsInfo map[int64]*NodeChannelInfo
+	}
+	type args struct {
+		opSet ChannelOpSet
+	}
+	tests := []struct {
+		name    string
+		fields  fields
+		args    args
+		wantErr bool
+	}{
+		{
+			"test more than 128 operations",
+			fields{
+				&mockTxnKv{},
+				map[int64]*NodeChannelInfo{
+					1: genNodeChannelInfos(1, 500),
+					2: {NodeID: 2},
+				},
+			},
+			args{
+				genChannelOperations(1, 2, 250),
+			},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &ChannelStore{
+				store:        tt.fields.store,
+				channelsInfo: tt.fields.channelsInfo,
+			}
+			err := c.Update(tt.args.opSet)
+			assert.Equal(t, tt.wantErr, err != nil)
+		})
+	}
+}