enhance: implement balancer at streaming coord

internal/metastore/catalog.go

@@ -9,6 +9,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/indexpb"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
@@ -186,3 +187,14 @@ type QueryCoordCatalog interface {
 	RemoveCollectionTarget(collectionID int64) error
 	GetCollectionTargets() (map[int64]*querypb.CollectionTarget, error)
 }
+
+// StreamingCoordCataLog is the interface for streamingcoord catalog
+type StreamingCoordCataLog interface {
+	// physical channel watch related
+
+	// ListPChannel list all pchannels on milvus.
+	ListPChannel(ctx context.Context) ([]*streamingpb.PChannelMeta, error)
+
+	// SavePChannel save a pchannel info to metastore.
+	SavePChannels(ctx context.Context, info []*streamingpb.PChannelMeta) error
+}

internal/metastore/kv/streamingcoord/constant.go

@@ -0,0 +1,6 @@
+package streamingcoord
+
+const (
+	MetaPrefix   = "streamingcoord-meta"
+	PChannelMeta = MetaPrefix + "/pchannel-meta"
+)

internal/metastore/kv/streamingcoord/kv_catalog.go

@@ -0,0 +1,62 @@
+package streamingcoord
+
+import (
+	"context"
+
+	"github.com/cockroachdb/errors"
+	"github.com/golang/protobuf/proto"
+
+	"github.com/milvus-io/milvus/internal/metastore"
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/pkg/kv"
+)
+
+// NewCataLog creates a new catalog instance
+func NewCataLog(metaKV kv.MetaKv) metastore.StreamingCoordCataLog {
+	return &catalog{
+		metaKV: metaKV,
+	}
+}
+
+// catalog is a kv based catalog.
+type catalog struct {
+	metaKV kv.MetaKv
+}
+
+// ListPChannels returns all pchannels
+func (c *catalog) ListPChannel(ctx context.Context) ([]*streamingpb.PChannelMeta, error) {
+	keys, values, err := c.metaKV.LoadWithPrefix(PChannelMeta)
+	if err != nil {
+		return nil, err
+	}
+
+	infos := make([]*streamingpb.PChannelMeta, 0, len(values))
+	for k, value := range values {
+		info := &streamingpb.PChannelMeta{}
+		err = proto.Unmarshal([]byte(value), info)
+		if err != nil {
+			return nil, errors.Wrapf(err, "unmarshal pchannel %s failed", keys[k])
+		}
+		infos = append(infos, info)
+	}
+	return infos, nil
+}
+
+// SavePChannels saves a pchannel
+func (c *catalog) SavePChannels(ctx context.Context, infos []*streamingpb.PChannelMeta) error {
+	kvs := make(map[string]string, len(infos))
+	for _, info := range infos {
+		key := buildPChannelInfoPath(info.GetChannel().GetName())
+		v, err := proto.Marshal(info)
+		if err != nil {
+			return errors.Wrapf(err, "marshal pchannel %s failed", info.GetChannel().GetName())
+		}
+		kvs[key] = string(v)
+	}
+	return c.metaKV.MultiSave(kvs)
+}
+
+// buildPChannelInfoPath builds the path for pchannel info.
+func buildPChannelInfoPath(name string) string {
+	return PChannelMeta + "/" + name
+}

internal/metastore/kv/streamingcoord/kv_catalog_test.go

@@ -0,0 +1,66 @@
+package streamingcoord
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/pkg/mocks/mock_kv"
+)
+
+func TestCatalog(t *testing.T) {
+	kv := mock_kv.NewMockMetaKv(t)
+
+	kvStorage := make(map[string]string)
+	kv.EXPECT().LoadWithPrefix(mock.Anything).RunAndReturn(func(s string) ([]string, []string, error) {
+		keys := make([]string, 0, len(kvStorage))
+		vals := make([]string, 0, len(kvStorage))
+		for k, v := range kvStorage {
+			keys = append(keys, k)
+			vals = append(vals, v)
+		}
+		return keys, vals, nil
+	})
+	kv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(kvs map[string]string) error {
+		for k, v := range kvs {
+			kvStorage[k] = v
+		}
+		return nil
+	})
+
+	catalog := NewCataLog(kv)
+	metas, err := catalog.ListPChannel(context.Background())
+	assert.NoError(t, err)
+	assert.Empty(t, metas)
+
+	err = catalog.SavePChannels(context.Background(), []*streamingpb.PChannelMeta{
+		{
+			Channel: &streamingpb.PChannelInfo{Name: "test", Term: 1},
+			Node:    &streamingpb.StreamingNodeInfo{ServerId: 1},
+		},
+		{
+			Channel: &streamingpb.PChannelInfo{Name: "test2", Term: 1},
+			Node:    &streamingpb.StreamingNodeInfo{ServerId: 1},
+		},
+	})
+	assert.NoError(t, err)
+
+	metas, err = catalog.ListPChannel(context.Background())
+	assert.NoError(t, err)
+	assert.Len(t, metas, 2)
+
+	// error path.
+	kv.EXPECT().LoadWithPrefix(mock.Anything).Unset()
+	kv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, errors.New("load error"))
+	metas, err = catalog.ListPChannel(context.Background())
+	assert.Error(t, err)
+	assert.Nil(t, metas)
+
+	kv.EXPECT().MultiSave(mock.Anything).Unset()
+	kv.EXPECT().MultiSave(mock.Anything).Return(errors.New("save error"))
+	assert.Error(t, err)
+}