Integrate message stream

Signed-off-by: groot <yihua.mo@zilliz.com>

cmd/datanode/main.go

@@ -8,6 +8,7 @@ import (
 	"syscall"
 
 	distributed "github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 )
 
 func main() {
@@ -15,7 +16,9 @@ func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	dn, err := distributed.NewDataNode(ctx)
+	msFactory := pulsarms.NewFactory()
+
+	dn, err := distributed.NewDataNode(ctx, msFactory)
 	if err != nil {
 		panic(err)
 	}

cmd/dataservice/main.go

@@ -8,12 +8,15 @@ import (
 	"syscall"
 
 	"github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 )
 
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 
-	svr, err := components.NewDataService(ctx)
+	msFactory := pulsarms.NewFactory()
+
+	svr, err := components.NewDataService(ctx, msFactory)
 	if err != nil {
 		panic(err)
 	}

cmd/distributed/components/data_node.go

@@ -12,6 +12,7 @@ import (
 	dsc "github.com/zilliztech/milvus-distributed/internal/distributed/dataservice"
 	msc "github.com/zilliztech/milvus-distributed/internal/distributed/masterservice"
 
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
 	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
 )
@@ -24,12 +25,12 @@ type DataNode struct {
 	dataService   *dsc.Client
 }
 
-func NewDataNode(ctx context.Context) (*DataNode, error) {
+func NewDataNode(ctx context.Context, factory msgstream.Factory) (*DataNode, error) {
 
 	const retry = 10
 	const interval = 200
 
-	svr, err := dnc.New(ctx)
+	svr, err := dnc.New(ctx, factory)
 	if err != nil {
 		panic(err)
 	}

cmd/distributed/components/data_service.go

@@ -9,6 +9,7 @@ import (
 
 	ms "github.com/zilliztech/milvus-distributed/internal/distributed/masterservice"
 	"github.com/zilliztech/milvus-distributed/internal/masterservice"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
 	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
 
 	"github.com/zilliztech/milvus-distributed/internal/distributed/dataservice"
@@ -20,8 +21,8 @@ type DataService struct {
 	masterClient *ms.GrpcClient
 }
 
-func NewDataService(ctx context.Context) (*DataService, error) {
-	service := dataservice.NewGrpcService(ctx)
+func NewDataService(ctx context.Context, factory msgstream.Factory) (*DataService, error) {
+	service := dataservice.NewGrpcService(ctx, factory)
 
 	masterservice.Params.Init()
 	client, err := ms.NewGrpcClient(fmt.Sprintf("%s:%d", masterservice.Params.Address, masterservice.Params.Port), 30*time.Second)

cmd/distributed/components/master_service.go

@@ -15,6 +15,7 @@ import (
 	qsc "github.com/zilliztech/milvus-distributed/internal/distributed/queryservice/client"
 	is "github.com/zilliztech/milvus-distributed/internal/indexservice"
 	ms "github.com/zilliztech/milvus-distributed/internal/masterservice"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
 	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
 	qs "github.com/zilliztech/milvus-distributed/internal/queryservice"
@@ -30,10 +31,10 @@ type MasterService struct {
 	queryService *qsc.Client
 }
 
-func NewMasterService(ctx context.Context) (*MasterService, error) {
+func NewMasterService(ctx context.Context, factory msgstream.Factory) (*MasterService, error) {
 	const reTryCnt = 3
 
-	svr, err := msc.NewGrpcServer(ctx)
+	svr, err := msc.NewGrpcServer(ctx, factory)
 	if err != nil {
 		return nil, err
 	}

cmd/distributed/components/proxy_node.go

@@ -3,6 +3,8 @@ package components
 import (
 	"context"
 
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
+
 	grpcproxynode "github.com/zilliztech/milvus-distributed/internal/distributed/proxynode"
 )
 
@@ -10,9 +12,9 @@ type ProxyNode struct {
 	svr *grpcproxynode.Server
 }
 
-func NewProxyNode(ctx context.Context) (*ProxyNode, error) {
+func NewProxyNode(ctx context.Context, factory msgstream.Factory) (*ProxyNode, error) {
 	n := &ProxyNode{}
-	svr, err := grpcproxynode.NewServer(ctx)
+	svr, err := grpcproxynode.NewServer(ctx, factory)
 	if err != nil {
 		return nil, err
 	}

cmd/distributed/components/proxy_service.go

@@ -3,6 +3,8 @@ package components
 import (
 	"context"
 
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
+
 	grpcproxyservice "github.com/zilliztech/milvus-distributed/internal/distributed/proxyservice"
 )
 
@@ -10,9 +12,9 @@ type ProxyService struct {
 	svr *grpcproxyservice.Server
 }
 
-func NewProxyService(ctx context.Context) (*ProxyService, error) {
+func NewProxyService(ctx context.Context, factory msgstream.Factory) (*ProxyService, error) {
 	service := &ProxyService{}
-	svr, err := grpcproxyservice.NewServer(ctx)
+	svr, err := grpcproxyservice.NewServer(ctx, factory)
 	if err != nil {
 		return nil, err
 	}

cmd/distributed/components/query_node.go

@@ -11,6 +11,7 @@ import (
 	msc "github.com/zilliztech/milvus-distributed/internal/distributed/masterservice"
 	qns "github.com/zilliztech/milvus-distributed/internal/distributed/querynode"
 	qsc "github.com/zilliztech/milvus-distributed/internal/distributed/queryservice/client"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
 
 	ds "github.com/zilliztech/milvus-distributed/internal/dataservice"
 	is "github.com/zilliztech/milvus-distributed/internal/indexservice"
@@ -32,11 +33,11 @@ type QueryNode struct {
 	queryService  *qsc.Client
 }
 
-func NewQueryNode(ctx context.Context) (*QueryNode, error) {
+func NewQueryNode(ctx context.Context, factory msgstream.Factory) (*QueryNode, error) {
 	const retry = 10
 	const interval = 500
 
-	svr, err := qns.NewServer(ctx)
+	svr, err := qns.NewServer(ctx, factory)
 	if err != nil {
 		panic(err)
 	}

cmd/distributed/components/query_service.go

@@ -12,6 +12,7 @@ import (
 
 	ds "github.com/zilliztech/milvus-distributed/internal/dataservice"
 	ms "github.com/zilliztech/milvus-distributed/internal/masterservice"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
 	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
 	"github.com/zilliztech/milvus-distributed/internal/queryservice"
@@ -25,12 +26,12 @@ type QueryService struct {
 	masterService *msc.GrpcClient
 }
 
-func NewQueryService(ctx context.Context) (*QueryService, error) {
+func NewQueryService(ctx context.Context, factory msgstream.Factory) (*QueryService, error) {
 	const retry = 10
 	const interval = 200
 
 	queryservice.Params.Init()
-	svr, err := qs.NewServer(ctx)
+	svr, err := qs.NewServer(ctx, factory)
 	if err != nil {
 		panic(err)
 	}

cmd/distributed/main.go

@@ -64,7 +64,7 @@ func run(serverType string) error {
 	default:
 		return errors.Errorf("unknown server type = %s", serverType)
 	}
-	role.Run()
+	role.Run(false)
 	return nil
 }
 

cmd/distributed/roles/roles.go

@@ -9,8 +9,18 @@ import (
 	"syscall"
 
 	"github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/rmqms"
 )
 
+func newMsgFactory(localMsg bool) msgstream.Factory {
+	if localMsg {
+		return rmqms.NewFactory()
+	}
+	return pulsarms.NewFactory()
+}
+
 type MilvusRoles struct {
 	EnableMaster           bool `env:"ENABLE_MASTER"`
 	EnableProxyService     bool `env:"ENABLE_PROXY_SERVICE"`
@@ -42,7 +52,7 @@ func (mr *MilvusRoles) EnvValue(env string) bool {
 	return false
 }
 
-func (mr *MilvusRoles) Run() {
+func (mr *MilvusRoles) Run(localMsg bool) {
 	if !mr.HasAnyRole() {
 		log.Printf("set the roles please ...")
 		return
@@ -55,8 +65,9 @@ func (mr *MilvusRoles) Run() {
 	if mr.EnableMaster {
 		log.Print("start as master service")
 		go func() {
+			factory := newMsgFactory(localMsg)
 			var err error
-			masterService, err = components.NewMasterService(ctx)
+			masterService, err = components.NewMasterService(ctx, factory)
 			if err != nil {
 				panic(err)
 			}
@@ -68,8 +79,9 @@ func (mr *MilvusRoles) Run() {
 	if mr.EnableProxyService {
 		log.Print("start as proxy service")
 		go func() {
+			factory := newMsgFactory(localMsg)
 			var err error
-			proxyService, err = components.NewProxyService(ctx)
+			proxyService, err = components.NewProxyService(ctx, factory)
 			if err != nil {
 				panic(err)
 			}
@@ -81,8 +93,9 @@ func (mr *MilvusRoles) Run() {
 	if mr.EnableProxyNode {
 		log.Print("start as proxy node")
 		go func() {
+			factory := newMsgFactory(localMsg)
 			var err error
-			proxyNode, err = components.NewProxyNode(ctx)
+			proxyNode, err = components.NewProxyNode(ctx, factory)
 			if err != nil {
 				panic(err)
 			}
@@ -94,8 +107,9 @@ func (mr *MilvusRoles) Run() {
 	if mr.EnableQueryService {
 		log.Print("start as query service")
 		go func() {
+			factory := newMsgFactory(localMsg)
 			var err error
-			queryService, err = components.NewQueryService(ctx)
+			queryService, err = components.NewQueryService(ctx, factory)
 			if err != nil {
 				panic(err)
 			}
@@ -107,8 +121,9 @@ func (mr *MilvusRoles) Run() {
 	if mr.EnableQueryNode {
 		log.Print("start as query node")
 		go func() {
+			factory := newMsgFactory(localMsg)
 			var err error
-			queryNode, err = components.NewQueryNode(ctx)
+			queryNode, err = components.NewQueryNode(ctx, factory)
 			if err != nil {
 				panic(err)
 			}
@@ -120,8 +135,9 @@ func (mr *MilvusRoles) Run() {
 	if mr.EnableDataService {
 		log.Print("start as data service")
 		go func() {
+			factory := newMsgFactory(localMsg)
 			var err error
-			dataService, err = components.NewDataService(ctx)
+			dataService, err = components.NewDataService(ctx, factory)
 			if err != nil {
 				panic(err)
 			}
@@ -133,8 +149,9 @@ func (mr *MilvusRoles) Run() {
 	if mr.EnableDataNode {
 		log.Print("start as data node")
 		go func() {
+			factory := newMsgFactory(localMsg)
 			var err error
-			dataNode, err = components.NewDataNode(ctx)
+			dataNode, err = components.NewDataNode(ctx, factory)
 			if err != nil {
 				panic(err)
 			}

cmd/masterservice/main.go

@@ -8,13 +8,15 @@ import (
 	"syscall"
 
 	distributed "github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 )
 
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	ms, err := distributed.NewMasterService(ctx)
+	msFactory := pulsarms.NewFactory()
+	ms, err := distributed.NewMasterService(ctx, msFactory)
 	if err != nil {
 		panic(err)
 	}

cmd/proxy/node/proxy_node.go

@@ -8,13 +8,15 @@ import (
 	"syscall"
 
 	"github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 
 	"go.uber.org/zap"
 )
 
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
-	n, err := components.NewProxyNode(ctx)
+	msFactory := pulsarms.NewFactory()
+	n, err := components.NewProxyNode(ctx, msFactory)
 	if err != nil {
 		log.Print("create server failed", zap.Error(err))
 	}

cmd/proxy/service/proxy_service.go

@@ -8,13 +8,15 @@ import (
 	"syscall"
 
 	"github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 
 	"go.uber.org/zap"
 )
 
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
-	s, err := components.NewProxyService(ctx)
+	msFactory := pulsarms.NewFactory()
+	s, err := components.NewProxyService(ctx, msFactory)
 	if err != nil {
 		log.Fatal("create proxy service error: " + err.Error())
 	}

cmd/querynode/querynode.go

@@ -8,13 +8,15 @@ import (
 	"syscall"
 
 	distributed "github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 )
 
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	svr, err := distributed.NewQueryNode(ctx)
+	msFactory := pulsarms.NewFactory()
+	svr, err := distributed.NewQueryNode(ctx, msFactory)
 
 	if err != nil {
 		panic(err)

cmd/queryservice/queryservice.go

@@ -8,13 +8,16 @@ import (
 	"syscall"
 
 	distributed "github.com/zilliztech/milvus-distributed/cmd/distributed/components"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 )
 
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	svr, err := distributed.NewQueryService(ctx)
+	msFactory := pulsarms.NewFactory()
+
+	svr, err := distributed.NewQueryService(ctx, msFactory)
 	if err != nil {
 		panic(err)
 	}

cmd/singlenode/main.go

@@ -20,5 +20,5 @@ func initRoles(roles *roles.MilvusRoles) {
 func main() {
 	var roles roles.MilvusRoles
 	initRoles(&roles)
-	roles.Run()
+	roles.Run(false)
 }

go.mod

@@ -27,6 +27,7 @@ require (
 	github.com/klauspost/compress v1.10.11 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/minio/minio-go/v7 v7.0.5
+	github.com/mitchellh/mapstructure v1.1.2
 	github.com/modern-go/reflect2 v1.0.1
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
 	github.com/oklog/run v1.1.0

internal/datanode/data_node.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/zilliztech/milvus-distributed/internal/errors"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
 	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
@@ -70,10 +71,12 @@ type (
 		replica   Replica
 
 		closer io.Closer
+
+		msFactory msgstream.Factory
 	}
 )
 
-func NewDataNode(ctx context.Context) *DataNode {
+func NewDataNode(ctx context.Context, factory msgstream.Factory) *DataNode {
 
 	Params.Init()
 	ctx2, cancel2 := context.WithCancel(ctx)
@@ -89,6 +92,7 @@ func NewDataNode(ctx context.Context) *DataNode {
 		masterService:   nil,
 		dataService:     nil,
 		replica:         nil,
+		msFactory:       factory,
 	}
 
 	node.State.Store(internalpb2.StateCode_INITIALIZING)
@@ -165,7 +169,7 @@ func (node *DataNode) Init() error {
 	chanSize := 100
 	node.flushChan = make(chan *flushMsg, chanSize)
 
-	node.dataSyncService = newDataSyncService(node.ctx, node.flushChan, replica, alloc)
+	node.dataSyncService = newDataSyncService(node.ctx, node.flushChan, replica, alloc, node.msFactory)
 	node.dataSyncService.init()
 	node.metaService = newMetaService(node.ctx, replica, node.masterService)
 

internal/datanode/data_sync_service.go

@@ -5,6 +5,7 @@ import (
 	"log"
 
 	etcdkv "github.com/zilliztech/milvus-distributed/internal/kv/etcd"
+	"github.com/zilliztech/milvus-distributed/internal/msgstream"
 	"github.com/zilliztech/milvus-distributed/internal/util/flowgraph"
 	"go.etcd.io/etcd/clientv3"
 )
@@ -15,16 +16,18 @@ type dataSyncService struct {
 	flushChan   chan *flushMsg
 	replica     Replica
 	idAllocator allocator
+	msFactory   msgstream.Factory
 }
 
 func newDataSyncService(ctx context.Context, flushChan chan *flushMsg,
-	replica Replica, alloc allocator) *dataSyncService {
+	replica Replica, alloc allocator, factory msgstream.Factory) *dataSyncService {
 	service := &dataSyncService{
 		ctx:         ctx,
 		fg:          nil,
 		flushChan:   flushChan,
 		replica:     replica,
 		idAllocator: alloc,
+		msFactory:   factory,
 	}
 	return service
 }
@@ -65,12 +68,21 @@ func (dsService *dataSyncService) initNodes() {
 
 	dsService.fg = flowgraph.NewTimeTickedFlowGraph(dsService.ctx)
 
-	var dmStreamNode Node = newDmInputNode(dsService.ctx)
-	var ddStreamNode Node = newDDInputNode(dsService.ctx)
+	m := map[string]interface{}{
+		"PulsarAddress":  Params.PulsarAddress,
+		"ReceiveBufSize": 1024,
+		"PulsarBufSize":  1024}
+	err = dsService.msFactory.SetParams(m)
+	if err != nil {
+		panic(err)
+	}
+
+	var dmStreamNode Node = newDmInputNode(dsService.ctx, dsService.msFactory)
+	var ddStreamNode Node = newDDInputNode(dsService.ctx, dsService.msFactory)
 
 	var filterDmNode Node = newFilteredDmNode()
 	var ddNode Node = newDDNode(dsService.ctx, mt, dsService.flushChan, dsService.replica, dsService.idAllocator)
-	var insertBufferNode Node = newInsertBufferNode(dsService.ctx, mt, dsService.replica, dsService.idAllocator)
+	var insertBufferNode Node = newInsertBufferNode(dsService.ctx, mt, dsService.replica, dsService.idAllocator, dsService.msFactory)
 	var gcNode Node = newGCNode(dsService.replica)
 
 	dsService.fg.AddNode(&dmStreamNode)