From 7a3a7213804342405714830e2991922554f5c6de Mon Sep 17 00:00:00 2001 From: jaime Date: Thu, 3 Mar 2022 21:57:56 +0800 Subject: [PATCH] Reconstruct mqstream module (#15784) Signed-off-by: yun.zhang --- cmd/components/data_coord.go | 2 +- cmd/components/data_node.go | 2 +- cmd/components/proxy.go | 2 +- cmd/components/query_coord.go | 2 +- cmd/components/query_node.go | 2 +- cmd/components/root_coord.go | 2 +- cmd/roles/roles.go | 9 +- internal/datacoord/channel_manager.go | 2 +- internal/datacoord/server.go | 6 +- internal/datacoord/server_test.go | 2 +- internal/datanode/data_node.go | 2 +- internal/datanode/data_node_test.go | 2 +- internal/datanode/data_sync_service.go | 2 +- internal/datanode/data_sync_service_test.go | 2 +- internal/datanode/flow_graph_dd_node.go | 2 +- internal/datanode/flow_graph_dd_node_test.go | 2 +- internal/datanode/flow_graph_delete_node.go | 2 +- .../datanode/flow_graph_delete_node_test.go | 2 +- .../flow_graph_dmstream_input_node_test.go | 6 +- .../datanode/flow_graph_insert_buffer_node.go | 2 +- .../flow_graph_insert_buffer_node_test.go | 2 +- internal/datanode/flow_graph_message.go | 2 +- internal/datanode/mock_test.go | 2 +- internal/distributed/datacoord/service.go | 2 +- internal/distributed/datanode/service.go | 2 +- internal/distributed/proxy/service.go | 2 +- internal/distributed/querycoord/service.go | 2 +- internal/distributed/querynode/service.go | 2 +- internal/distributed/rootcoord/service.go | 2 +- .../distributed/rootcoord/service_test.go | 2 +- .../mqimpl/rocksmq/client}/client.go | 8 +- .../mqimpl/rocksmq/client}/client_impl.go | 10 +- .../rocksmq/client}/client_impl_test.go | 16 +- .../mqimpl/rocksmq/client}/consumer.go | 6 +- .../mqimpl/rocksmq/client}/consumer_impl.go | 2 +- .../rocksmq/client}/consumer_impl_test.go | 6 +- .../mqimpl/rocksmq/client}/error.go | 2 +- .../mqimpl/rocksmq/client}/producer.go | 2 +- .../mqimpl/rocksmq/client}/producer_impl.go | 5 +- .../rocksmq/client}/producer_impl_test.go | 2 +- .../mqimpl/rocksmq/client}/reader.go | 2 +- .../mqimpl/rocksmq/client}/reader_impl.go | 2 +- .../rocksmq/client}/reader_impl_test.go | 4 +- .../mqimpl/rocksmq/client}/test_helper.go | 10 +- .../mqimpl/rocksmq/server}/global_rmq.go | 2 +- .../mqimpl/rocksmq/server}/global_rmq_test.go | 2 +- .../mqimpl/rocksmq/server}/rocksmq.go | 2 +- .../mqimpl/rocksmq/server}/rocksmq_impl.go | 2 +- .../rocksmq/server}/rocksmq_impl_test.go | 2 +- .../mqimpl/rocksmq/server}/rocksmq_reader.go | 2 +- .../rocksmq/server}/rocksmq_retention.go | 2 +- .../rocksmq/server}/rocksmq_retention_test.go | 2 +- internal/{ => mq}/msgstream/OWNERS | 0 internal/{ => mq}/msgstream/mq_factory.go | 22 +-- .../{ => mq}/msgstream/mq_factory_test.go | 0 internal/{ => mq}/msgstream/mq_msgstream.go | 78 ++++---- .../{ => mq}/msgstream/mq_msgstream_test.go | 104 +++++------ .../msgstream/mqwrapper}/client.go | 2 +- .../msgstream/mqwrapper}/consumer.go | 2 +- .../mqclient => mq/msgstream/mqwrapper}/id.go | 2 +- .../msgstream/mqwrapper}/message.go | 2 +- .../msgstream/mqwrapper}/producer.go | 2 +- .../mqwrapper/pulsar}/pulsar_client.go | 26 +-- .../mqwrapper/pulsar}/pulsar_client_test.go | 84 +++++---- .../mqwrapper/pulsar}/pulsar_consumer.go | 22 +-- .../mqwrapper/pulsar}/pulsar_consumer_test.go | 18 +- .../msgstream/mqwrapper/pulsar}/pulsar_id.go | 14 +- .../mqwrapper/pulsar}/pulsar_id_test.go | 8 +- .../mqwrapper/pulsar}/pulsar_message.go | 7 +- .../mqwrapper/pulsar}/pulsar_producer.go | 8 +- .../mqwrapper/pulsar}/pulsar_producer_test.go | 10 +- .../mqwrapper/pulsar}/pulsar_reader.go | 8 +- .../mqwrapper/pulsar}/pulsar_reader_test.go | 18 +- .../msgstream/mqwrapper}/reader.go | 2 +- .../msgstream/mqwrapper/rmq}/rmq_client.go | 46 +++-- .../mqwrapper/rmq}/rmq_client_test.go | 57 +++--- .../msgstream/mqwrapper/rmq}/rmq_consumer.go | 26 +-- .../msgstream/mqwrapper/rmq}/rmq_id.go | 9 +- .../msgstream/mqwrapper/rmq}/rmq_id_test.go | 2 +- .../msgstream/mqwrapper/rmq}/rmq_message.go | 11 +- .../msgstream/mqwrapper/rmq}/rmq_producer.go | 14 +- .../msgstream/mqwrapper/rmq}/rmq_reader.go | 14 +- internal/{ => mq}/msgstream/msg.go | 0 internal/{ => mq}/msgstream/msg_test.go | 0 internal/{ => mq}/msgstream/msgstream.go | 6 +- internal/{ => mq}/msgstream/repack_func.go | 0 internal/{ => mq}/msgstream/trace.go | 0 internal/{ => mq}/msgstream/unmarshal.go | 0 internal/{ => mq}/msgstream/unmarshal_test.go | 0 internal/proxy/channels_mgr.go | 2 +- internal/proxy/impl.go | 2 +- internal/proxy/mock_test.go | 6 +- internal/proxy/proxy.go | 2 +- internal/proxy/proxy_test.go | 2 +- internal/proxy/repack_func.go | 2 +- internal/proxy/repack_func_test.go | 2 +- internal/proxy/task.go | 2 +- internal/proxy/task_scheduler.go | 6 +- internal/proxy/task_test.go | 2 +- internal/querycoord/channel_unsubscribe.go | 2 +- .../querycoord/channel_unsubscribe_test.go | 2 +- internal/querycoord/cluster_test.go | 2 +- internal/querycoord/meta.go | 2 +- internal/querycoord/query_coord.go | 2 +- internal/querycoord/query_coord_test.go | 2 +- internal/querycoord/segment_allocator_test.go | 2 +- internal/querynode/data_sync_service.go | 2 +- .../querynode/flow_graph_delete_node_test.go | 2 +- .../flow_graph_filter_delete_node.go | 2 +- .../flow_graph_filter_delete_node_test.go | 2 +- .../querynode/flow_graph_filter_dm_node.go | 2 +- .../flow_graph_filter_dm_node_test.go | 2 +- internal/querynode/flow_graph_insert_node.go | 2 +- .../querynode/flow_graph_insert_node_test.go | 2 +- internal/querynode/flow_graph_message.go | 2 +- internal/querynode/flow_graph_message_test.go | 2 +- internal/querynode/flow_graph_query_node.go | 6 +- internal/querynode/mock_test.go | 2 +- internal/querynode/query_collection.go | 2 +- internal/querynode/query_collection_test.go | 2 +- internal/querynode/query_node.go | 2 +- internal/querynode/query_node_test.go | 2 +- internal/querynode/query_service.go | 2 +- internal/querynode/query_service_test.go | 2 +- internal/querynode/segment_loader.go | 2 +- internal/querynode/stats_service.go | 2 +- internal/querynode/stats_service_test.go | 2 +- internal/querynode/streaming.go | 2 +- internal/querynode/task_test.go | 2 +- internal/rootcoord/dml_channels.go | 2 +- internal/rootcoord/dml_channels_test.go | 6 +- internal/rootcoord/root_coord.go | 2 +- internal/rootcoord/root_coord_test.go | 2 +- internal/rootcoord/timestamp_test.go | 2 +- internal/rootcoord/timeticksync.go | 2 +- internal/rootcoord/timeticksync_test.go | 2 +- internal/rootcoord/util.go | 2 +- internal/rootcoord/util_test.go | 2 +- internal/util/flowgraph/input_node.go | 2 +- internal/util/flowgraph/input_node_test.go | 2 +- internal/util/flowgraph/message.go | 2 +- internal/util/flowgraph/message_test.go | 2 +- internal/util/flowgraph/node_test.go | 2 +- internal/util/paramtable/base_table.go | 170 ++++++++++-------- scripts/run_go_unittest.sh | 4 +- 145 files changed, 582 insertions(+), 508 deletions(-) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/client.go (87%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/client_impl.go (97%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/client_impl_test.go (95%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/consumer.go (95%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/consumer_impl.go (99%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/consumer_impl_test.go (97%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/error.go (98%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/producer.go (98%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/producer_impl.go (95%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/producer_impl_test.go (98%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/reader.go (99%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/reader_impl.go (99%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/reader_impl_test.go (98%) rename internal/{util/rocksmq/client/rocksmq => mq/mqimpl/rocksmq/client}/test_helper.go (86%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/global_rmq.go (99%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/global_rmq_test.go (99%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/rocksmq.go (99%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/rocksmq_impl.go (99%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/rocksmq_impl_test.go (99%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/rocksmq_reader.go (99%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/rocksmq_retention.go (99%) rename internal/{util/rocksmq/server/rocksmq => mq/mqimpl/rocksmq/server}/rocksmq_retention_test.go (99%) rename internal/{ => mq}/msgstream/OWNERS (100%) rename internal/{ => mq}/msgstream/mq_factory.go (85%) rename internal/{ => mq}/msgstream/mq_factory_test.go (100%) rename internal/{ => mq}/msgstream/mq_msgstream.go (91%) rename internal/{ => mq}/msgstream/mq_msgstream_test.go (94%) rename internal/{util/mqclient => mq/msgstream/mqwrapper}/client.go (98%) rename internal/{util/mqclient => mq/msgstream/mqwrapper}/consumer.go (99%) rename internal/{util/mqclient => mq/msgstream/mqwrapper}/id.go (98%) rename internal/{util/mqclient => mq/msgstream/mqwrapper}/message.go (98%) rename internal/{util/mqclient => mq/msgstream/mqwrapper}/producer.go (98%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_client.go (81%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_client_test.go (88%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_consumer.go (90%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_consumer_test.go (81%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_id.go (84%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_id_test.go (93%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_message.go (87%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_producer.go (83%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_producer_test.go (84%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_reader.go (87%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/pulsar}/pulsar_reader_test.go (83%) rename internal/{util/mqclient => mq/msgstream/mqwrapper}/reader.go (99%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_client.go (64%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_client_test.go (76%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_consumer.go (76%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_id.go (89%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_id_test.go (98%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_message.go (82%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_producer.go (73%) rename internal/{util/mqclient => mq/msgstream/mqwrapper/rmq}/rmq_reader.go (67%) rename internal/{ => mq}/msgstream/msg.go (100%) rename internal/{ => mq}/msgstream/msg_test.go (100%) rename internal/{ => mq}/msgstream/msgstream.go (95%) rename internal/{ => mq}/msgstream/repack_func.go (100%) rename internal/{ => mq}/msgstream/trace.go (100%) rename internal/{ => mq}/msgstream/unmarshal.go (100%) rename internal/{ => mq}/msgstream/unmarshal_test.go (100%) diff --git a/cmd/components/data_coord.go b/cmd/components/data_coord.go index f223f39bc5..c055817a7b 100644 --- a/cmd/components/data_coord.go +++ b/cmd/components/data_coord.go @@ -23,7 +23,7 @@ import ( "github.com/milvus-io/milvus/internal/proto/internalpb" grpcdatacoordclient "github.com/milvus-io/milvus/internal/distributed/datacoord" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // DataCoord implements grpc server of DataCoord server diff --git a/cmd/components/data_node.go b/cmd/components/data_node.go index 05e8956fbf..3c8dc595cc 100644 --- a/cmd/components/data_node.go +++ b/cmd/components/data_node.go @@ -23,7 +23,7 @@ import ( grpcdatanode "github.com/milvus-io/milvus/internal/distributed/datanode" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // DataNode implements DataNode grpc server diff --git a/cmd/components/proxy.go b/cmd/components/proxy.go index 75b3bae390..2b192435ca 100644 --- a/cmd/components/proxy.go +++ b/cmd/components/proxy.go @@ -23,7 +23,7 @@ import ( "github.com/milvus-io/milvus/internal/proto/internalpb" grpcproxy "github.com/milvus-io/milvus/internal/distributed/proxy" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // Proxy implements Proxy grpc server diff --git a/cmd/components/query_coord.go b/cmd/components/query_coord.go index 9eef565a3b..8c86e5c877 100644 --- a/cmd/components/query_coord.go +++ b/cmd/components/query_coord.go @@ -23,7 +23,7 @@ import ( "github.com/milvus-io/milvus/internal/proto/internalpb" grpcquerycoord "github.com/milvus-io/milvus/internal/distributed/querycoord" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // QueryCoord implements QueryCoord grpc server diff --git a/cmd/components/query_node.go b/cmd/components/query_node.go index dc02e0fdf8..4af426f6b4 100644 --- a/cmd/components/query_node.go +++ b/cmd/components/query_node.go @@ -23,7 +23,7 @@ import ( "github.com/milvus-io/milvus/internal/proto/internalpb" grpcquerynode "github.com/milvus-io/milvus/internal/distributed/querynode" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // QueryNode implements QueryNode grpc server diff --git a/cmd/components/root_coord.go b/cmd/components/root_coord.go index 1b586842a1..78aae63bcd 100644 --- a/cmd/components/root_coord.go +++ b/cmd/components/root_coord.go @@ -24,7 +24,7 @@ import ( "github.com/milvus-io/milvus/internal/proto/internalpb" rc "github.com/milvus-io/milvus/internal/distributed/rootcoord" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/opentracing/opentracing-go" ) diff --git a/cmd/roles/roles.go b/cmd/roles/roles.go index d132286338..99c9216df3 100644 --- a/cmd/roles/roles.go +++ b/cmd/roles/roles.go @@ -27,6 +27,8 @@ import ( "syscall" "time" + rocksmqimpl "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" + "go.uber.org/zap" "github.com/milvus-io/milvus/cmd/components" @@ -36,7 +38,7 @@ import ( "github.com/milvus-io/milvus/internal/indexnode" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proxy" "github.com/milvus-io/milvus/internal/querycoord" @@ -47,7 +49,6 @@ import ( logutil "github.com/milvus-io/milvus/internal/util/logutil" "github.com/milvus-io/milvus/internal/util/metricsinfo" "github.com/milvus-io/milvus/internal/util/paramtable" - "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" "github.com/milvus-io/milvus/internal/util/trace" "github.com/milvus-io/milvus/internal/util/typeutil" ) @@ -62,12 +63,12 @@ func newMsgFactory(localMsg bool) msgstream.Factory { } func initRocksmq() error { - err := rocksmq.InitRocksMQ() + err := rocksmqimpl.InitRocksMQ() return err } func stopRocksmq() { - rocksmq.CloseRocksMQ() + rocksmqimpl.CloseRocksMQ() } // MilvusRoles decides which components are brought up with Milvus. diff --git a/internal/datacoord/channel_manager.go b/internal/datacoord/channel_manager.go index 00408318eb..365a8bc020 100644 --- a/internal/datacoord/channel_manager.go +++ b/internal/datacoord/channel_manager.go @@ -24,7 +24,7 @@ import ( "github.com/milvus-io/milvus/internal/kv" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/datapb" "go.uber.org/zap" "stathat.com/c/consistent" diff --git a/internal/datacoord/server.go b/internal/datacoord/server.go index b1a3614665..34192361f4 100644 --- a/internal/datacoord/server.go +++ b/internal/datacoord/server.go @@ -30,14 +30,14 @@ import ( rootcoordclient "github.com/milvus-io/milvus/internal/distributed/rootcoord/client" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/milvuspb" "github.com/milvus-io/milvus/internal/types" "github.com/milvus-io/milvus/internal/util/logutil" "github.com/milvus-io/milvus/internal/util/metricsinfo" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/milvus-io/milvus/internal/util/paramtable" "github.com/milvus-io/milvus/internal/util/retry" "github.com/milvus-io/milvus/internal/util/sessionutil" @@ -449,7 +449,7 @@ func (s *Server) startDataNodeTtLoop(ctx context.Context) { return } ttMsgStream.AsConsumerWithPosition([]string{Params.MsgChannelCfg.DataCoordTimeTick}, - Params.MsgChannelCfg.DataCoordSubName, mqclient.SubscriptionPositionLatest) + Params.MsgChannelCfg.DataCoordSubName, mqwrapper.SubscriptionPositionLatest) log.Info("DataCoord creates the timetick channel consumer", zap.String("timeTickChannel", Params.MsgChannelCfg.DataCoordTimeTick), zap.String("subscription", Params.MsgChannelCfg.DataCoordSubName)) diff --git a/internal/datacoord/server_test.go b/internal/datacoord/server_test.go index 7866030ce5..5bec2e4e0d 100644 --- a/internal/datacoord/server_test.go +++ b/internal/datacoord/server_test.go @@ -33,7 +33,7 @@ import ( "github.com/milvus-io/milvus/internal/common" memkv "github.com/milvus-io/milvus/internal/kv/mem" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/datanode/data_node.go b/internal/datanode/data_node.go index a6ed6bf5a0..24ba7a8e70 100644 --- a/internal/datanode/data_node.go +++ b/internal/datanode/data_node.go @@ -40,7 +40,7 @@ import ( miniokv "github.com/milvus-io/milvus/internal/kv/minio" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/datanode/data_node_test.go b/internal/datanode/data_node_test.go index 9ec512894a..bc8cf34cc6 100644 --- a/internal/datanode/data_node_test.go +++ b/internal/datanode/data_node_test.go @@ -31,7 +31,7 @@ import ( "github.com/milvus-io/milvus/internal/common" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/types" "github.com/milvus-io/milvus/internal/proto/commonpb" diff --git a/internal/datanode/data_sync_service.go b/internal/datanode/data_sync_service.go index 6d289c4e09..7ffb0f426d 100644 --- a/internal/datanode/data_sync_service.go +++ b/internal/datanode/data_sync_service.go @@ -24,7 +24,7 @@ import ( "github.com/milvus-io/milvus/internal/kv" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/types" "github.com/milvus-io/milvus/internal/util/flowgraph" diff --git a/internal/datanode/data_sync_service_test.go b/internal/datanode/data_sync_service_test.go index a8acc0b2ec..bcaa21d758 100644 --- a/internal/datanode/data_sync_service_test.go +++ b/internal/datanode/data_sync_service_test.go @@ -30,7 +30,7 @@ import ( "github.com/milvus-io/milvus/internal/common" memkv "github.com/milvus-io/milvus/internal/kv/mem" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/datanode/flow_graph_dd_node.go b/internal/datanode/flow_graph_dd_node.go index b0dbe4980b..0423524d9d 100644 --- a/internal/datanode/flow_graph_dd_node.go +++ b/internal/datanode/flow_graph_dd_node.go @@ -24,7 +24,7 @@ import ( "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/datanode/flow_graph_dd_node_test.go b/internal/datanode/flow_graph_dd_node_test.go index 636fc5d05a..9a15a7ea9f 100644 --- a/internal/datanode/flow_graph_dd_node_test.go +++ b/internal/datanode/flow_graph_dd_node_test.go @@ -24,7 +24,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/datanode/flow_graph_delete_node.go b/internal/datanode/flow_graph_delete_node.go index 214e5d23b9..90e230eea1 100644 --- a/internal/datanode/flow_graph_delete_node.go +++ b/internal/datanode/flow_graph_delete_node.go @@ -27,7 +27,7 @@ import ( "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/storage" "github.com/milvus-io/milvus/internal/util/trace" diff --git a/internal/datanode/flow_graph_delete_node_test.go b/internal/datanode/flow_graph_delete_node_test.go index 42cd8bc627..56b36ee511 100644 --- a/internal/datanode/flow_graph_delete_node_test.go +++ b/internal/datanode/flow_graph_delete_node_test.go @@ -25,7 +25,7 @@ import ( "github.com/bits-and-blooms/bloom/v3" "github.com/milvus-io/milvus/internal/common" memkv "github.com/milvus-io/milvus/internal/kv/mem" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/schemapb" "github.com/milvus-io/milvus/internal/util/flowgraph" "github.com/stretchr/testify/assert" diff --git a/internal/datanode/flow_graph_dmstream_input_node_test.go b/internal/datanode/flow_graph_dmstream_input_node_test.go index bc14b02c1b..39640eec2e 100644 --- a/internal/datanode/flow_graph_dmstream_input_node_test.go +++ b/internal/datanode/flow_graph_dmstream_input_node_test.go @@ -21,9 +21,9 @@ import ( "errors" "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/stretchr/testify/assert" ) @@ -69,7 +69,7 @@ func (mtm *mockTtMsgStream) Chan() <-chan *msgstream.MsgPack { func (mtm *mockTtMsgStream) AsProducer(channels []string) {} func (mtm *mockTtMsgStream) AsReader(channels []string, subName string) {} func (mtm *mockTtMsgStream) AsConsumer(channels []string, subName string) {} -func (mtm *mockTtMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) { +func (mtm *mockTtMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqwrapper.SubscriptionInitialPosition) { } func (mtm *mockTtMsgStream) SetRepackFunc(repackFunc msgstream.RepackFunc) {} func (mtm *mockTtMsgStream) ComputeProduceChannelIndexes(tsMsgs []msgstream.TsMsg) [][]int32 { diff --git a/internal/datanode/flow_graph_insert_buffer_node.go b/internal/datanode/flow_graph_insert_buffer_node.go index db41adc938..eea70e6fec 100644 --- a/internal/datanode/flow_graph_insert_buffer_node.go +++ b/internal/datanode/flow_graph_insert_buffer_node.go @@ -34,7 +34,7 @@ import ( "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/datanode/flow_graph_insert_buffer_node_test.go b/internal/datanode/flow_graph_insert_buffer_node_test.go index 333991097a..5610fba035 100644 --- a/internal/datanode/flow_graph_insert_buffer_node_test.go +++ b/internal/datanode/flow_graph_insert_buffer_node_test.go @@ -27,7 +27,7 @@ import ( "github.com/stretchr/testify/require" memkv "github.com/milvus-io/milvus/internal/kv/mem" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/types" "github.com/milvus-io/milvus/internal/util/flowgraph" diff --git a/internal/datanode/flow_graph_message.go b/internal/datanode/flow_graph_message.go index 5221c077a9..fb2f2d6b76 100644 --- a/internal/datanode/flow_graph_message.go +++ b/internal/datanode/flow_graph_message.go @@ -17,7 +17,7 @@ package datanode import ( - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/util/flowgraph" ) diff --git a/internal/datanode/mock_test.go b/internal/datanode/mock_test.go index e6e1fdf26b..84ded2e0a4 100644 --- a/internal/datanode/mock_test.go +++ b/internal/datanode/mock_test.go @@ -30,7 +30,7 @@ import ( etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" s "github.com/milvus-io/milvus/internal/storage" "github.com/milvus-io/milvus/internal/types" "github.com/milvus-io/milvus/internal/util/tsoutil" diff --git a/internal/distributed/datacoord/service.go b/internal/distributed/datacoord/service.go index 567ee63894..b800884be8 100644 --- a/internal/distributed/datacoord/service.go +++ b/internal/distributed/datacoord/service.go @@ -29,7 +29,7 @@ import ( grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus" "github.com/milvus-io/milvus/internal/datacoord" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/distributed/datanode/service.go b/internal/distributed/datanode/service.go index f0ed2275a7..d8ef54e854 100644 --- a/internal/distributed/datanode/service.go +++ b/internal/distributed/datanode/service.go @@ -31,7 +31,7 @@ import ( dcc "github.com/milvus-io/milvus/internal/distributed/datacoord/client" rcc "github.com/milvus-io/milvus/internal/distributed/rootcoord/client" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/distributed/proxy/service.go b/internal/distributed/proxy/service.go index 65becfad4d..1679d7e2aa 100644 --- a/internal/distributed/proxy/service.go +++ b/internal/distributed/proxy/service.go @@ -34,7 +34,7 @@ import ( qcc "github.com/milvus-io/milvus/internal/distributed/querycoord/client" rcc "github.com/milvus-io/milvus/internal/distributed/rootcoord/client" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" diff --git a/internal/distributed/querycoord/service.go b/internal/distributed/querycoord/service.go index 152aa9423e..1bd853a6af 100644 --- a/internal/distributed/querycoord/service.go +++ b/internal/distributed/querycoord/service.go @@ -29,7 +29,7 @@ import ( icc "github.com/milvus-io/milvus/internal/distributed/indexcoord/client" rcc "github.com/milvus-io/milvus/internal/distributed/rootcoord/client" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" diff --git a/internal/distributed/querynode/service.go b/internal/distributed/querynode/service.go index ddf296710e..0fbd75736f 100644 --- a/internal/distributed/querynode/service.go +++ b/internal/distributed/querynode/service.go @@ -27,7 +27,7 @@ import ( ot "github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" diff --git a/internal/distributed/rootcoord/service.go b/internal/distributed/rootcoord/service.go index 1bbdf16e70..b5bf75af5e 100644 --- a/internal/distributed/rootcoord/service.go +++ b/internal/distributed/rootcoord/service.go @@ -27,7 +27,7 @@ import ( ot "github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing" pnc "github.com/milvus-io/milvus/internal/distributed/proxy/client" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/distributed/rootcoord/service_test.go b/internal/distributed/rootcoord/service_test.go index 48ac14dd15..ad11ee09e2 100644 --- a/internal/distributed/rootcoord/service_test.go +++ b/internal/distributed/rootcoord/service_test.go @@ -30,7 +30,7 @@ import ( "github.com/golang/protobuf/proto" rcc "github.com/milvus-io/milvus/internal/distributed/rootcoord/client" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/etcdpb" diff --git a/internal/util/rocksmq/client/rocksmq/client.go b/internal/mq/mqimpl/rocksmq/client/client.go similarity index 87% rename from internal/util/rocksmq/client/rocksmq/client.go rename to internal/mq/mqimpl/rocksmq/client/client.go index 7046acca3b..097e8d5bc5 100644 --- a/internal/util/rocksmq/client/rocksmq/client.go +++ b/internal/mq/mqimpl/rocksmq/client/client.go @@ -9,17 +9,17 @@ // 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 rocksmq +package client import ( - server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" ) // RocksMQ is the type server.RocksMQ type RocksMQ = server.RocksMQ // NewClient returns a rocksmq client -func NewClient(options ClientOptions) (Client, error) { +func NewClient(options Options) (Client, error) { if options.Server == nil { options.Server = server.Rmq } @@ -27,7 +27,7 @@ func NewClient(options ClientOptions) (Client, error) { } // ClientOptions is the options of a client -type ClientOptions struct { +type Options struct { Server RocksMQ } diff --git a/internal/util/rocksmq/client/rocksmq/client_impl.go b/internal/mq/mqimpl/rocksmq/client/client_impl.go similarity index 97% rename from internal/util/rocksmq/client/rocksmq/client_impl.go rename to internal/mq/mqimpl/rocksmq/client/client_impl.go index 3e457a801c..2b84313d76 100644 --- a/internal/util/rocksmq/client/rocksmq/client_impl.go +++ b/internal/mq/mqimpl/rocksmq/client/client_impl.go @@ -9,15 +9,15 @@ // 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 rocksmq +package client import ( "reflect" "sync" - "github.com/milvus-io/milvus/internal/log" - server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" + "github.com/milvus-io/milvus/internal/log" "go.uber.org/zap" ) @@ -30,7 +30,8 @@ type client struct { closeOnce sync.Once } -func newClient(options ClientOptions) (*client, error) { +func newClient(options Options) (*client, error) { + if options.Server == nil { return nil, newError(InvalidConfiguration, "options.Server is nil") } @@ -48,6 +49,7 @@ func newClient(options ClientOptions) (*client, error) { func (c *client) CreateProducer(options ProducerOptions) (Producer, error) { // Create a producer producer, err := newProducer(c, options) + if err != nil { return nil, err } diff --git a/internal/util/rocksmq/client/rocksmq/client_impl_test.go b/internal/mq/mqimpl/rocksmq/client/client_impl_test.go similarity index 95% rename from internal/util/rocksmq/client/rocksmq/client_impl_test.go rename to internal/mq/mqimpl/rocksmq/client/client_impl_test.go index 30a8777ca0..1af179c5f1 100644 --- a/internal/util/rocksmq/client/rocksmq/client_impl_test.go +++ b/internal/mq/mqimpl/rocksmq/client/client_impl_test.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "os" @@ -22,7 +22,7 @@ import ( var rmqPath = "/tmp/rocksmq_client" func TestClient(t *testing.T) { - client, err := NewClient(ClientOptions{}) + client, err := NewClient(Options{}) assert.NotNil(t, client) assert.Nil(t, err) } @@ -34,7 +34,7 @@ func TestClient_CreateProducer(t *testing.T) { assert.Error(t, err) ///////////////////////////////////////////////// - client, err := NewClient(ClientOptions{ + client, err := NewClient(Options{ Server: newMockRocksMQ(), }) assert.NoError(t, err) @@ -50,7 +50,7 @@ func TestClient_CreateProducer(t *testing.T) { rmqPathTest := rmqPath + "/test_client1" rmq := newRocksMQ(t, rmqPathTest) defer removePath(rmqPath) - client1, err := NewClient(ClientOptions{ + client1, err := NewClient(Options{ Server: rmq, }) assert.NoError(t, err) @@ -72,7 +72,7 @@ func TestClient_CreateProducer(t *testing.T) { } func TestClient_Subscribe(t *testing.T) { - client, err := NewClient(ClientOptions{ + client, err := NewClient(Options{ Server: newMockRocksMQ(), }) assert.NoError(t, err) @@ -90,7 +90,7 @@ func TestClient_Subscribe(t *testing.T) { rmqPathTest := rmqPath + "/test_client2" rmq := newRocksMQ(t, rmqPathTest) defer removePath(rmqPath) - client1, err := NewClient(ClientOptions{ + client1, err := NewClient(Options{ Server: rmq, }) assert.NoError(t, err) @@ -131,7 +131,7 @@ func TestClient_SeekLatest(t *testing.T) { rmqPathTest := rmqPath + "/seekLatest" rmq := newRocksMQ(t, rmqPathTest) defer removePath(rmqPath) - client, err := NewClient(ClientOptions{ + client, err := NewClient(Options{ Server: rmq, }) assert.NoError(t, err) @@ -202,7 +202,7 @@ func TestClient_consume(t *testing.T) { rmqPathTest := rmqPath + "/test_client3" rmq := newRocksMQ(t, rmqPathTest) defer removePath(rmqPath) - client, err := NewClient(ClientOptions{ + client, err := NewClient(Options{ Server: rmq, }) assert.NoError(t, err) diff --git a/internal/util/rocksmq/client/rocksmq/consumer.go b/internal/mq/mqimpl/rocksmq/client/consumer.go similarity index 95% rename from internal/util/rocksmq/client/rocksmq/consumer.go rename to internal/mq/mqimpl/rocksmq/client/consumer.go index 0467637ccb..16fdcfa4c5 100644 --- a/internal/util/rocksmq/client/rocksmq/consumer.go +++ b/internal/mq/mqimpl/rocksmq/client/consumer.go @@ -9,9 +9,11 @@ // 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 rocksmq +package client -import server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" +import ( + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" +) // SubscriptionInitialPosition is the initial subscription position type SubscriptionInitialPosition int diff --git a/internal/util/rocksmq/client/rocksmq/consumer_impl.go b/internal/mq/mqimpl/rocksmq/client/consumer_impl.go similarity index 99% rename from internal/util/rocksmq/client/rocksmq/consumer_impl.go rename to internal/mq/mqimpl/rocksmq/client/consumer_impl.go index d6fbe1116b..0d8bf9a7a8 100644 --- a/internal/util/rocksmq/client/rocksmq/consumer_impl.go +++ b/internal/mq/mqimpl/rocksmq/client/consumer_impl.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "sync" diff --git a/internal/util/rocksmq/client/rocksmq/consumer_impl_test.go b/internal/mq/mqimpl/rocksmq/client/consumer_impl_test.go similarity index 97% rename from internal/util/rocksmq/client/rocksmq/consumer_impl_test.go rename to internal/mq/mqimpl/rocksmq/client/consumer_impl_test.go index 9d83a9951f..ffdedb3c03 100644 --- a/internal/util/rocksmq/client/rocksmq/consumer_impl_test.go +++ b/internal/mq/mqimpl/rocksmq/client/consumer_impl_test.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "os" @@ -52,7 +52,7 @@ func TestConsumer_newConsumer(t *testing.T) { rmqPathTest := rmqPath + "/test_consumer1" rmq := newRocksMQ(t, rmqPathTest) defer removePath(rmqPath) - client, err := newClient(ClientOptions{ + client, err := newClient(Options{ Server: rmq, }) assert.NoError(t, err) @@ -123,7 +123,7 @@ func TestConsumer_Seek(t *testing.T) { rmqPathTest := rmqPath + "/test_consumer2" rmq := newRocksMQ(t, rmqPathTest) defer removePath(rmqPath) - client, err := newClient(ClientOptions{ + client, err := newClient(Options{ Server: rmq, }) assert.NoError(t, err) diff --git a/internal/util/rocksmq/client/rocksmq/error.go b/internal/mq/mqimpl/rocksmq/client/error.go similarity index 98% rename from internal/util/rocksmq/client/rocksmq/error.go rename to internal/mq/mqimpl/rocksmq/client/error.go index bf8caf9ad7..d1391bb208 100644 --- a/internal/util/rocksmq/client/rocksmq/error.go +++ b/internal/mq/mqimpl/rocksmq/client/error.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import "fmt" diff --git a/internal/util/rocksmq/client/rocksmq/producer.go b/internal/mq/mqimpl/rocksmq/client/producer.go similarity index 98% rename from internal/util/rocksmq/client/rocksmq/producer.go rename to internal/mq/mqimpl/rocksmq/client/producer.go index 01e48c8ddd..93c1a28faf 100644 --- a/internal/util/rocksmq/client/rocksmq/producer.go +++ b/internal/mq/mqimpl/rocksmq/client/producer.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client // ProducerOptions is the options of a producer type ProducerOptions struct { diff --git a/internal/util/rocksmq/client/rocksmq/producer_impl.go b/internal/mq/mqimpl/rocksmq/client/producer_impl.go similarity index 95% rename from internal/util/rocksmq/client/rocksmq/producer_impl.go rename to internal/mq/mqimpl/rocksmq/client/producer_impl.go index 981a32d403..4c59e9e2f6 100644 --- a/internal/util/rocksmq/client/rocksmq/producer_impl.go +++ b/internal/mq/mqimpl/rocksmq/client/producer_impl.go @@ -9,12 +9,11 @@ // 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 rocksmq +package client import ( "github.com/milvus-io/milvus/internal/log" - server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" - + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" "go.uber.org/zap" ) diff --git a/internal/util/rocksmq/client/rocksmq/producer_impl_test.go b/internal/mq/mqimpl/rocksmq/client/producer_impl_test.go similarity index 98% rename from internal/util/rocksmq/client/rocksmq/producer_impl_test.go rename to internal/mq/mqimpl/rocksmq/client/producer_impl_test.go index f53af07062..36a2aa79a0 100644 --- a/internal/util/rocksmq/client/rocksmq/producer_impl_test.go +++ b/internal/mq/mqimpl/rocksmq/client/producer_impl_test.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "testing" diff --git a/internal/util/rocksmq/client/rocksmq/reader.go b/internal/mq/mqimpl/rocksmq/client/reader.go similarity index 99% rename from internal/util/rocksmq/client/rocksmq/reader.go rename to internal/mq/mqimpl/rocksmq/client/reader.go index c000e263f2..5ae6825f69 100644 --- a/internal/util/rocksmq/client/rocksmq/reader.go +++ b/internal/mq/mqimpl/rocksmq/client/reader.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "context" diff --git a/internal/util/rocksmq/client/rocksmq/reader_impl.go b/internal/mq/mqimpl/rocksmq/client/reader_impl.go similarity index 99% rename from internal/util/rocksmq/client/rocksmq/reader_impl.go rename to internal/mq/mqimpl/rocksmq/client/reader_impl.go index 8bb2507ec1..0e5e6c75f7 100644 --- a/internal/util/rocksmq/client/rocksmq/reader_impl.go +++ b/internal/mq/mqimpl/rocksmq/client/reader_impl.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "context" diff --git a/internal/util/rocksmq/client/rocksmq/reader_impl_test.go b/internal/mq/mqimpl/rocksmq/client/reader_impl_test.go similarity index 98% rename from internal/util/rocksmq/client/rocksmq/reader_impl_test.go rename to internal/mq/mqimpl/rocksmq/client/reader_impl_test.go index 4dc4b60527..9658e4a899 100644 --- a/internal/util/rocksmq/client/rocksmq/reader_impl_test.go +++ b/internal/mq/mqimpl/rocksmq/client/reader_impl_test.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "context" @@ -45,7 +45,7 @@ func TestReader_Next(t *testing.T) { rmqPathTest := rmqPath + "/test_reader" rmq := newRocksMQ(t, rmqPathTest) defer removePath(rmqPath) - client, err := newClient(ClientOptions{ + client, err := newClient(Options{ Server: rmq, }) assert.NoError(t, err) diff --git a/internal/util/rocksmq/client/rocksmq/test_helper.go b/internal/mq/mqimpl/rocksmq/client/test_helper.go similarity index 86% rename from internal/util/rocksmq/client/rocksmq/test_helper.go rename to internal/mq/mqimpl/rocksmq/client/test_helper.go index 6c708f0bf0..5bc036fa39 100644 --- a/internal/util/rocksmq/client/rocksmq/test_helper.go +++ b/internal/mq/mqimpl/rocksmq/client/test_helper.go @@ -9,7 +9,7 @@ // 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 rocksmq +package client import ( "fmt" @@ -17,9 +17,9 @@ import ( "testing" "time" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" + "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" - server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" "github.com/stretchr/testify/assert" "go.uber.org/zap" @@ -39,7 +39,7 @@ func newMockRocksMQ() server.RocksMQ { } func newMockClient() *client { - client, _ := newClient(ClientOptions{ + client, _ := newClient(Options{ Server: newMockRocksMQ(), }) return client @@ -47,7 +47,7 @@ func newMockClient() *client { func newRocksMQ(t *testing.T, rmqPath string) server.RocksMQ { rocksdbPath := rmqPath - rmq, err := rocksmq.NewRocksMQ(rocksdbPath, nil) + rmq, err := server.NewRocksMQ(rocksdbPath, nil) assert.NoError(t, err) return rmq } diff --git a/internal/util/rocksmq/server/rocksmq/global_rmq.go b/internal/mq/mqimpl/rocksmq/server/global_rmq.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/global_rmq.go rename to internal/mq/mqimpl/rocksmq/server/global_rmq.go index fba0ff38e7..3670e5b935 100644 --- a/internal/util/rocksmq/server/rocksmq/global_rmq.go +++ b/internal/mq/mqimpl/rocksmq/server/global_rmq.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package rocksmq +package server import ( "errors" diff --git a/internal/util/rocksmq/server/rocksmq/global_rmq_test.go b/internal/mq/mqimpl/rocksmq/server/global_rmq_test.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/global_rmq_test.go rename to internal/mq/mqimpl/rocksmq/server/global_rmq_test.go index 71b5ccfcb0..a447c052a9 100644 --- a/internal/util/rocksmq/server/rocksmq/global_rmq_test.go +++ b/internal/mq/mqimpl/rocksmq/server/global_rmq_test.go @@ -9,7 +9,7 @@ // 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 rocksmq +package server import ( "log" diff --git a/internal/util/rocksmq/server/rocksmq/rocksmq.go b/internal/mq/mqimpl/rocksmq/server/rocksmq.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/rocksmq.go rename to internal/mq/mqimpl/rocksmq/server/rocksmq.go index b506853031..5f5d1b6601 100644 --- a/internal/util/rocksmq/server/rocksmq/rocksmq.go +++ b/internal/mq/mqimpl/rocksmq/server/rocksmq.go @@ -9,7 +9,7 @@ // 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 rocksmq +package server import "context" diff --git a/internal/util/rocksmq/server/rocksmq/rocksmq_impl.go b/internal/mq/mqimpl/rocksmq/server/rocksmq_impl.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/rocksmq_impl.go rename to internal/mq/mqimpl/rocksmq/server/rocksmq_impl.go index 974c4f9c49..0e0660fecf 100644 --- a/internal/util/rocksmq/server/rocksmq/rocksmq_impl.go +++ b/internal/mq/mqimpl/rocksmq/server/rocksmq_impl.go @@ -9,7 +9,7 @@ // 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 rocksmq +package server import ( "context" diff --git a/internal/util/rocksmq/server/rocksmq/rocksmq_impl_test.go b/internal/mq/mqimpl/rocksmq/server/rocksmq_impl_test.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/rocksmq_impl_test.go rename to internal/mq/mqimpl/rocksmq/server/rocksmq_impl_test.go index 927927a707..02d7691cd4 100644 --- a/internal/util/rocksmq/server/rocksmq/rocksmq_impl_test.go +++ b/internal/mq/mqimpl/rocksmq/server/rocksmq_impl_test.go @@ -9,7 +9,7 @@ // 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 rocksmq +package server import ( "context" diff --git a/internal/util/rocksmq/server/rocksmq/rocksmq_reader.go b/internal/mq/mqimpl/rocksmq/server/rocksmq_reader.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/rocksmq_reader.go rename to internal/mq/mqimpl/rocksmq/server/rocksmq_reader.go index 04cd0cf48b..ec213cb094 100644 --- a/internal/util/rocksmq/server/rocksmq/rocksmq_reader.go +++ b/internal/mq/mqimpl/rocksmq/server/rocksmq_reader.go @@ -9,7 +9,7 @@ // 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 rocksmq +package server import ( "context" diff --git a/internal/util/rocksmq/server/rocksmq/rocksmq_retention.go b/internal/mq/mqimpl/rocksmq/server/rocksmq_retention.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/rocksmq_retention.go rename to internal/mq/mqimpl/rocksmq/server/rocksmq_retention.go index 12e3670045..1cbc0a8c62 100644 --- a/internal/util/rocksmq/server/rocksmq/rocksmq_retention.go +++ b/internal/mq/mqimpl/rocksmq/server/rocksmq_retention.go @@ -9,7 +9,7 @@ // 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 rocksmq +package server import ( "fmt" diff --git a/internal/util/rocksmq/server/rocksmq/rocksmq_retention_test.go b/internal/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go similarity index 99% rename from internal/util/rocksmq/server/rocksmq/rocksmq_retention_test.go rename to internal/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go index fff512ae28..de9dc873ac 100644 --- a/internal/util/rocksmq/server/rocksmq/rocksmq_retention_test.go +++ b/internal/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go @@ -9,7 +9,7 @@ // 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 rocksmq +package server import ( "os" diff --git a/internal/msgstream/OWNERS b/internal/mq/msgstream/OWNERS similarity index 100% rename from internal/msgstream/OWNERS rename to internal/mq/msgstream/OWNERS diff --git a/internal/msgstream/mq_factory.go b/internal/mq/msgstream/mq_factory.go similarity index 85% rename from internal/msgstream/mq_factory.go rename to internal/mq/msgstream/mq_factory.go index 356258de8d..f0e34b2086 100644 --- a/internal/msgstream/mq_factory.go +++ b/internal/mq/msgstream/mq_factory.go @@ -19,12 +19,12 @@ package msgstream import ( "context" - "github.com/apache/pulsar-client-go/pulsar" - "github.com/mitchellh/mapstructure" + rmqimplserver "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" - "github.com/milvus-io/milvus/internal/util/mqclient" - "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" - rocksmqserver "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" + "github.com/apache/pulsar-client-go/pulsar" + puslarmqwrapper "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/pulsar" + rmqwrapper "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/rmq" + "github.com/mitchellh/mapstructure" ) // PmsFactory is a pulsar msgstream factory that implemented Factory interface(msgstream.go) @@ -47,7 +47,7 @@ func (f *PmsFactory) SetParams(params map[string]interface{}) error { // NewMsgStream is used to generate a new Msgstream object func (f *PmsFactory) NewMsgStream(ctx context.Context) (MsgStream, error) { - pulsarClient, err := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: f.PulsarAddress}) + pulsarClient, err := puslarmqwrapper.NewClient(pulsar.ClientOptions{URL: f.PulsarAddress}) if err != nil { return nil, err } @@ -56,7 +56,7 @@ func (f *PmsFactory) NewMsgStream(ctx context.Context) (MsgStream, error) { // NewTtMsgStream is used to generate a new TtMsgstream object func (f *PmsFactory) NewTtMsgStream(ctx context.Context) (MsgStream, error) { - pulsarClient, err := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: f.PulsarAddress}) + pulsarClient, err := puslarmqwrapper.NewClient(pulsar.ClientOptions{URL: f.PulsarAddress}) if err != nil { return nil, err } @@ -97,7 +97,7 @@ func (f *RmsFactory) SetParams(params map[string]interface{}) error { // NewMsgStream is used to generate a new Msgstream object func (f *RmsFactory) NewMsgStream(ctx context.Context) (MsgStream, error) { - rmqClient, err := mqclient.NewRmqClient(rocksmq.ClientOptions{Server: rocksmqserver.Rmq}) + rmqClient, err := rmqwrapper.NewClientWithDefaultOptions() if err != nil { return nil, err } @@ -106,7 +106,7 @@ func (f *RmsFactory) NewMsgStream(ctx context.Context) (MsgStream, error) { // NewTtMsgStream is used to generate a new TtMsgstream object func (f *RmsFactory) NewTtMsgStream(ctx context.Context) (MsgStream, error) { - rmqClient, err := mqclient.NewRmqClient(rocksmq.ClientOptions{Server: rocksmqserver.Rmq}) + rmqClient, err := rmqwrapper.NewClientWithDefaultOptions() if err != nil { return nil, err } @@ -115,7 +115,7 @@ func (f *RmsFactory) NewTtMsgStream(ctx context.Context) (MsgStream, error) { // NewQueryMsgStream is used to generate a new QueryMsgstream object func (f *RmsFactory) NewQueryMsgStream(ctx context.Context) (MsgStream, error) { - rmqClient, err := mqclient.NewRmqClient(rocksmq.ClientOptions{Server: rocksmqserver.Rmq}) + rmqClient, err := rmqwrapper.NewClientWithDefaultOptions() if err != nil { return nil, err } @@ -130,6 +130,6 @@ func NewRmsFactory() Factory { RmqBufSize: 1024, } - rocksmqserver.InitRocksMQ() + rmqimplserver.InitRocksMQ() return f } diff --git a/internal/msgstream/mq_factory_test.go b/internal/mq/msgstream/mq_factory_test.go similarity index 100% rename from internal/msgstream/mq_factory_test.go rename to internal/mq/msgstream/mq_factory_test.go diff --git a/internal/msgstream/mq_msgstream.go b/internal/mq/msgstream/mq_msgstream.go similarity index 91% rename from internal/msgstream/mq_msgstream.go rename to internal/mq/msgstream/mq_msgstream.go index 52f82a2eec..0f0421f657 100644 --- a/internal/msgstream/mq_msgstream.go +++ b/internal/mq/msgstream/mq_msgstream.go @@ -28,9 +28,9 @@ import ( "go.uber.org/zap" "github.com/milvus-io/milvus/internal/log" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/milvus-io/milvus/internal/util/retry" "github.com/milvus-io/milvus/internal/util/trace" "github.com/opentracing/opentracing-go" @@ -40,12 +40,12 @@ var _ MsgStream = (*mqMsgStream)(nil) type mqMsgStream struct { ctx context.Context - client mqclient.Client - producers map[string]mqclient.Producer + client mqwrapper.Client + producers map[string]mqwrapper.Producer producerChannels []string - consumers map[string]mqclient.Consumer + consumers map[string]mqwrapper.Consumer consumerChannels []string - readers map[string]mqclient.Reader + readers map[string]mqwrapper.Reader readerChannels []string repackFunc RepackFunc unmarshal UnmarshalDispatcher @@ -62,13 +62,13 @@ type mqMsgStream struct { func NewMqMsgStream(ctx context.Context, receiveBufSize int64, bufSize int64, - client mqclient.Client, + client mqwrapper.Client, unmarshal UnmarshalDispatcher) (*mqMsgStream, error) { streamCtx, streamCancel := context.WithCancel(ctx) - producers := make(map[string]mqclient.Producer) - consumers := make(map[string]mqclient.Consumer) - readers := make(map[string]mqclient.Reader) + producers := make(map[string]mqwrapper.Producer) + consumers := make(map[string]mqwrapper.Consumer) + readers := make(map[string]mqwrapper.Reader) producerChannels := make([]string, 0) consumerChannels := make([]string, 0) readerChannels := make([]string, 0) @@ -104,7 +104,7 @@ func (ms *mqMsgStream) AsProducer(channels []string) { break } fn := func() error { - pp, err := ms.client.CreateProducer(mqclient.ProducerOptions{Topic: channel}) + pp, err := ms.client.CreateProducer(mqwrapper.ProducerOptions{Topic: channel}) if err != nil { return err } @@ -128,21 +128,21 @@ func (ms *mqMsgStream) AsProducer(channels []string) { // AsConsumer Create consumer to receive message from channels func (ms *mqMsgStream) AsConsumer(channels []string, subName string) { - ms.AsConsumerWithPosition(channels, subName, mqclient.SubscriptionPositionEarliest) + ms.AsConsumerWithPosition(channels, subName, mqwrapper.SubscriptionPositionEarliest) } // AsConsumerWithPosition Create consumer to receive message from channels, with initial position // if initial position is set to latest, last message in the channel is exclusive -func (ms *mqMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) { +func (ms *mqMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqwrapper.SubscriptionInitialPosition) { for _, channel := range channels { if _, ok := ms.consumers[channel]; ok { continue } fn := func() error { - pc, err := ms.client.Subscribe(mqclient.ConsumerOptions{ + pc, err := ms.client.Subscribe(mqwrapper.ConsumerOptions{ Topic: channel, SubscriptionName: subName, - Type: mqclient.Exclusive, + Type: mqwrapper.Exclusive, SubscriptionInitialPosition: position, BufSize: ms.bufSize, }) @@ -175,7 +175,7 @@ func (ms *mqMsgStream) AsReader(channels []string, subName string) { break } fn := func() error { - r, err := ms.client.CreateReader(mqclient.ReaderOptions{ + r, err := ms.client.CreateReader(mqwrapper.ReaderOptions{ Topic: channel, StartMessageID: ms.client.EarliestMessageID(), SubscriptionRolePrefix: subName, @@ -303,7 +303,7 @@ func (ms *mqMsgStream) Produce(msgPack *MsgPack) error { return err } - msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}} + msg := &mqwrapper.ProducerMessage{Payload: m, Properties: map[string]string{}} trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties) @@ -369,7 +369,7 @@ func (ms *mqMsgStream) ProduceMark(msgPack *MsgPack) (map[string][]MessageID, er return ids, err } - msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}} + msg := &mqwrapper.ProducerMessage{Payload: m, Properties: map[string]string{}} trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties) @@ -412,7 +412,7 @@ func (ms *mqMsgStream) Broadcast(msgPack *MsgPack) error { return err } - msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}} + msg := &mqwrapper.ProducerMessage{Payload: m, Properties: map[string]string{}} trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties) @@ -454,7 +454,7 @@ func (ms *mqMsgStream) BroadcastMark(msgPack *MsgPack) (map[string][]MessageID, return ids, err } - msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}} + msg := &mqwrapper.ProducerMessage{Payload: m, Properties: map[string]string{}} trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties) @@ -491,7 +491,7 @@ func (ms *mqMsgStream) Consume() *MsgPack { } } -func (ms *mqMsgStream) getTsMsgFromConsumerMsg(msg mqclient.Message) (TsMsg, error) { +func (ms *mqMsgStream) getTsMsgFromConsumerMsg(msg mqwrapper.Message) (TsMsg, error) { header := commonpb.MsgHeader{} if msg.Payload() == nil { return nil, fmt.Errorf("failed to unmarshal message header, payload is empty") @@ -517,7 +517,7 @@ func (ms *mqMsgStream) getTsMsgFromConsumerMsg(msg mqclient.Message) (TsMsg, err return tsMsg, nil } -func (ms *mqMsgStream) receiveMsg(consumer mqclient.Consumer) { +func (ms *mqMsgStream) receiveMsg(consumer mqwrapper.Consumer) { defer ms.wait.Done() for { @@ -649,10 +649,10 @@ var _ MsgStream = (*MqTtMsgStream)(nil) // MqTtMsgStream is a msgstream that contains timeticks type MqTtMsgStream struct { mqMsgStream - chanMsgBuf map[mqclient.Consumer][]TsMsg - chanMsgPos map[mqclient.Consumer]*internalpb.MsgPosition - chanStopChan map[mqclient.Consumer]chan bool - chanTtMsgTime map[mqclient.Consumer]Timestamp + chanMsgBuf map[mqwrapper.Consumer][]TsMsg + chanMsgPos map[mqwrapper.Consumer]*internalpb.MsgPosition + chanStopChan map[mqwrapper.Consumer]chan bool + chanTtMsgTime map[mqwrapper.Consumer]Timestamp chanMsgBufMutex *sync.Mutex chanTtMsgTimeMutex *sync.RWMutex chanWaitGroup *sync.WaitGroup @@ -664,16 +664,16 @@ type MqTtMsgStream struct { func NewMqTtMsgStream(ctx context.Context, receiveBufSize int64, bufSize int64, - client mqclient.Client, + client mqwrapper.Client, unmarshal UnmarshalDispatcher) (*MqTtMsgStream, error) { msgStream, err := NewMqMsgStream(ctx, receiveBufSize, bufSize, client, unmarshal) if err != nil { return nil, err } - chanMsgBuf := make(map[mqclient.Consumer][]TsMsg) - chanMsgPos := make(map[mqclient.Consumer]*internalpb.MsgPosition) - chanStopChan := make(map[mqclient.Consumer]chan bool) - chanTtMsgTime := make(map[mqclient.Consumer]Timestamp) + chanMsgBuf := make(map[mqwrapper.Consumer][]TsMsg) + chanMsgPos := make(map[mqwrapper.Consumer]*internalpb.MsgPosition) + chanStopChan := make(map[mqwrapper.Consumer]chan bool) + chanTtMsgTime := make(map[mqwrapper.Consumer]Timestamp) syncConsumer := make(chan int, 1) return &MqTtMsgStream{ @@ -689,7 +689,7 @@ func NewMqTtMsgStream(ctx context.Context, }, nil } -func (ms *MqTtMsgStream) addConsumer(consumer mqclient.Consumer, channel string) { +func (ms *MqTtMsgStream) addConsumer(consumer mqwrapper.Consumer, channel string) { if len(ms.consumers) == 0 { ms.syncConsumer <- 1 } @@ -707,20 +707,20 @@ func (ms *MqTtMsgStream) addConsumer(consumer mqclient.Consumer, channel string) // AsConsumer subscribes channels as consumer for a MsgStream func (ms *MqTtMsgStream) AsConsumer(channels []string, subName string) { - ms.AsConsumerWithPosition(channels, subName, mqclient.SubscriptionPositionEarliest) + ms.AsConsumerWithPosition(channels, subName, mqwrapper.SubscriptionPositionEarliest) } // AsConsumerWithPosition subscribes channels as consumer for a MsgStream and seeks to a certain position. -func (ms *MqTtMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) { +func (ms *MqTtMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqwrapper.SubscriptionInitialPosition) { for _, channel := range channels { if _, ok := ms.consumers[channel]; ok { continue } fn := func() error { - pc, err := ms.client.Subscribe(mqclient.ConsumerOptions{ + pc, err := ms.client.Subscribe(mqwrapper.ConsumerOptions{ Topic: channel, SubscriptionName: subName, - Type: mqclient.Exclusive, + Type: mqwrapper.Exclusive, SubscriptionInitialPosition: position, BufSize: ms.bufSize, }) @@ -778,7 +778,7 @@ func (ms *MqTtMsgStream) Close() { func (ms *MqTtMsgStream) bufMsgPackToChannel() { defer ms.wait.Done() - chanTtMsgSync := make(map[mqclient.Consumer]bool) + chanTtMsgSync := make(map[mqwrapper.Consumer]bool) // block here until addConsumer if _, ok := <-ms.syncConsumer; !ok { @@ -876,7 +876,7 @@ func (ms *MqTtMsgStream) bufMsgPackToChannel() { } // Save all msgs into chanMsgBuf[] till receive one ttMsg -func (ms *MqTtMsgStream) consumeToTtMsg(consumer mqclient.Consumer) { +func (ms *MqTtMsgStream) consumeToTtMsg(consumer mqwrapper.Consumer) { defer ms.chanWaitGroup.Done() for { select { @@ -923,7 +923,7 @@ func (ms *MqTtMsgStream) consumeToTtMsg(consumer mqclient.Consumer) { } // return true only when all channels reach same timetick -func (ms *MqTtMsgStream) allChanReachSameTtMsg(chanTtMsgSync map[mqclient.Consumer]bool) (Timestamp, bool) { +func (ms *MqTtMsgStream) allChanReachSameTtMsg(chanTtMsgSync map[mqwrapper.Consumer]bool) (Timestamp, bool) { tsMap := make(map[Timestamp]int) var maxTime Timestamp for _, t := range ms.chanTtMsgTime { @@ -950,7 +950,7 @@ func (ms *MqTtMsgStream) allChanReachSameTtMsg(chanTtMsgSync map[mqclient.Consum // Seek to the specified position func (ms *MqTtMsgStream) Seek(msgPositions []*internalpb.MsgPosition) error { - var consumer mqclient.Consumer + var consumer mqwrapper.Consumer var mp *MsgPosition var err error fn := func() error { diff --git a/internal/msgstream/mq_msgstream_test.go b/internal/mq/msgstream/mq_msgstream_test.go similarity index 94% rename from internal/msgstream/mq_msgstream_test.go rename to internal/mq/msgstream/mq_msgstream_test.go index 18c6d8ab41..06f4b3ce0b 100644 --- a/internal/msgstream/mq_msgstream_test.go +++ b/internal/mq/msgstream/mq_msgstream_test.go @@ -29,21 +29,23 @@ import ( "time" "unsafe" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" + "github.com/apache/pulsar-client-go/pulsar" + pulsarwrapper "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/pulsar" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/rmq" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/milvus-io/milvus/internal/allocator" "github.com/milvus-io/milvus/internal/common" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/util/etcd" "github.com/milvus-io/milvus/internal/util/funcutil" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/milvus-io/milvus/internal/util/paramtable" - client "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" - "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" ) var Params paramtable.BaseTable @@ -60,12 +62,12 @@ type fixture struct { } type parameters struct { - client mqclient.Client + client mqwrapper.Client } func (f *fixture) setup() []parameters { pulsarAddress, _ := Params.Load("_PulsarAddress") - pulsarClient, err := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, err := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) assert.Nil(f.t, err) rocksdbName := "/tmp/rocksmq_unittest_" + f.t.Name() @@ -82,12 +84,12 @@ func (f *fixture) setup() []parameters { f.etcdKV = etcdkv.NewEtcdKV(etcdCli, "/etcd/test/root") idAllocator := allocator.NewGlobalIDAllocator("dummy", f.etcdKV) _ = idAllocator.Initialize() - err = rocksmq.InitRmq(rocksdbName, idAllocator) + err = server.InitRmq(rocksdbName, idAllocator) if err != nil { log.Fatalf("InitRmq error = %v", err) } - rmqClient, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient, _ := rmq.NewClientWithDefaultOptions() parameters := []parameters{ {pulsarClient}, {rmqClient}, @@ -98,7 +100,7 @@ func (f *fixture) setup() []parameters { func (f *fixture) teardown() { rocksdbName := "/tmp/rocksmq_unittest_" + f.t.Name() - rocksmq.CloseRocksMQ() + server.CloseRocksMQ() f.etcdKV.Close() _ = os.RemoveAll(rocksdbName) _ = os.RemoveAll(rocksdbName + "_meta_kv") @@ -111,14 +113,14 @@ func Test_NewMqMsgStream(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { _, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) }(parameters[i].client) } } -// TODO(wxyu): add a mock implement of mqclient.Client, then inject errors to improve coverage +// TODO(wxyu): add a mock implement of mqwrapper.Client, then inject errors to improve coverage func TestMqMsgStream_AsProducer(t *testing.T) { f := &fixture{t: t} parameters := f.setup() @@ -126,7 +128,7 @@ func TestMqMsgStream_AsProducer(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -136,7 +138,7 @@ func TestMqMsgStream_AsProducer(t *testing.T) { } } -// TODO(wxyu): add a mock implement of mqclient.Client, then inject errors to improve coverage +// TODO(wxyu): add a mock implement of mqwrapper.Client, then inject errors to improve coverage func TestMqMsgStream_AsConsumer(t *testing.T) { f := &fixture{t: t} parameters := f.setup() @@ -144,7 +146,7 @@ func TestMqMsgStream_AsConsumer(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -162,7 +164,7 @@ func TestMqMsgStream_ComputeProduceChannelIndexes(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -206,7 +208,7 @@ func TestMqMsgStream_GetProduceChannels(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -229,7 +231,7 @@ func TestMqMsgStream_Produce(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -272,7 +274,7 @@ func TestMqMsgStream_Broadcast(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -290,7 +292,7 @@ func TestMqMsgStream_Consume(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { // Consume return nil when ctx canceled var wg sync.WaitGroup ctx, cancel := context.WithCancel(context.Background()) @@ -317,7 +319,7 @@ func TestMqMsgStream_Chan(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -334,7 +336,7 @@ func TestMqMsgStream_SeekNotSubscribed(t *testing.T) { factory := &ProtoUDFactory{} for i := range parameters { - func(client mqclient.Client) { + func(client mqwrapper.Client) { m, err := NewMqMsgStream(context.Background(), 100, 100, client, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -538,12 +540,12 @@ func TestStream_PulsarMsgStream_InsertRepackFunc(t *testing.T) { factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) inputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) inputStream.AsProducer(producerChannels) inputStream.Start() - pulsarClient2, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient2, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient2, factory.NewUnmarshalDispatcher()) outputStream.AsConsumer(consumerChannels, consumerSubName) outputStream.Start() @@ -591,12 +593,12 @@ func TestStream_PulsarMsgStream_DeleteRepackFunc(t *testing.T) { msgPack.Msgs = append(msgPack.Msgs, deleteMsg) factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) inputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) inputStream.AsProducer(producerChannels) inputStream.Start() - pulsarClient2, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient2, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient2, factory.NewUnmarshalDispatcher()) outputStream.AsConsumer(consumerChannels, consumerSubName) outputStream.Start() @@ -624,12 +626,12 @@ func TestStream_PulsarMsgStream_DefaultRepackFunc(t *testing.T) { msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_QueryNodeStats, 4)) factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) inputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) inputStream.AsProducer(producerChannels) inputStream.Start() - pulsarClient2, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient2, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient2, factory.NewUnmarshalDispatcher()) outputStream.AsConsumer(consumerChannels, consumerSubName) outputStream.Start() @@ -1103,7 +1105,7 @@ func TestStream_MqMsgStream_Seek(t *testing.T) { outputStream.Close() factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream2, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) outputStream2.AsConsumer(consumerChannels, consumerSubName) outputStream2.Seek([]*internalpb.MsgPosition{seekPosition}) @@ -1143,7 +1145,7 @@ func TestStream_MqMsgStream_SeekInvalidMessage(t *testing.T) { } factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream2, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) outputStream2.AsConsumer(consumerChannels, funcutil.RandomString(8)) defer outputStream2.Close() @@ -1200,7 +1202,7 @@ func TestStream_RMqMsgStream_SeekInvalidMessage(t *testing.T) { outputStream.Close() factory := ProtoUDFactory{} - rmqClient2, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient2, _ := rmq.NewClientWithDefaultOptions() outputStream2, _ := NewMqMsgStream(context.Background(), 100, 100, rmqClient2, factory.NewUnmarshalDispatcher()) outputStream2.AsConsumer(consumerChannels, funcutil.RandomString(8)) @@ -1252,9 +1254,9 @@ func TestStream_MqMsgStream_SeekLatest(t *testing.T) { err := inputStream.Produce(msgPack) assert.Nil(t, err) factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream2, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) - outputStream2.AsConsumerWithPosition(consumerChannels, consumerSubName, mqclient.SubscriptionPositionLatest) + outputStream2.AsConsumerWithPosition(consumerChannels, consumerSubName, mqwrapper.SubscriptionPositionLatest) outputStream2.Start() msgPack.Msgs = nil @@ -1353,7 +1355,7 @@ func initRmq(name string) *etcdkv.EtcdKV { idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV) _ = idAllocator.Initialize() - err = rocksmq.InitRmq(name, idAllocator) + err = server.InitRmq(name, idAllocator) if err != nil { log.Fatalf("InitRmq error = %v", err) @@ -1362,7 +1364,7 @@ func initRmq(name string) *etcdkv.EtcdKV { } func Close(rocksdbName string, intputStream, outputStream MsgStream, etcdKV *etcdkv.EtcdKV) { - rocksmq.CloseRocksMQ() + server.CloseRocksMQ() intputStream.Close() outputStream.Close() etcdKV.Close() @@ -1377,7 +1379,7 @@ func initRmqStream(producerChannels []string, opts ...RepackFunc) (MsgStream, MsgStream) { factory := ProtoUDFactory{} - rmqClient, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient, _ := rmq.NewClientWithDefaultOptions() inputStream, _ := NewMqMsgStream(context.Background(), 100, 100, rmqClient, factory.NewUnmarshalDispatcher()) inputStream.AsProducer(producerChannels) for _, opt := range opts { @@ -1386,7 +1388,7 @@ func initRmqStream(producerChannels []string, inputStream.Start() var input MsgStream = inputStream - rmqClient2, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient2, _ := rmq.NewClientWithDefaultOptions() outputStream, _ := NewMqMsgStream(context.Background(), 100, 100, rmqClient2, factory.NewUnmarshalDispatcher()) outputStream.AsConsumer(consumerChannels, consumerGroupName) outputStream.Start() @@ -1401,7 +1403,7 @@ func initRmqTtStream(producerChannels []string, opts ...RepackFunc) (MsgStream, MsgStream) { factory := ProtoUDFactory{} - rmqClient, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient, _ := rmq.NewClientWithDefaultOptions() inputStream, _ := NewMqMsgStream(context.Background(), 100, 100, rmqClient, factory.NewUnmarshalDispatcher()) inputStream.AsProducer(producerChannels) for _, opt := range opts { @@ -1410,7 +1412,7 @@ func initRmqTtStream(producerChannels []string, inputStream.Start() var input MsgStream = inputStream - rmqClient2, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient2, _ := rmq.NewClientWithDefaultOptions() outputStream, _ := NewMqTtMsgStream(context.Background(), 100, 100, rmqClient2, factory.NewUnmarshalDispatcher()) outputStream.AsConsumer(consumerChannels, consumerGroupName) outputStream.Start() @@ -1559,7 +1561,7 @@ func TestStream_RmqTtMsgStream_Seek(t *testing.T) { factory := ProtoUDFactory{} - rmqClient, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient, _ := rmq.NewClientWithDefaultOptions() outputStream, _ = NewMqTtMsgStream(context.Background(), 100, 100, rmqClient, factory.NewUnmarshalDispatcher()) consumerSubName = funcutil.RandomString(8) outputStream.AsConsumer(consumerChannels, consumerSubName) @@ -1589,7 +1591,7 @@ func TestStream_BroadcastMark(t *testing.T) { producerChannels := []string{c1, c2} factory := ProtoUDFactory{} - pulsarClient, err := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, err := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) assert.Nil(t, err) outputStream, err := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -1651,7 +1653,7 @@ func TestStream_ProduceMark(t *testing.T) { producerChannels := []string{c1, c2} factory := ProtoUDFactory{} - pulsarClient, err := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, err := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) assert.Nil(t, err) outputStream, err := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) assert.Nil(t, err) @@ -1737,13 +1739,13 @@ func (t *MarshalFailTsMsg) Unmarshal(_ MarshalType) (TsMsg, error) { return nil, errors.New("mocked error") } -var _ mqclient.Producer = (*mockSendFailProducer)(nil) +var _ mqwrapper.Producer = (*mockSendFailProducer)(nil) type mockSendFailProducer struct { - mqclient.Producer + mqwrapper.Producer } -func (p *mockSendFailProducer) Send(_ context.Context, _ *mqclient.ProducerMessage) (MessageID, error) { +func (p *mockSendFailProducer) Send(_ context.Context, _ *mqwrapper.ProducerMessage) (MessageID, error) { return nil, errors.New("mocked error") } @@ -1928,7 +1930,7 @@ func getTimeTickMsgPack(reqID UniqueID) *MsgPack { func getPulsarInputStream(pulsarAddress string, producerChannels []string, opts ...RepackFunc) MsgStream { factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) inputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) inputStream.AsProducer(producerChannels) for _, opt := range opts { @@ -1940,7 +1942,7 @@ func getPulsarInputStream(pulsarAddress string, producerChannels []string, opts func getPulsarOutputStream(pulsarAddress string, consumerChannels []string, consumerSubName string) MsgStream { factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) outputStream.AsConsumer(consumerChannels, consumerSubName) outputStream.Start() @@ -1949,7 +1951,7 @@ func getPulsarOutputStream(pulsarAddress string, consumerChannels []string, cons func getPulsarReader(pulsarAddress string, consumerChannels []string) MsgStream { factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream, _ := NewMqMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) outputStream.AsReader(consumerChannels, "pulsar-reader-prefix-") return outputStream @@ -1957,7 +1959,7 @@ func getPulsarReader(pulsarAddress string, consumerChannels []string) MsgStream func getPulsarTtOutputStream(pulsarAddress string, consumerChannels []string, consumerSubName string) MsgStream { factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream, _ := NewMqTtMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) outputStream.AsConsumer(consumerChannels, consumerSubName) outputStream.Start() @@ -1966,7 +1968,7 @@ func getPulsarTtOutputStream(pulsarAddress string, consumerChannels []string, co func getPulsarTtOutputStreamAndSeek(pulsarAddress string, positions []*MsgPosition) MsgStream { factory := ProtoUDFactory{} - pulsarClient, _ := mqclient.GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pulsarClient, _ := pulsarwrapper.NewClient(pulsar.ClientOptions{URL: pulsarAddress}) outputStream, _ := NewMqTtMsgStream(context.Background(), 100, 100, pulsarClient, factory.NewUnmarshalDispatcher()) consumerName := []string{} for _, c := range positions { @@ -2017,7 +2019,7 @@ func TestStream_RmqTtMsgStream_AsConsumerWithPosition(t *testing.T) { etcdKV := initRmq(rocksdbName) factory := ProtoUDFactory{} - rmqClient, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient, _ := rmq.NewClientWithDefaultOptions() otherInputStream, _ := NewMqMsgStream(context.Background(), 100, 100, rmqClient, factory.NewUnmarshalDispatcher()) otherInputStream.AsProducer([]string{"root_timetick"}) @@ -2032,9 +2034,9 @@ func TestStream_RmqTtMsgStream_AsConsumerWithPosition(t *testing.T) { inputStream.Produce(getTimeTickMsgPack(int64(i))) } - rmqClient2, _ := mqclient.NewRmqClient(client.ClientOptions{Server: rocksmq.Rmq}) + rmqClient2, _ := rmq.NewClientWithDefaultOptions() outputStream, _ := NewMqMsgStream(context.Background(), 100, 100, rmqClient2, factory.NewUnmarshalDispatcher()) - outputStream.AsConsumerWithPosition(consumerChannels, consumerSubName, mqclient.SubscriptionPositionLatest) + outputStream.AsConsumerWithPosition(consumerChannels, consumerSubName, mqwrapper.SubscriptionPositionLatest) outputStream.Start() inputStream.Produce(getTimeTickMsgPack(1000)) diff --git a/internal/util/mqclient/client.go b/internal/mq/msgstream/mqwrapper/client.go similarity index 98% rename from internal/util/mqclient/client.go rename to internal/mq/msgstream/mqwrapper/client.go index a3bd0af746..081139f4f7 100644 --- a/internal/util/mqclient/client.go +++ b/internal/mq/msgstream/mqwrapper/client.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package mqwrapper // Client is the interface that provides operations of message queues type Client interface { diff --git a/internal/util/mqclient/consumer.go b/internal/mq/msgstream/mqwrapper/consumer.go similarity index 99% rename from internal/util/mqclient/consumer.go rename to internal/mq/msgstream/mqwrapper/consumer.go index 730cfe33dd..967e5806ed 100644 --- a/internal/util/mqclient/consumer.go +++ b/internal/mq/msgstream/mqwrapper/consumer.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package mqwrapper // SubscriptionInitialPosition is the type of a subscription initial position type SubscriptionInitialPosition int diff --git a/internal/util/mqclient/id.go b/internal/mq/msgstream/mqwrapper/id.go similarity index 98% rename from internal/util/mqclient/id.go rename to internal/mq/msgstream/mqwrapper/id.go index 3775e881a5..114ff2cfbd 100644 --- a/internal/util/mqclient/id.go +++ b/internal/mq/msgstream/mqwrapper/id.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package mqwrapper // MessageID is the interface that provides operations of message is type MessageID interface { diff --git a/internal/util/mqclient/message.go b/internal/mq/msgstream/mqwrapper/message.go similarity index 98% rename from internal/util/mqclient/message.go rename to internal/mq/msgstream/mqwrapper/message.go index 8d69bf1b8e..dbb13484e5 100644 --- a/internal/util/mqclient/message.go +++ b/internal/mq/msgstream/mqwrapper/message.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package mqwrapper // Message is the interface that provides operations of a consumer type Message interface { diff --git a/internal/util/mqclient/producer.go b/internal/mq/msgstream/mqwrapper/producer.go similarity index 98% rename from internal/util/mqclient/producer.go rename to internal/mq/msgstream/mqwrapper/producer.go index 50e39584b8..b229e00b44 100644 --- a/internal/util/mqclient/producer.go +++ b/internal/mq/msgstream/mqwrapper/producer.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package mqwrapper import "context" diff --git a/internal/util/mqclient/pulsar_client.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_client.go similarity index 81% rename from internal/util/mqclient/pulsar_client.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_client.go index 5115323808..1c919f7815 100644 --- a/internal/util/mqclient/pulsar_client.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_client.go @@ -14,12 +14,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "errors" "sync" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/util/retry" @@ -51,9 +53,9 @@ func isPulsarError(err error, result ...pulsar.Result) bool { return false } -// GetPulsarClientInstance creates a pulsarClient object +// NewClient creates a pulsarClient object // according to the parameter opts of type pulsar.ClientOptions -func GetPulsarClientInstance(opts pulsar.ClientOptions) (*pulsarClient, error) { +func NewClient(opts pulsar.ClientOptions) (*pulsarClient, error) { once.Do(func() { c, err := pulsar.NewClient(opts) if err != nil { @@ -67,7 +69,7 @@ func GetPulsarClientInstance(opts pulsar.ClientOptions) (*pulsarClient, error) { } // CreateProducer create a pulsar producer from options -func (pc *pulsarClient) CreateProducer(options ProducerOptions) (Producer, error) { +func (pc *pulsarClient) CreateProducer(options mqwrapper.ProducerOptions) (mqwrapper.Producer, error) { opts := pulsar.ProducerOptions{Topic: options.Topic} pp, err := pc.client.CreateProducer(opts) if err != nil { @@ -81,7 +83,7 @@ func (pc *pulsarClient) CreateProducer(options ProducerOptions) (Producer, error } // CreateReader creates a pulsar reader instance -func (pc *pulsarClient) CreateReader(options ReaderOptions) (Reader, error) { +func (pc *pulsarClient) CreateReader(options mqwrapper.ReaderOptions) (mqwrapper.Reader, error) { opts := pulsar.ReaderOptions{ Topic: options.Topic, StartMessageID: options.StartMessageID.(*pulsarID).messageID, @@ -100,7 +102,7 @@ func (pc *pulsarClient) CreateReader(options ReaderOptions) (Reader, error) { } // Subscribe creates a pulsar consumer instance and subscribe a topic -func (pc *pulsarClient) Subscribe(options ConsumerOptions) (Consumer, error) { +func (pc *pulsarClient) Subscribe(options mqwrapper.ConsumerOptions) (mqwrapper.Consumer, error) { receiveChannel := make(chan pulsar.ConsumerMessage, options.BufSize) consumer, err := pc.client.Subscribe(pulsar.ConsumerOptions{ Topic: options.Topic, @@ -117,9 +119,9 @@ func (pc *pulsarClient) Subscribe(options ConsumerOptions) (Consumer, error) { return nil, err } - pConsumer := &PulsarConsumer{c: consumer, closeCh: make(chan struct{})} + pConsumer := &Consumer{c: consumer, closeCh: make(chan struct{})} // prevent seek to earliest patch applied when using latest position options - if options.SubscriptionInitialPosition == SubscriptionPositionLatest { + if options.SubscriptionInitialPosition == mqwrapper.SubscriptionPositionLatest { pConsumer.AtLatest = true } @@ -127,14 +129,14 @@ func (pc *pulsarClient) Subscribe(options ConsumerOptions) (Consumer, error) { } // EarliestMessageID returns the earliest message id -func (pc *pulsarClient) EarliestMessageID() MessageID { +func (pc *pulsarClient) EarliestMessageID() mqwrapper.MessageID { msgID := pulsar.EarliestMessageID() return &pulsarID{messageID: msgID} } // StringToMsgID converts the string id to MessageID type -func (pc *pulsarClient) StringToMsgID(id string) (MessageID, error) { - pID, err := StringToPulsarMsgID(id) +func (pc *pulsarClient) StringToMsgID(id string) (mqwrapper.MessageID, error) { + pID, err := stringToMsgID(id) if err != nil { return nil, err } @@ -142,7 +144,7 @@ func (pc *pulsarClient) StringToMsgID(id string) (MessageID, error) { } // BytesToMsgID converts []byte id to MessageID type -func (pc *pulsarClient) BytesToMsgID(id []byte) (MessageID, error) { +func (pc *pulsarClient) BytesToMsgID(id []byte) (mqwrapper.MessageID, error) { pID, err := DeserializePulsarMsgID(id) if err != nil { return nil, err diff --git a/internal/util/mqclient/pulsar_client_test.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_client_test.go similarity index 88% rename from internal/util/mqclient/pulsar_client_test.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_client_test.go index 5526bb73ae..26e973219d 100644 --- a/internal/util/mqclient/pulsar_client_test.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_client_test.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "bytes" @@ -23,19 +23,31 @@ import ( "errors" "fmt" "math/rand" + "os" "testing" "time" "unsafe" - "github.com/milvus-io/milvus/internal/util/retry" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/milvus-io/milvus/internal/util/paramtable" + "go.uber.org/zap" "github.com/apache/pulsar-client-go/pulsar" "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" + "github.com/milvus-io/milvus/internal/util/retry" "github.com/stretchr/testify/assert" ) +var Params paramtable.BaseTable + +func TestMain(m *testing.M) { + Params.Init() + exitCode := m.Run() + os.Exit(exitCode) +} + func IntToBytes(n int) []byte { tmp := int32(n) bytesBuffer := bytes.NewBuffer([]byte{}) @@ -51,14 +63,14 @@ func BytesToInt(b []byte) int { } func Produce(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, arr []int) { - producer, err := pc.CreateProducer(ProducerOptions{Topic: topic}) + producer, err := pc.CreateProducer(mqwrapper.ProducerOptions{Topic: topic}) assert.Nil(t, err) assert.NotNil(t, producer) log.Info("Produce start") for _, v := range arr { - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: IntToBytes(v), Properties: map[string]string{}, } @@ -70,7 +82,7 @@ func Produce(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, log.Info("Produce done") } -func VerifyMessage(t *testing.T, msg Message) { +func VerifyMessage(t *testing.T, msg mqwrapper.Message) { pload := BytesToInt(msg.Payload()) log.Info("RECV", zap.Any("v", pload)) pm := msg.(*pulsarMessage) @@ -82,13 +94,13 @@ func VerifyMessage(t *testing.T, msg Message) { } // Consume1 will consume random messages and record the last MessageID it received -func Consume1(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, c chan MessageID, total *int) { - consumer, err := pc.Subscribe(ConsumerOptions{ +func Consume1(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, c chan mqwrapper.MessageID, total *int) { + consumer, err := pc.Subscribe(mqwrapper.ConsumerOptions{ Topic: topic, SubscriptionName: subName, BufSize: 1024, - Type: KeyShared, - SubscriptionInitialPosition: SubscriptionPositionEarliest, + Type: mqwrapper.KeyShared, + SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest, }) assert.Nil(t, err) assert.NotNil(t, consumer) @@ -100,7 +112,7 @@ func Consume1(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, rand.Seed(time.Now().UnixNano()) cnt := 1 + rand.Int()%5 - var msg Message + var msg mqwrapper.Message for i := 0; i < cnt; i++ { select { case <-ctx.Done(): @@ -120,13 +132,13 @@ func Consume1(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, } // Consume2 will consume messages from specified MessageID -func Consume2(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, msgID MessageID, total *int) { - consumer, err := pc.Subscribe(ConsumerOptions{ +func Consume2(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, msgID mqwrapper.MessageID, total *int) { + consumer, err := pc.Subscribe(mqwrapper.ConsumerOptions{ Topic: topic, SubscriptionName: subName, BufSize: 1024, - Type: KeyShared, - SubscriptionInitialPosition: SubscriptionPositionEarliest, + Type: mqwrapper.KeyShared, + SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest, }) assert.Nil(t, err) assert.NotNil(t, consumer) @@ -156,12 +168,12 @@ func Consume2(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, } func Consume3(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, total *int) { - consumer, err := pc.Subscribe(ConsumerOptions{ + consumer, err := pc.Subscribe(mqwrapper.ConsumerOptions{ Topic: topic, SubscriptionName: subName, BufSize: 1024, - Type: KeyShared, - SubscriptionInitialPosition: SubscriptionPositionEarliest, + Type: mqwrapper.KeyShared, + SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest, }) assert.Nil(t, err) assert.NotNil(t, consumer) @@ -185,7 +197,7 @@ func Consume3(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, func TestPulsarClient_Consume1(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer pc.Close() assert.NoError(t, err) assert.NotNil(t, pc) @@ -194,7 +206,7 @@ func TestPulsarClient_Consume1(t *testing.T) { topic := fmt.Sprintf("test-topic-%d", rand.Int()) subName := fmt.Sprintf("test-subname-%d", rand.Int()) arr := []int{111, 222, 333, 444, 555, 666, 777} - c := make(chan MessageID, 1) + c := make(chan mqwrapper.MessageID, 1) ctx, cancel := context.WithCancel(context.Background()) @@ -233,7 +245,7 @@ func TestPulsarClient_Consume1(t *testing.T) { log.Info("main done") } -func Consume21(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, c chan MessageID, total *int) { +func Consume21(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, c chan mqwrapper.MessageID, total *int) { consumer, err := pc.client.Subscribe(pulsar.ConsumerOptions{ Topic: topic, SubscriptionName: subName, @@ -271,7 +283,7 @@ func Consume21(ctx context.Context, t *testing.T, pc *pulsarClient, topic string } // Consume2 will consume messages from specified MessageID -func Consume22(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, msgID MessageID, total *int) { +func Consume22(ctx context.Context, t *testing.T, pc *pulsarClient, topic string, subName string, msgID mqwrapper.MessageID, total *int) { consumer, err := pc.client.Subscribe(pulsar.ConsumerOptions{ Topic: topic, SubscriptionName: subName, @@ -336,7 +348,7 @@ func Consume23(ctx context.Context, t *testing.T, pc *pulsarClient, topic string func TestPulsarClient_Consume2(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer pc.Close() assert.NoError(t, err) assert.NotNil(t, pc) @@ -345,7 +357,7 @@ func TestPulsarClient_Consume2(t *testing.T) { topic := fmt.Sprintf("test-topic-%d", rand.Int()) subName := fmt.Sprintf("test-subname-%d", rand.Int()) arr := []int{111, 222, 333, 444, 555, 666, 777} - c := make(chan MessageID, 1) + c := make(chan mqwrapper.MessageID, 1) ctx, cancel := context.WithCancel(context.Background()) @@ -386,7 +398,7 @@ func TestPulsarClient_Consume2(t *testing.T) { func TestPulsarClient_SeekPosition(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer pc.Close() assert.NoError(t, err) assert.NotNil(t, pc) @@ -396,15 +408,15 @@ func TestPulsarClient_SeekPosition(t *testing.T) { topic := fmt.Sprintf("test-topic-%d", rand.Int()) subName := fmt.Sprintf("test-subname-%d", rand.Int()) - producer, err := pc.CreateProducer(ProducerOptions{Topic: topic}) + producer, err := pc.CreateProducer(mqwrapper.ProducerOptions{Topic: topic}) assert.Nil(t, err) assert.NotNil(t, producer) log.Info("Produce start") - ids := []MessageID{} + ids := []mqwrapper.MessageID{} arr := []int{1, 2, 3} for _, v := range arr { - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: IntToBytes(v), Properties: map[string]string{}, } @@ -459,7 +471,7 @@ func TestPulsarClient_SeekPosition(t *testing.T) { func TestPulsarClient_SeekLatest(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer pc.Close() assert.NoError(t, err) assert.NotNil(t, pc) @@ -469,7 +481,7 @@ func TestPulsarClient_SeekLatest(t *testing.T) { topic := fmt.Sprintf("test-topic-%d", rand.Int()) subName := fmt.Sprintf("test-subname-%d", rand.Int()) - producer, err := pc.CreateProducer(ProducerOptions{Topic: topic}) + producer, err := pc.CreateProducer(mqwrapper.ProducerOptions{Topic: topic}) assert.Nil(t, err) assert.NotNil(t, producer) @@ -477,7 +489,7 @@ func TestPulsarClient_SeekLatest(t *testing.T) { arr := []int{1, 2, 3} for _, v := range arr { - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: IntToBytes(v), Properties: map[string]string{}, } @@ -510,7 +522,7 @@ func TestPulsarClient_SeekLatest(t *testing.T) { loop = false case <-time.After(2 * time.Second): log.Info("after 2 seconds") - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: IntToBytes(4), Properties: map[string]string{}, } @@ -522,7 +534,7 @@ func TestPulsarClient_SeekLatest(t *testing.T) { func TestPulsarClient_EarliestMessageID(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - client, _ := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + client, _ := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer client.Close() mid := client.EarliestMessageID() @@ -531,11 +543,11 @@ func TestPulsarClient_EarliestMessageID(t *testing.T) { func TestPulsarClient_StringToMsgID(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - client, _ := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + client, _ := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer client.Close() mid := client.EarliestMessageID() - str := PulsarMsgIDToString(mid) + str := msgIDToString(mid) res, err := client.StringToMsgID(str) assert.Nil(t, err) @@ -549,7 +561,7 @@ func TestPulsarClient_StringToMsgID(t *testing.T) { func TestPulsarClient_BytesToMsgID(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - client, _ := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + client, _ := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer client.Close() mid := pulsar.EarliestMessageID() @@ -657,7 +669,7 @@ func TestPulsarClient_SubscribeExclusiveFail(t *testing.T) { client: &mockPulsarClient{}, } - _, err := pc.Subscribe(ConsumerOptions{}) + _, err := pc.Subscribe(mqwrapper.ConsumerOptions{}) assert.Error(t, err) assert.True(t, retry.IsUnRecoverable(err)) }) diff --git a/internal/util/mqclient/pulsar_consumer.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer.go similarity index 90% rename from internal/util/mqclient/pulsar_consumer.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer.go index 1b208a422f..7aaec11649 100644 --- a/internal/util/mqclient/pulsar_consumer.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "context" @@ -22,17 +22,19 @@ import ( "time" "unsafe" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/util/retry" "go.uber.org/zap" ) -// PulsarConsumer consumes from pulsar -type PulsarConsumer struct { +// Consumer consumes from pulsar +type Consumer struct { c pulsar.Consumer pulsar.Reader - msgChannel chan Message + msgChannel chan mqwrapper.Message hasSeek bool AtLatest bool closeCh chan struct{} @@ -42,15 +44,15 @@ type PulsarConsumer struct { } // Subscription get a subscription for the consumer -func (pc *PulsarConsumer) Subscription() string { +func (pc *Consumer) Subscription() string { return pc.c.Subscription() } // Chan returns a message channel -func (pc *PulsarConsumer) Chan() <-chan Message { +func (pc *Consumer) Chan() <-chan mqwrapper.Message { if pc.msgChannel == nil { pc.once.Do(func() { - pc.msgChannel = make(chan Message, 256) + pc.msgChannel = make(chan mqwrapper.Message, 256) // this part handles msgstream expectation when the consumer is not seeked // pulsar's default behavior is setting postition to the earliest pointer when client of the same subscription pointer is not acked // yet, our message stream is to setting to the very start point of the topic @@ -90,7 +92,7 @@ func (pc *PulsarConsumer) Chan() <-chan Message { // Seek seek consume position to the pointed messageID, // the pointed messageID will be consumed after the seek in pulsar -func (pc *PulsarConsumer) Seek(id MessageID, inclusive bool) error { +func (pc *Consumer) Seek(id mqwrapper.MessageID, inclusive bool) error { messageID := id.(*pulsarID).messageID err := pc.c.Seek(messageID) if err == nil { @@ -102,13 +104,13 @@ func (pc *PulsarConsumer) Seek(id MessageID, inclusive bool) error { } // Ack the consumption of a single message -func (pc *PulsarConsumer) Ack(message Message) { +func (pc *Consumer) Ack(message mqwrapper.Message) { pm := message.(*pulsarMessage) pc.c.Ack(pm.msg) } // Close the consumer and stop the broker to push more messages -func (pc *PulsarConsumer) Close() { +func (pc *Consumer) Close() { pc.closeOnce.Do(func() { defer pc.c.Close() // Unsubscribe for the consumer diff --git a/internal/util/mqclient/pulsar_consumer_test.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer_test.go similarity index 81% rename from internal/util/mqclient/pulsar_consumer_test.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer_test.go index d9739a71d0..74b6dd2769 100644 --- a/internal/util/mqclient/pulsar_consumer_test.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer_test.go @@ -14,19 +14,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "fmt" "testing" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" "github.com/stretchr/testify/assert" ) func TestPulsarConsumer_Subscription(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) assert.Nil(t, err) defer pc.Close() @@ -34,8 +36,8 @@ func TestPulsarConsumer_Subscription(t *testing.T) { consumer, err := pc.client.Subscribe(pulsar.ConsumerOptions{ Topic: "Topic", SubscriptionName: "SubName", - Type: pulsar.SubscriptionType(Exclusive), - SubscriptionInitialPosition: pulsar.SubscriptionInitialPosition(SubscriptionPositionEarliest), + Type: pulsar.SubscriptionType(mqwrapper.Exclusive), + SubscriptionInitialPosition: pulsar.SubscriptionInitialPosition(mqwrapper.SubscriptionPositionEarliest), MessageChannel: receiveChannel, }) assert.Nil(t, err) @@ -58,15 +60,15 @@ func Test_PatchEarliestMessageID(t *testing.T) { func TestPulsarConsumer_Close(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) assert.Nil(t, err) receiveChannel := make(chan pulsar.ConsumerMessage, 100) consumer, err := pc.client.Subscribe(pulsar.ConsumerOptions{ Topic: "Topic-1", SubscriptionName: "SubName-1", - Type: pulsar.SubscriptionType(Exclusive), - SubscriptionInitialPosition: pulsar.SubscriptionInitialPosition(SubscriptionPositionEarliest), + Type: pulsar.SubscriptionType(mqwrapper.Exclusive), + SubscriptionInitialPosition: pulsar.SubscriptionInitialPosition(mqwrapper.SubscriptionPositionEarliest), MessageChannel: receiveChannel, }) assert.Nil(t, err) @@ -75,7 +77,7 @@ func TestPulsarConsumer_Close(t *testing.T) { str := consumer.Subscription() assert.NotNil(t, str) - pulsarConsumer := &PulsarConsumer{c: consumer, closeCh: make(chan struct{})} + pulsarConsumer := &Consumer{c: consumer, closeCh: make(chan struct{})} pulsarConsumer.Close() // test double close diff --git a/internal/util/mqclient/pulsar_id.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_id.go similarity index 84% rename from internal/util/mqclient/pulsar_id.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_id.go index d5ad6a0200..f371a940bf 100644 --- a/internal/util/mqclient/pulsar_id.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_id.go @@ -14,11 +14,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "strings" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" ) @@ -28,7 +30,7 @@ type pulsarID struct { // Check if pulsarID implements pulsar.MessageID and MessageID interface var _ pulsar.MessageID = &pulsarID{} -var _ MessageID = &pulsarID{} +var _ mqwrapper.MessageID = &pulsarID{} func (pid *pulsarID) Serialize() []byte { return pid.messageID.Serialize() @@ -60,12 +62,12 @@ func DeserializePulsarMsgID(messageID []byte) (pulsar.MessageID, error) { return pulsar.DeserializeMessageID(messageID) } -// PulsarMsgIDToString is used to convert a message ID to string -func PulsarMsgIDToString(messageID pulsar.MessageID) string { +// msgIDToString is used to convert a message ID to string +func msgIDToString(messageID pulsar.MessageID) string { return strings.ToValidUTF8(string(messageID.Serialize()), "") } -// StringToPulsarMsgID is used to convert a string to message ID -func StringToPulsarMsgID(msgString string) (pulsar.MessageID, error) { +// StringToMsgID is used to convert a string to message ID +func stringToMsgID(msgString string) (pulsar.MessageID, error) { return pulsar.DeserializeMessageID([]byte(msgString)) } diff --git a/internal/util/mqclient/pulsar_id_test.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_id_test.go similarity index 93% rename from internal/util/mqclient/pulsar_id_test.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_id_test.go index da26cf2ac8..b434bb0aa6 100644 --- a/internal/util/mqclient/pulsar_id_test.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_id_test.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "testing" @@ -59,7 +59,7 @@ func Test_DeserializePulsarMsgID(t *testing.T) { func Test_PulsarMsgIDToString(t *testing.T) { mid := pulsar.EarliestMessageID() - str := PulsarMsgIDToString(mid) + str := msgIDToString(mid) assert.NotNil(t, str) assert.NotZero(t, len(str)) } @@ -67,8 +67,8 @@ func Test_PulsarMsgIDToString(t *testing.T) { func Test_StringToPulsarMsgID(t *testing.T) { mid := pulsar.EarliestMessageID() - str := PulsarMsgIDToString(mid) - res, err := StringToPulsarMsgID(str) + str := msgIDToString(mid) + res, err := stringToMsgID(str) assert.Nil(t, err) assert.NotNil(t, res) } diff --git a/internal/util/mqclient/pulsar_message.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_message.go similarity index 87% rename from internal/util/mqclient/pulsar_message.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_message.go index f8f28c558a..66a1c3967f 100644 --- a/internal/util/mqclient/pulsar_message.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_message.go @@ -14,14 +14,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "github.com/apache/pulsar-client-go/pulsar" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" ) // Check pulsarMessage implements ConsumerMessage -var _ Message = (*pulsarMessage)(nil) +var _ mqwrapper.Message = (*pulsarMessage)(nil) type pulsarMessage struct { msg pulsar.Message @@ -39,7 +40,7 @@ func (pm *pulsarMessage) Payload() []byte { return pm.msg.Payload() } -func (pm *pulsarMessage) ID() MessageID { +func (pm *pulsarMessage) ID() mqwrapper.MessageID { id := pm.msg.ID() pid := &pulsarID{messageID: id} return pid diff --git a/internal/util/mqclient/pulsar_producer.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer.go similarity index 83% rename from internal/util/mqclient/pulsar_producer.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer.go index ad3ece1352..6a3a8217d4 100644 --- a/internal/util/mqclient/pulsar_producer.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer.go @@ -14,16 +14,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "context" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" ) // implementation assertion -var _ Producer = (*pulsarProducer)(nil) +var _ mqwrapper.Producer = (*pulsarProducer)(nil) type pulsarProducer struct { p pulsar.Producer @@ -34,7 +36,7 @@ func (pp *pulsarProducer) Topic() string { return pp.p.Topic() } -func (pp *pulsarProducer) Send(ctx context.Context, message *ProducerMessage) (MessageID, error) { +func (pp *pulsarProducer) Send(ctx context.Context, message *mqwrapper.ProducerMessage) (mqwrapper.MessageID, error) { ppm := &pulsar.ProducerMessage{Payload: message.Payload, Properties: message.Properties} pmID, err := pp.p.Send(ctx, ppm) return &pulsarID{messageID: pmID}, err diff --git a/internal/util/mqclient/pulsar_producer_test.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer_test.go similarity index 84% rename from internal/util/mqclient/pulsar_producer_test.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer_test.go index 79dbd60989..fcfc96a150 100644 --- a/internal/util/mqclient/pulsar_producer_test.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer_test.go @@ -14,32 +14,34 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "context" "testing" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" "github.com/stretchr/testify/assert" ) func TestPulsarProducer(t *testing.T) { pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) defer pc.Close() assert.NoError(t, err) assert.NotNil(t, pc) topic := "TEST" - producer, err := pc.CreateProducer(ProducerOptions{Topic: topic}) + producer, err := pc.CreateProducer(mqwrapper.ProducerOptions{Topic: topic}) assert.Nil(t, err) assert.NotNil(t, producer) pulsarProd := producer.(*pulsarProducer) assert.Equal(t, pulsarProd.Topic(), topic) - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: []byte{}, Properties: map[string]string{}, } diff --git a/internal/util/mqclient/pulsar_reader.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader.go similarity index 87% rename from internal/util/mqclient/pulsar_reader.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader.go index b8c526f1f4..04e067822f 100644 --- a/internal/util/mqclient/pulsar_reader.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader.go @@ -14,11 +14,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "context" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" ) @@ -33,7 +35,7 @@ func (pr *pulsarReader) Topic() string { } // Next read the next message in the topic, blocking until a message is available -func (pr *pulsarReader) Next(ctx context.Context) (Message, error) { +func (pr *pulsarReader) Next(ctx context.Context) (mqwrapper.Message, error) { pm, err := pr.r.Next(ctx) if err != nil { return nil, err @@ -51,7 +53,7 @@ func (pr *pulsarReader) Close() { pr.r.Close() } -func (pr *pulsarReader) Seek(id MessageID) error { +func (pr *pulsarReader) Seek(id mqwrapper.MessageID) error { messageID := id.(*pulsarID).messageID err := pr.r.Seek(messageID) if err != nil { diff --git a/internal/util/mqclient/pulsar_reader_test.go b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader_test.go similarity index 83% rename from internal/util/mqclient/pulsar_reader_test.go rename to internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader_test.go index 1ace70286c..2065bea0e6 100644 --- a/internal/util/mqclient/pulsar_reader_test.go +++ b/internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader_test.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package pulsar import ( "context" @@ -23,6 +23,8 @@ import ( "testing" "time" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + "github.com/apache/pulsar-client-go/pulsar" "github.com/stretchr/testify/assert" ) @@ -30,23 +32,23 @@ import ( func TestPulsarReader(t *testing.T) { ctx := context.Background() pulsarAddress, _ := Params.Load("_PulsarAddress") - pc, err := GetPulsarClientInstance(pulsar.ClientOptions{URL: pulsarAddress}) + pc, err := NewClient(pulsar.ClientOptions{URL: pulsarAddress}) assert.Nil(t, err) defer pc.Close() rand.Seed(time.Now().UnixNano()) topic := fmt.Sprintf("test-%d", rand.Int()) - producer, err := pc.CreateProducer(ProducerOptions{Topic: topic}) + producer, err := pc.CreateProducer(mqwrapper.ProducerOptions{Topic: topic}) assert.Nil(t, err) assert.NotNil(t, producer) defer producer.Close() const N = 10 - var seekID MessageID + var seekID mqwrapper.MessageID for i := 0; i < N; i++ { - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: []byte(fmt.Sprintf("helloworld-%d", i)), Properties: map[string]string{}, } @@ -58,7 +60,7 @@ func TestPulsarReader(t *testing.T) { } } - reader, err := pc.CreateReader(ReaderOptions{ + reader, err := pc.CreateReader(mqwrapper.ReaderOptions{ Topic: topic, StartMessageID: pc.EarliestMessageID(), }) @@ -75,7 +77,7 @@ func TestPulsarReader(t *testing.T) { assert.NotNil(t, revMsg) } - readerOfStartMessageID, err := pc.CreateReader(ReaderOptions{ + readerOfStartMessageID, err := pc.CreateReader(mqwrapper.ReaderOptions{ Topic: topic, StartMessageID: seekID, StartMessageIDInclusive: true, @@ -90,7 +92,7 @@ func TestPulsarReader(t *testing.T) { assert.NotNil(t, revMsg) } - readerOfSeek, err := pc.CreateReader(ReaderOptions{ + readerOfSeek, err := pc.CreateReader(mqwrapper.ReaderOptions{ Topic: topic, StartMessageID: pc.EarliestMessageID(), }) diff --git a/internal/util/mqclient/reader.go b/internal/mq/msgstream/mqwrapper/reader.go similarity index 99% rename from internal/util/mqclient/reader.go rename to internal/mq/msgstream/mqwrapper/reader.go index 1b86934e22..bf3714e456 100644 --- a/internal/util/mqclient/reader.go +++ b/internal/mq/msgstream/mqwrapper/reader.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package mqwrapper import ( "context" diff --git a/internal/util/mqclient/rmq_client.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_client.go similarity index 64% rename from internal/util/mqclient/rmq_client.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_client.go index 76ff94c8ae..3446298f38 100644 --- a/internal/util/mqclient/rmq_client.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_client.go @@ -14,26 +14,34 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package rmq import ( "errors" "strconv" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" + + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "go.uber.org/zap" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" ) // rmqClient contains a rocksmq client type rmqClient struct { - client rocksmq.Client + client client.Client } -// NewRmqClient returns a new rmqClient object -func NewRmqClient(opts rocksmq.ClientOptions) (*rmqClient, error) { - c, err := rocksmq.NewClient(opts) +func NewClientWithDefaultOptions() (*rmqClient, error) { + option := client.Options{Server: server.Rmq} + return NewClient(option) +} + +// NewClient returns a new rmqClient object +func NewClient(opts client.Options) (*rmqClient, error) { + c, err := client.NewClient(opts) if err != nil { log.Error("Failed to set rmq client: ", zap.Error(err)) return nil, err @@ -42,8 +50,8 @@ func NewRmqClient(opts rocksmq.ClientOptions) (*rmqClient, error) { } // CreateProducer creates a producer for rocksmq client -func (rc *rmqClient) CreateProducer(options ProducerOptions) (Producer, error) { - rmqOpts := rocksmq.ProducerOptions{Topic: options.Topic} +func (rc *rmqClient) CreateProducer(options mqwrapper.ProducerOptions) (mqwrapper.Producer, error) { + rmqOpts := client.ProducerOptions{Topic: options.Topic} pp, err := rc.client.CreateProducer(rmqOpts) if err != nil { return nil, err @@ -53,8 +61,8 @@ func (rc *rmqClient) CreateProducer(options ProducerOptions) (Producer, error) { } // CreateReader creates a rocksmq reader from reader options -func (rc *rmqClient) CreateReader(options ReaderOptions) (Reader, error) { - opts := rocksmq.ReaderOptions{ +func (rc *rmqClient) CreateReader(options mqwrapper.ReaderOptions) (mqwrapper.Reader, error) { + opts := client.ReaderOptions{ Topic: options.Topic, StartMessageID: options.StartMessageID.(*rmqID).messageID, StartMessageIDInclusive: options.StartMessageIDInclusive, @@ -72,32 +80,32 @@ func (rc *rmqClient) CreateReader(options ReaderOptions) (Reader, error) { } // Subscribe subscribes a consumer in rmq client -func (rc *rmqClient) Subscribe(options ConsumerOptions) (Consumer, error) { - receiveChannel := make(chan rocksmq.Message, options.BufSize) +func (rc *rmqClient) Subscribe(options mqwrapper.ConsumerOptions) (mqwrapper.Consumer, error) { + receiveChannel := make(chan client.Message, options.BufSize) - cli, err := rc.client.Subscribe(rocksmq.ConsumerOptions{ + cli, err := rc.client.Subscribe(client.ConsumerOptions{ Topic: options.Topic, SubscriptionName: options.SubscriptionName, MessageChannel: receiveChannel, - SubscriptionInitialPosition: rocksmq.SubscriptionInitialPosition(options.SubscriptionInitialPosition), + SubscriptionInitialPosition: client.SubscriptionInitialPosition(options.SubscriptionInitialPosition), }) if err != nil { return nil, err } - rConsumer := &RmqConsumer{c: cli, closeCh: make(chan struct{})} + rConsumer := &Consumer{c: cli, closeCh: make(chan struct{})} return rConsumer, nil } // EarliestMessageID returns the earliest message ID for rmq client -func (rc *rmqClient) EarliestMessageID() MessageID { - rID := rocksmq.EarliestMessageID() +func (rc *rmqClient) EarliestMessageID() mqwrapper.MessageID { + rID := client.EarliestMessageID() return &rmqID{messageID: rID} } // StringToMsgID converts string id to MessageID -func (rc *rmqClient) StringToMsgID(id string) (MessageID, error) { +func (rc *rmqClient) StringToMsgID(id string) (mqwrapper.MessageID, error) { rID, err := strconv.ParseInt(id, 10, 64) if err != nil { return nil, err @@ -106,7 +114,7 @@ func (rc *rmqClient) StringToMsgID(id string) (MessageID, error) { } // BytesToMsgID converts a byte array to messageID -func (rc *rmqClient) BytesToMsgID(id []byte) (MessageID, error) { +func (rc *rmqClient) BytesToMsgID(id []byte) (mqwrapper.MessageID, error) { rID, err := DeserializeRmqID(id) if err != nil { return nil, err diff --git a/internal/util/mqclient/rmq_client_test.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_client_test.go similarity index 76% rename from internal/util/mqclient/rmq_client_test.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_client_test.go index 002edd8922..0671bf3bbb 100644 --- a/internal/util/mqclient/rmq_client_test.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_client_test.go @@ -14,21 +14,24 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package rmq import ( "context" + "fmt" "os" "testing" "time" + rocksmqimplclient "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client" + rocksmqimplserver "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" + "github.com/apache/pulsar-client-go/pulsar" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" + pulsarwrapper "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/pulsar" "github.com/milvus-io/milvus/internal/util/paramtable" "github.com/stretchr/testify/assert" - - "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" - rocksmq1 "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" ) var Params paramtable.BaseTable @@ -38,30 +41,31 @@ func TestMain(m *testing.M) { path := "/tmp/milvus/rdb_data" os.Setenv("ROCKSMQ_PATH", path) defer os.RemoveAll(path) - _ = rocksmq1.InitRocksMQ() + _ = rocksmqimplserver.InitRocksMQ() exitCode := m.Run() - defer rocksmq1.CloseRocksMQ() + defer rocksmqimplserver.CloseRocksMQ() os.Exit(exitCode) } func Test_NewRmqClient(t *testing.T) { - opts := rocksmq.ClientOptions{} - client, err := NewRmqClient(opts) + client, err := createRmqClient() defer client.Close() assert.Nil(t, err) assert.NotNil(t, client) } func TestRmqClient_CreateProducer(t *testing.T) { - opts := rocksmq.ClientOptions{} - client, err := NewRmqClient(opts) + opts := rocksmqimplclient.Options{} + client, err := NewClient(opts) defer client.Close() assert.Nil(t, err) assert.NotNil(t, client) topic := "TestRmqClient_CreateProducer" - proOpts := ProducerOptions{Topic: topic} + proOpts := mqwrapper.ProducerOptions{Topic: topic} producer, err := client.CreateProducer(proOpts) + fmt.Println("===========producer:", producer, err) + defer producer.Close() assert.Nil(t, err) assert.NotNil(t, producer) @@ -70,38 +74,37 @@ func TestRmqClient_CreateProducer(t *testing.T) { defer rmqProducer.Close() assert.Equal(t, rmqProducer.Topic(), topic) - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: []byte{}, Properties: nil, } _, err = rmqProducer.Send(context.TODO(), msg) assert.Nil(t, err) - invalidOpts := ProducerOptions{Topic: ""} + invalidOpts := mqwrapper.ProducerOptions{Topic: ""} producer, e := client.CreateProducer(invalidOpts) assert.Nil(t, producer) assert.Error(t, e) } func TestRmqClient_Subscribe(t *testing.T) { - opts := rocksmq.ClientOptions{} - client, err := NewRmqClient(opts) + client, err := createRmqClient() defer client.Close() assert.Nil(t, err) assert.NotNil(t, client) topic := "TestRmqClient_Subscribe" - proOpts := ProducerOptions{Topic: topic} + proOpts := mqwrapper.ProducerOptions{Topic: topic} producer, err := client.CreateProducer(proOpts) defer producer.Close() assert.Nil(t, err) assert.NotNil(t, producer) subName := "subName" - consumerOpts := ConsumerOptions{ + consumerOpts := mqwrapper.ConsumerOptions{ Topic: "", SubscriptionName: subName, - SubscriptionInitialPosition: SubscriptionPositionEarliest, + SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest, BufSize: 1024, } @@ -116,7 +119,7 @@ func TestRmqClient_Subscribe(t *testing.T) { assert.NotNil(t, consumer) assert.Equal(t, consumer.Subscription(), subName) - msg := &ProducerMessage{ + msg := &mqwrapper.ProducerMessage{ Payload: []byte{1}, Properties: nil, } @@ -148,8 +151,7 @@ func TestRmqClient_Subscribe(t *testing.T) { } func TestRmqClient_EarliestMessageID(t *testing.T) { - opts := rocksmq.ClientOptions{} - client, _ := NewRmqClient(opts) + client, _ := createRmqClient() defer client.Close() mid := client.EarliestMessageID() @@ -157,8 +159,7 @@ func TestRmqClient_EarliestMessageID(t *testing.T) { } func TestRmqClient_StringToMsgID(t *testing.T) { - opts := rocksmq.ClientOptions{} - client, _ := NewRmqClient(opts) + client, _ := createRmqClient() defer client.Close() str := "5" @@ -173,14 +174,18 @@ func TestRmqClient_StringToMsgID(t *testing.T) { } func TestRmqClient_BytesToMsgID(t *testing.T) { - opts := rocksmq.ClientOptions{} - client, _ := NewRmqClient(opts) + client, _ := createRmqClient() defer client.Close() mid := pulsar.EarliestMessageID() - binary := SerializePulsarMsgID(mid) + binary := pulsarwrapper.SerializePulsarMsgID(mid) res, err := client.BytesToMsgID(binary) assert.Nil(t, err) assert.NotNil(t, res) } + +func createRmqClient() (*rmqClient, error) { + opts := rocksmqimplclient.Options{} + return NewClient(opts) +} diff --git a/internal/util/mqclient/rmq_consumer.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_consumer.go similarity index 76% rename from internal/util/mqclient/rmq_consumer.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_consumer.go index bfd680d15b..ab5dd55d6d 100644 --- a/internal/util/mqclient/rmq_consumer.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_consumer.go @@ -14,19 +14,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package rmq import ( "sync" "sync/atomic" - "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client" + + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" ) -// RmqConsumer is a client that used to consume messages from rocksmq -type RmqConsumer struct { - c rocksmq.Consumer - msgChannel chan Message +// Consumer is a client that used to consume messages from rocksmq +type Consumer struct { + c client.Consumer + msgChannel chan mqwrapper.Message closeCh chan struct{} once sync.Once skip int32 @@ -34,15 +36,15 @@ type RmqConsumer struct { } // Subscription returns the subscription name of this consumer -func (rc *RmqConsumer) Subscription() string { +func (rc *Consumer) Subscription() string { return rc.c.Subscription() } // Chan returns a channel to read messages from rocksmq -func (rc *RmqConsumer) Chan() <-chan Message { +func (rc *Consumer) Chan() <-chan mqwrapper.Message { if rc.msgChannel == nil { rc.once.Do(func() { - rc.msgChannel = make(chan Message, 256) + rc.msgChannel = make(chan mqwrapper.Message, 256) rc.wg.Add(1) go func() { defer rc.wg.Done() @@ -71,7 +73,7 @@ func (rc *RmqConsumer) Chan() <-chan Message { } // Seek is used to seek the position in rocksmq topic -func (rc *RmqConsumer) Seek(id MessageID, inclusive bool) error { +func (rc *Consumer) Seek(id mqwrapper.MessageID, inclusive bool) error { msgID := id.(*rmqID).messageID // skip the first message when consume if !inclusive { @@ -81,11 +83,11 @@ func (rc *RmqConsumer) Seek(id MessageID, inclusive bool) error { } // Ack is used to ask a rocksmq message -func (rc *RmqConsumer) Ack(message Message) { +func (rc *Consumer) Ack(message mqwrapper.Message) { } // Close is used to free the resources of this consumer -func (rc *RmqConsumer) Close() { +func (rc *Consumer) Close() { close(rc.closeCh) rc.wg.Wait() } diff --git a/internal/util/mqclient/rmq_id.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_id.go similarity index 89% rename from internal/util/mqclient/rmq_id.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_id.go index a6bece88ad..fb95ae8a2f 100644 --- a/internal/util/mqclient/rmq_id.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_id.go @@ -14,20 +14,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package rmq import ( "github.com/milvus-io/milvus/internal/common" - "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" ) // rmqID wraps message ID for rocksmq type rmqID struct { - messageID rocksmq.UniqueID + messageID server.UniqueID } // Check if rmqID implements MessageID interface -var _ MessageID = &rmqID{} +var _ mqwrapper.MessageID = &rmqID{} // Serialize convert rmq message id to []byte func (rid *rmqID) Serialize() []byte { diff --git a/internal/util/mqclient/rmq_id_test.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_id_test.go similarity index 98% rename from internal/util/mqclient/rmq_id_test.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_id_test.go index d240fc5204..03fa38f748 100644 --- a/internal/util/mqclient/rmq_id_test.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_id_test.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mqclient +package rmq import ( "testing" diff --git a/internal/util/mqclient/rmq_message.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_message.go similarity index 82% rename from internal/util/mqclient/rmq_message.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_message.go index d759fcea83..6748438906 100644 --- a/internal/util/mqclient/rmq_message.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_message.go @@ -9,18 +9,19 @@ // 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 mqclient +package rmq import ( - "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" ) // Check rmqMessage implements ConsumerMessage -var _ Message = (*rmqMessage)(nil) +var _ mqwrapper.Message = (*rmqMessage)(nil) // rmqMessage wraps the message for rocksmq type rmqMessage struct { - msg rocksmq.Message + msg client.Message } // Topic returns the topic name of rocksmq message @@ -39,6 +40,6 @@ func (rm *rmqMessage) Payload() []byte { } // ID returns the id of rocksmq message -func (rm *rmqMessage) ID() MessageID { +func (rm *rmqMessage) ID() mqwrapper.MessageID { return &rmqID{messageID: rm.msg.MsgID} } diff --git a/internal/util/mqclient/rmq_producer.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_producer.go similarity index 73% rename from internal/util/mqclient/rmq_producer.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_producer.go index 7a769409e1..ce85b37810 100644 --- a/internal/util/mqclient/rmq_producer.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_producer.go @@ -9,19 +9,21 @@ // 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 mqclient +package rmq import ( "context" - "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client" + + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" ) -var _ Producer = (*rmqProducer)(nil) +var _ mqwrapper.Producer = (*rmqProducer)(nil) // rmqProducer contains a rocksmq producer type rmqProducer struct { - p rocksmq.Producer + p client.Producer } // Topic returns the topic of rmq producer @@ -30,8 +32,8 @@ func (rp *rmqProducer) Topic() string { } // Send send the producer messages to rocksmq -func (rp *rmqProducer) Send(ctx context.Context, message *ProducerMessage) (MessageID, error) { - pm := &rocksmq.ProducerMessage{Payload: message.Payload} +func (rp *rmqProducer) Send(ctx context.Context, message *mqwrapper.ProducerMessage) (mqwrapper.MessageID, error) { + pm := &client.ProducerMessage{Payload: message.Payload} id, err := rp.p.Send(pm) return &rmqID{messageID: id}, err } diff --git a/internal/util/mqclient/rmq_reader.go b/internal/mq/msgstream/mqwrapper/rmq/rmq_reader.go similarity index 67% rename from internal/util/mqclient/rmq_reader.go rename to internal/mq/msgstream/mqwrapper/rmq/rmq_reader.go index 148678584d..16fd8728ed 100644 --- a/internal/util/mqclient/rmq_reader.go +++ b/internal/mq/msgstream/mqwrapper/rmq/rmq_reader.go @@ -1,16 +1,18 @@ -package mqclient +package rmq import ( "context" - "github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq" + "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client" + + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" ) -var _ Reader = (*rmqReader)(nil) +var _ mqwrapper.Reader = (*rmqReader)(nil) // rmqReader contains a rocksmq reader type rmqReader struct { - r rocksmq.Reader + r client.Reader } // Topic returns the topic name of a reader @@ -19,7 +21,7 @@ func (rr *rmqReader) Topic() string { } // Next returns the next message of reader, blocking until a message is available -func (rr *rmqReader) Next(ctx context.Context) (Message, error) { +func (rr *rmqReader) Next(ctx context.Context) (mqwrapper.Message, error) { rMsg, err := rr.r.Next(ctx) if err != nil { return nil, err @@ -34,7 +36,7 @@ func (rr *rmqReader) HasNext() bool { } // Seek seeks the reader position to id -func (rr *rmqReader) Seek(id MessageID) error { +func (rr *rmqReader) Seek(id mqwrapper.MessageID) error { msgID := id.(*rmqID).messageID return rr.r.Seek(msgID) } diff --git a/internal/msgstream/msg.go b/internal/mq/msgstream/msg.go similarity index 100% rename from internal/msgstream/msg.go rename to internal/mq/msgstream/msg.go diff --git a/internal/msgstream/msg_test.go b/internal/mq/msgstream/msg_test.go similarity index 100% rename from internal/msgstream/msg_test.go rename to internal/mq/msgstream/msg_test.go diff --git a/internal/msgstream/msgstream.go b/internal/mq/msgstream/msgstream.go similarity index 95% rename from internal/msgstream/msgstream.go rename to internal/mq/msgstream/msgstream.go index 3b815a8974..5303d5bffb 100644 --- a/internal/msgstream/msgstream.go +++ b/internal/mq/msgstream/msgstream.go @@ -19,8 +19,8 @@ package msgstream import ( "context" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/milvus-io/milvus/internal/util/typeutil" ) @@ -37,7 +37,7 @@ type IntPrimaryKey = typeutil.IntPrimaryKey type MsgPosition = internalpb.MsgPosition // MessageID is an alias for short -type MessageID = mqclient.MessageID +type MessageID = mqwrapper.MessageID // MsgPack represents a batch of msg in msgstream type MsgPack struct { @@ -59,7 +59,7 @@ type MsgStream interface { AsProducer(channels []string) AsConsumer(channels []string, subName string) AsReader(channels []string, subName string) - AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) + AsConsumerWithPosition(channels []string, subName string, position mqwrapper.SubscriptionInitialPosition) SetRepackFunc(repackFunc RepackFunc) ComputeProduceChannelIndexes(tsMsgs []TsMsg) [][]int32 GetProduceChannels() []string diff --git a/internal/msgstream/repack_func.go b/internal/mq/msgstream/repack_func.go similarity index 100% rename from internal/msgstream/repack_func.go rename to internal/mq/msgstream/repack_func.go diff --git a/internal/msgstream/trace.go b/internal/mq/msgstream/trace.go similarity index 100% rename from internal/msgstream/trace.go rename to internal/mq/msgstream/trace.go diff --git a/internal/msgstream/unmarshal.go b/internal/mq/msgstream/unmarshal.go similarity index 100% rename from internal/msgstream/unmarshal.go rename to internal/mq/msgstream/unmarshal.go diff --git a/internal/msgstream/unmarshal_test.go b/internal/mq/msgstream/unmarshal_test.go similarity index 100% rename from internal/msgstream/unmarshal_test.go rename to internal/mq/msgstream/unmarshal_test.go diff --git a/internal/proxy/channels_mgr.go b/internal/proxy/channels_mgr.go index 106b8912f1..7cd1ee4a3f 100644 --- a/internal/proxy/channels_mgr.go +++ b/internal/proxy/channels_mgr.go @@ -27,7 +27,7 @@ import ( "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" "github.com/milvus-io/milvus/internal/proto/querypb" diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go index ca80aef290..62bb6ebecd 100644 --- a/internal/proxy/impl.go +++ b/internal/proxy/impl.go @@ -30,7 +30,7 @@ import ( "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/proxy/mock_test.go b/internal/proxy/mock_test.go index 526aa41737..253cc6af7c 100644 --- a/internal/proxy/mock_test.go +++ b/internal/proxy/mock_test.go @@ -22,13 +22,13 @@ import ( "sync" "time" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/rootcoordpb" "github.com/milvus-io/milvus/internal/proto/schemapb" "github.com/milvus-io/milvus/internal/util/funcutil" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/milvus-io/milvus/internal/util/uniquegenerator" ) @@ -292,7 +292,7 @@ func (ms *simpleMockMsgStream) HasNext(channelName string) bool { return true } -func (ms *simpleMockMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) { +func (ms *simpleMockMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqwrapper.SubscriptionInitialPosition) { } func (ms *simpleMockMsgStream) ComputeProduceChannelIndexes(tsMsgs []msgstream.TsMsg) [][]int32 { diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index b15cba620d..5b0584a832 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -29,7 +29,7 @@ import ( "github.com/milvus-io/milvus/internal/allocator" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/querypb" diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index dd150d6b3e..738d411a50 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -95,7 +95,7 @@ import ( "github.com/milvus-io/milvus/internal/querycoord" "github.com/stretchr/testify/assert" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" diff --git a/internal/proxy/repack_func.go b/internal/proxy/repack_func.go index fce3cf8507..5db22bb479 100644 --- a/internal/proxy/repack_func.go +++ b/internal/proxy/repack_func.go @@ -19,7 +19,7 @@ package proxy import ( "fmt" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // insertRepackFunc deprecated, use defaultInsertRepackFunc instead. diff --git a/internal/proxy/repack_func_test.go b/internal/proxy/repack_func_test.go index 72b91fa056..131583430f 100644 --- a/internal/proxy/repack_func_test.go +++ b/internal/proxy/repack_func_test.go @@ -20,7 +20,7 @@ import ( "math/rand" "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/stretchr/testify/assert" ) diff --git a/internal/proxy/task.go b/internal/proxy/task.go index 352621cbb4..0772292224 100644 --- a/internal/proxy/task.go +++ b/internal/proxy/task.go @@ -38,7 +38,7 @@ import ( "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/indexpb" diff --git a/internal/proxy/task_scheduler.go b/internal/proxy/task_scheduler.go index 7c6024e07e..0f991cd95b 100644 --- a/internal/proxy/task_scheduler.go +++ b/internal/proxy/task_scheduler.go @@ -31,10 +31,10 @@ import ( "go.uber.org/zap" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/milvus-io/milvus/internal/util/trace" "github.com/opentracing/opentracing-go" oplog "github.com/opentracing/opentracing-go/log" @@ -818,7 +818,7 @@ func (sched *taskScheduler) collectResultLoop() { queryResultMsgStream, _ := sched.msFactory.NewQueryMsgStream(sched.ctx) // proxy didn't need to walk through all the search results in channel, because it no longer has client connections. consumeSubName := fmt.Sprintf("%s-%d", Params.MsgChannelCfg.ProxySubName, Params.ProxyCfg.ProxyID) - queryResultMsgStream.AsConsumerWithPosition(Params.ProxyCfg.SearchResultChannelNames, consumeSubName, mqclient.SubscriptionPositionLatest) + queryResultMsgStream.AsConsumerWithPosition(Params.ProxyCfg.SearchResultChannelNames, consumeSubName, mqwrapper.SubscriptionPositionLatest) log.Debug("Proxy", zap.Strings("SearchResultChannelNames", Params.ProxyCfg.SearchResultChannelNames), zap.Any("consumeSubName", consumeSubName)) diff --git a/internal/proxy/task_test.go b/internal/proxy/task_test.go index b36f4fa2c2..759ebbebe4 100644 --- a/internal/proxy/task_test.go +++ b/internal/proxy/task_test.go @@ -37,7 +37,7 @@ import ( "github.com/milvus-io/milvus/internal/allocator" "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" diff --git a/internal/querycoord/channel_unsubscribe.go b/internal/querycoord/channel_unsubscribe.go index 9b6cad9497..fdedd12add 100644 --- a/internal/querycoord/channel_unsubscribe.go +++ b/internal/querycoord/channel_unsubscribe.go @@ -28,7 +28,7 @@ import ( etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/internal/util/funcutil" ) diff --git a/internal/querycoord/channel_unsubscribe_test.go b/internal/querycoord/channel_unsubscribe_test.go index 4d7dd13a5c..ef90821141 100644 --- a/internal/querycoord/channel_unsubscribe_test.go +++ b/internal/querycoord/channel_unsubscribe_test.go @@ -25,7 +25,7 @@ import ( "github.com/stretchr/testify/assert" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/internal/util/etcd" ) diff --git a/internal/querycoord/cluster_test.go b/internal/querycoord/cluster_test.go index 8d0c47df7e..d1a8fa01e5 100644 --- a/internal/querycoord/cluster_test.go +++ b/internal/querycoord/cluster_test.go @@ -32,7 +32,7 @@ import ( "github.com/milvus-io/milvus/internal/kv" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/etcdpb" diff --git a/internal/querycoord/meta.go b/internal/querycoord/meta.go index a64b87b211..60ce145629 100644 --- a/internal/querycoord/meta.go +++ b/internal/querycoord/meta.go @@ -31,7 +31,7 @@ import ( "github.com/milvus-io/milvus/internal/kv" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/querycoord/query_coord.go b/internal/querycoord/query_coord.go index ebdad00315..d9c75cf073 100644 --- a/internal/querycoord/query_coord.go +++ b/internal/querycoord/query_coord.go @@ -33,7 +33,7 @@ import ( "github.com/milvus-io/milvus/internal/allocator" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/querypb" diff --git a/internal/querycoord/query_coord_test.go b/internal/querycoord/query_coord_test.go index 57c5a0901c..90253f7343 100644 --- a/internal/querycoord/query_coord_test.go +++ b/internal/querycoord/query_coord_test.go @@ -33,7 +33,7 @@ import ( etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/internal/util/etcd" diff --git a/internal/querycoord/segment_allocator_test.go b/internal/querycoord/segment_allocator_test.go index c8c60f7132..a4fcf4b432 100644 --- a/internal/querycoord/segment_allocator_test.go +++ b/internal/querycoord/segment_allocator_test.go @@ -23,7 +23,7 @@ import ( "github.com/stretchr/testify/assert" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/internal/util/etcd" "github.com/milvus-io/milvus/internal/util/sessionutil" diff --git a/internal/querynode/data_sync_service.go b/internal/querynode/data_sync_service.go index 0fb31942a2..848c26fd48 100644 --- a/internal/querynode/data_sync_service.go +++ b/internal/querynode/data_sync_service.go @@ -25,7 +25,7 @@ import ( "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // dataSyncService manages a lot of flow graphs diff --git a/internal/querynode/flow_graph_delete_node_test.go b/internal/querynode/flow_graph_delete_node_test.go index 5c4c7a8d51..f9e91444d7 100644 --- a/internal/querynode/flow_graph_delete_node_test.go +++ b/internal/querynode/flow_graph_delete_node_test.go @@ -23,7 +23,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/milvus-io/milvus/internal/common" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/util/flowgraph" ) diff --git a/internal/querynode/flow_graph_filter_delete_node.go b/internal/querynode/flow_graph_filter_delete_node.go index 8dcdff2211..b5a38d986b 100644 --- a/internal/querynode/flow_graph_filter_delete_node.go +++ b/internal/querynode/flow_graph_filter_delete_node.go @@ -23,7 +23,7 @@ import ( "go.uber.org/zap" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/util/flowgraph" "github.com/milvus-io/milvus/internal/util/trace" diff --git a/internal/querynode/flow_graph_filter_delete_node_test.go b/internal/querynode/flow_graph_filter_delete_node_test.go index 254baf8e30..a265df5d7b 100644 --- a/internal/querynode/flow_graph_filter_delete_node_test.go +++ b/internal/querynode/flow_graph_filter_delete_node_test.go @@ -22,7 +22,7 @@ import ( "github.com/stretchr/testify/assert" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/util/flowgraph" ) diff --git a/internal/querynode/flow_graph_filter_dm_node.go b/internal/querynode/flow_graph_filter_dm_node.go index 9e4ea0eb2f..58852806dc 100644 --- a/internal/querynode/flow_graph_filter_dm_node.go +++ b/internal/querynode/flow_graph_filter_dm_node.go @@ -23,7 +23,7 @@ import ( "go.uber.org/zap" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/util/flowgraph" "github.com/milvus-io/milvus/internal/util/trace" diff --git a/internal/querynode/flow_graph_filter_dm_node_test.go b/internal/querynode/flow_graph_filter_dm_node_test.go index a1990dc477..ee69d1a261 100644 --- a/internal/querynode/flow_graph_filter_dm_node_test.go +++ b/internal/querynode/flow_graph_filter_dm_node_test.go @@ -22,7 +22,7 @@ import ( "github.com/stretchr/testify/assert" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/querynode/flow_graph_insert_node.go b/internal/querynode/flow_graph_insert_node.go index c36c4f0fbf..aa9c974266 100644 --- a/internal/querynode/flow_graph_insert_node.go +++ b/internal/querynode/flow_graph_insert_node.go @@ -30,7 +30,7 @@ import ( "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/schemapb" "github.com/milvus-io/milvus/internal/util/flowgraph" diff --git a/internal/querynode/flow_graph_insert_node_test.go b/internal/querynode/flow_graph_insert_node_test.go index 8aeddc1ff2..f142fa06a1 100644 --- a/internal/querynode/flow_graph_insert_node_test.go +++ b/internal/querynode/flow_graph_insert_node_test.go @@ -24,7 +24,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/milvus-io/milvus/internal/common" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/util/flowgraph" ) diff --git a/internal/querynode/flow_graph_message.go b/internal/querynode/flow_graph_message.go index 8670d74166..c9ee444830 100644 --- a/internal/querynode/flow_graph_message.go +++ b/internal/querynode/flow_graph_message.go @@ -17,7 +17,7 @@ package querynode import ( - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/util/flowgraph" ) diff --git a/internal/querynode/flow_graph_message_test.go b/internal/querynode/flow_graph_message_test.go index 50dc61ae94..7acb8b6b75 100644 --- a/internal/querynode/flow_graph_message_test.go +++ b/internal/querynode/flow_graph_message_test.go @@ -21,7 +21,7 @@ import ( "github.com/stretchr/testify/assert" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) func TestFlowGraphMsg_insertMsg(t *testing.T) { diff --git a/internal/querynode/flow_graph_query_node.go b/internal/querynode/flow_graph_query_node.go index 167ff80f1b..dc6150feb5 100644 --- a/internal/querynode/flow_graph_query_node.go +++ b/internal/querynode/flow_graph_query_node.go @@ -25,10 +25,10 @@ import ( "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/util/flowgraph" - "github.com/milvus-io/milvus/internal/util/mqclient" ) // queryNodeFlowGraph is a TimeTickedFlowGraph in query node @@ -212,7 +212,7 @@ func (q *queryNodeFlowGraph) consumerFlowGraphLatest(channel Channel, subName Co if q.dmlStream == nil { return errors.New("null dml message stream in flow graph") } - q.dmlStream.AsConsumerWithPosition([]string{channel}, subName, mqclient.SubscriptionPositionLatest) + q.dmlStream.AsConsumerWithPosition([]string{channel}, subName, mqwrapper.SubscriptionPositionLatest) log.Debug("query node flow graph consumes from pChannel", zap.Any("collectionID", q.collectionID), zap.Any("channel", channel), diff --git a/internal/querynode/mock_test.go b/internal/querynode/mock_test.go index 5039be9829..364a16c713 100644 --- a/internal/querynode/mock_test.go +++ b/internal/querynode/mock_test.go @@ -31,7 +31,7 @@ import ( etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" minioKV "github.com/milvus-io/milvus/internal/kv/minio" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/etcdpb" diff --git a/internal/querynode/query_collection.go b/internal/querynode/query_collection.go index 0719ff082a..5431cca50f 100644 --- a/internal/querynode/query_collection.go +++ b/internal/querynode/query_collection.go @@ -32,7 +32,7 @@ import ( "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/etcdpb" diff --git a/internal/querynode/query_collection_test.go b/internal/querynode/query_collection_test.go index a1fe1976b8..6dfe643db1 100644 --- a/internal/querynode/query_collection_test.go +++ b/internal/querynode/query_collection_test.go @@ -33,7 +33,7 @@ import ( "github.com/milvus-io/milvus/internal/common" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" diff --git a/internal/querynode/query_node.go b/internal/querynode/query_node.go index a0ce55dc86..559b448f8e 100644 --- a/internal/querynode/query_node.go +++ b/internal/querynode/query_node.go @@ -46,7 +46,7 @@ import ( "github.com/milvus-io/milvus/internal/kv" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/internal/types" diff --git a/internal/querynode/query_node_test.go b/internal/querynode/query_node_test.go index bf1767ac9f..5879a25336 100644 --- a/internal/querynode/query_node_test.go +++ b/internal/querynode/query_node_test.go @@ -31,7 +31,7 @@ import ( "github.com/stretchr/testify/assert" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/etcdpb" "github.com/milvus-io/milvus/internal/proto/querypb" diff --git a/internal/querynode/query_service.go b/internal/querynode/query_service.go index 252c94a89a..305114c7ff 100644 --- a/internal/querynode/query_service.go +++ b/internal/querynode/query_service.go @@ -27,7 +27,7 @@ import ( "go.uber.org/zap" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/storage" ) diff --git a/internal/querynode/query_service_test.go b/internal/querynode/query_service_test.go index 1a2a611af8..77a6af5e46 100644 --- a/internal/querynode/query_service_test.go +++ b/internal/querynode/query_service_test.go @@ -27,7 +27,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/milvus-io/milvus/internal/common" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" diff --git a/internal/querynode/segment_loader.go b/internal/querynode/segment_loader.go index 940ccad9c0..21dd245c28 100644 --- a/internal/querynode/segment_loader.go +++ b/internal/querynode/segment_loader.go @@ -33,7 +33,7 @@ import ( minioKV "github.com/milvus-io/milvus/internal/kv/minio" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/querynode/stats_service.go b/internal/querynode/stats_service.go index 17a47fd521..d8b2e01234 100644 --- a/internal/querynode/stats_service.go +++ b/internal/querynode/stats_service.go @@ -21,7 +21,7 @@ import ( "time" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "go.uber.org/zap" diff --git a/internal/querynode/stats_service_test.go b/internal/querynode/stats_service_test.go index 10f4b30bc2..687a23d365 100644 --- a/internal/querynode/stats_service_test.go +++ b/internal/querynode/stats_service_test.go @@ -20,7 +20,7 @@ import ( "context" "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/stretchr/testify/assert" ) diff --git a/internal/querynode/streaming.go b/internal/querynode/streaming.go index 51a78b385f..8d2208109c 100644 --- a/internal/querynode/streaming.go +++ b/internal/querynode/streaming.go @@ -27,7 +27,7 @@ import ( etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/segcorepb" "github.com/milvus-io/milvus/internal/util/timerecord" ) diff --git a/internal/querynode/task_test.go b/internal/querynode/task_test.go index 34cda3b0d1..8e90d0f9c2 100644 --- a/internal/querynode/task_test.go +++ b/internal/querynode/task_test.go @@ -23,7 +23,7 @@ import ( "github.com/apache/pulsar-client-go/pulsar" "github.com/stretchr/testify/assert" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" diff --git a/internal/rootcoord/dml_channels.go b/internal/rootcoord/dml_channels.go index e996d475f3..04b37b2652 100644 --- a/internal/rootcoord/dml_channels.go +++ b/internal/rootcoord/dml_channels.go @@ -27,7 +27,7 @@ import ( "go.uber.org/zap" "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) type dmlMsgStream struct { diff --git a/internal/rootcoord/dml_channels_test.go b/internal/rootcoord/dml_channels_test.go index 086a33acd8..da476093dc 100644 --- a/internal/rootcoord/dml_channels_test.go +++ b/internal/rootcoord/dml_channels_test.go @@ -22,9 +22,9 @@ import ( "sync" "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper" "github.com/milvus-io/milvus/internal/util/funcutil" - "github.com/milvus-io/milvus/internal/util/mqclient" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -138,7 +138,7 @@ func (ms *FailMsgStream) Chan() <-chan *msgstream.MsgPack { return func (ms *FailMsgStream) AsProducer(channels []string) {} func (ms *FailMsgStream) AsConsumer(channels []string, subName string) {} func (ms *FailMsgStream) AsReader(channels []string, subName string) {} -func (ms *FailMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) { +func (ms *FailMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqwrapper.SubscriptionInitialPosition) { } func (ms *FailMsgStream) SetRepackFunc(repackFunc msgstream.RepackFunc) {} func (ms *FailMsgStream) ComputeProduceChannelIndexes(tsMsgs []msgstream.TsMsg) [][]int32 { return nil } diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go index cb096eb8c7..22e7a5e80c 100644 --- a/internal/rootcoord/root_coord.go +++ b/internal/rootcoord/root_coord.go @@ -37,7 +37,7 @@ import ( etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - ms "github.com/milvus-io/milvus/internal/msgstream" + ms "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/etcdpb" diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go index 3f0fc53fee..aa87590d23 100644 --- a/internal/rootcoord/root_coord_test.go +++ b/internal/rootcoord/root_coord_test.go @@ -32,7 +32,7 @@ import ( "github.com/milvus-io/milvus/internal/kv" etcdkv "github.com/milvus-io/milvus/internal/kv/etcd" memkv "github.com/milvus-io/milvus/internal/kv/mem" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/util/metricsinfo" diff --git a/internal/rootcoord/timestamp_test.go b/internal/rootcoord/timestamp_test.go index dc3816e9e7..6eabd97ec9 100644 --- a/internal/rootcoord/timestamp_test.go +++ b/internal/rootcoord/timestamp_test.go @@ -23,7 +23,7 @@ import ( "sync" "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/indexpb" diff --git a/internal/rootcoord/timeticksync.go b/internal/rootcoord/timeticksync.go index 6f206564a2..1da67efc98 100644 --- a/internal/rootcoord/timeticksync.go +++ b/internal/rootcoord/timeticksync.go @@ -25,7 +25,7 @@ import ( "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/metrics" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/util/sessionutil" diff --git a/internal/rootcoord/timeticksync_test.go b/internal/rootcoord/timeticksync_test.go index 3d56a74374..c72b6169d6 100644 --- a/internal/rootcoord/timeticksync_test.go +++ b/internal/rootcoord/timeticksync_test.go @@ -23,7 +23,7 @@ import ( "github.com/stretchr/testify/assert" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" ) diff --git a/internal/rootcoord/util.go b/internal/rootcoord/util.go index 63350b7a4e..3c49de1ab8 100644 --- a/internal/rootcoord/util.go +++ b/internal/rootcoord/util.go @@ -21,7 +21,7 @@ import ( "fmt" "github.com/golang/protobuf/proto" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/etcdpb" "github.com/milvus-io/milvus/internal/proto/schemapb" diff --git a/internal/rootcoord/util_test.go b/internal/rootcoord/util_test.go index db4256cc85..1dd4d1fa33 100644 --- a/internal/rootcoord/util_test.go +++ b/internal/rootcoord/util_test.go @@ -19,7 +19,7 @@ package rootcoord import ( "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/etcdpb" "github.com/milvus-io/milvus/internal/proto/schemapb" diff --git a/internal/util/flowgraph/input_node.go b/internal/util/flowgraph/input_node.go index a627c10835..6c391510ab 100644 --- a/internal/util/flowgraph/input_node.go +++ b/internal/util/flowgraph/input_node.go @@ -18,7 +18,7 @@ package flowgraph import ( "github.com/milvus-io/milvus/internal/log" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/util/trace" "github.com/opentracing/opentracing-go" oplog "github.com/opentracing/opentracing-go/log" diff --git a/internal/util/flowgraph/input_node_test.go b/internal/util/flowgraph/input_node_test.go index a8310a0e96..f1747af3a9 100644 --- a/internal/util/flowgraph/input_node_test.go +++ b/internal/util/flowgraph/input_node_test.go @@ -21,7 +21,7 @@ import ( "os" "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/stretchr/testify/assert" ) diff --git a/internal/util/flowgraph/message.go b/internal/util/flowgraph/message.go index bf37c76f2a..83e72416ef 100644 --- a/internal/util/flowgraph/message.go +++ b/internal/util/flowgraph/message.go @@ -17,7 +17,7 @@ package flowgraph import ( - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" ) // Msg is an abstract class that contains a method to get the time tick of this message diff --git a/internal/util/flowgraph/message_test.go b/internal/util/flowgraph/message_test.go index a148b8b5ed..8dc3269eff 100644 --- a/internal/util/flowgraph/message_test.go +++ b/internal/util/flowgraph/message_test.go @@ -20,7 +20,7 @@ import ( "context" "testing" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/stretchr/testify/assert" ) diff --git a/internal/util/flowgraph/node_test.go b/internal/util/flowgraph/node_test.go index 49ca990a25..4984c9d3ba 100644 --- a/internal/util/flowgraph/node_test.go +++ b/internal/util/flowgraph/node_test.go @@ -24,7 +24,7 @@ import ( "testing" "time" - "github.com/milvus-io/milvus/internal/msgstream" + "github.com/milvus-io/milvus/internal/mq/msgstream" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/stretchr/testify/assert" diff --git a/internal/util/paramtable/base_table.go b/internal/util/paramtable/base_table.go index e408dcbd89..e2c3ff7e6b 100644 --- a/internal/util/paramtable/base_table.go +++ b/internal/util/paramtable/base_table.go @@ -79,7 +79,7 @@ func (gp *BaseTable) Init() { gp.loadFromMilvusYaml() - gp.tryloadFromEnv() + gp.tryLoadFromEnv() gp.InitLogCfg() } @@ -124,80 +124,12 @@ func (gp *BaseTable) loadFromMilvusYaml() { } } -func (gp *BaseTable) tryloadFromEnv() { - // minio - minioAddress := os.Getenv("MINIO_ADDRESS") - if minioAddress == "" { - minioHost := gp.LoadWithDefault("minio.address", DefaultMinioHost) - port := gp.LoadWithDefault("minio.port", DefaultMinioPort) - minioAddress = minioHost + ":" + port - } - gp.Save("_MinioAddress", minioAddress) - - // etcd - etcdEndpoints := os.Getenv("ETCD_ENDPOINTS") - if etcdEndpoints == "" { - etcdEndpoints = gp.LoadWithDefault("etcd.endpoints", DefaultEtcdEndpoints) - } - gp.Save("_EtcdEndpoints", etcdEndpoints) - - // pulsar - pulsarAddress := os.Getenv("PULSAR_ADDRESS") - if pulsarAddress == "" { - pulsarHost := gp.LoadWithDefault("pulsar.address", DefaultPulsarHost) - port := gp.LoadWithDefault("pulsar.port", DefaultPulsarPort) - pulsarAddress = "pulsar://" + pulsarHost + ":" + port - } - gp.Save("_PulsarAddress", pulsarAddress) - - // rocksmq - rocksmqPath := os.Getenv("ROCKSMQ_PATH") - if rocksmqPath == "" { - rocksmqPath = gp.LoadWithDefault("rocksmq.path", DefaultRocksmqPath) - } - gp.Save("_RocksmqPath", rocksmqPath) - - insertBufferFlushSize := os.Getenv("DATA_NODE_IBUFSIZE") - if insertBufferFlushSize == "" { - insertBufferFlushSize = gp.LoadWithDefault("datanode.flush.insertBufSize", DefaultInsertBufferSize) - } - gp.Save("_DATANODE_INSERTBUFSIZE", insertBufferFlushSize) - - minioAccessKey := os.Getenv("MINIO_ACCESS_KEY") - if minioAccessKey == "" { - minioAccessKey = gp.LoadWithDefault("minio.accessKeyID", DefaultMinioAccessKey) - } - gp.Save("_MinioAccessKeyID", minioAccessKey) - - minioSecretKey := os.Getenv("MINIO_SECRET_KEY") - if minioSecretKey == "" { - minioSecretKey = gp.LoadWithDefault("minio.secretAccessKey", DefaultMinioSecretAccessKey) - } - gp.Save("_MinioSecretAccessKey", minioSecretKey) - - minioUseSSL := os.Getenv("MINIO_USE_SSL") - if minioUseSSL == "" { - minioUseSSL = gp.LoadWithDefault("minio.useSSL", DefaultMinioUseSSL) - } - gp.Save("_MinioUseSSL", minioUseSSL) - - minioBucketName := os.Getenv("MINIO_BUCKET_NAME") - if minioBucketName == "" { - minioBucketName = gp.LoadWithDefault("minio.bucketName", DefaultMinioBucketName) - } - gp.Save("_MinioBucketName", minioBucketName) - - // try to load environment start with ENV_PREFIX - for _, e := range os.Environ() { - parts := strings.SplitN(e, "=", 2) - if strings.Contains(parts[0], DefaultEnvPrefix) { - parts := strings.SplitN(e, "=", 2) - // remove the ENV PREFIX and use the rest as key - keyParts := strings.SplitAfterN(parts[0], ".", 2) - // mem kv throw no errors - gp.Save(keyParts[1], parts[1]) - } - } +func (gp *BaseTable) tryLoadFromEnv() { + gp.loadEtcdConfig() + gp.loadMinioConfig() + gp.loadMQConfig() + gp.loadDataNodeConfig() + gp.loadOtherEnvs() } // Load loads an object with @key. @@ -449,3 +381,91 @@ func (gp *BaseTable) SetLogger(id UniqueID) { gp.LogCfgFunc(gp.Log) } } + +func (gp *BaseTable) loadPulsarConfig() { + pulsarAddress := os.Getenv("PULSAR_ADDRESS") + if pulsarAddress == "" { + pulsarHost := gp.LoadWithDefault("pulsar.address", DefaultPulsarHost) + port := gp.LoadWithDefault("pulsar.port", DefaultPulsarPort) + pulsarAddress = "pulsar://" + pulsarHost + ":" + port + } + + gp.Save("_PulsarAddress", pulsarAddress) +} + +func (gp *BaseTable) loadRocksMQConfig() { + rocksmqPath := os.Getenv("ROCKSMQ_PATH") + if rocksmqPath == "" { + rocksmqPath = gp.LoadWithDefault("rocksmq.path", DefaultRocksmqPath) + } + gp.Save("_RocksmqPath", rocksmqPath) +} + +func (gp *BaseTable) loadMQConfig() { + gp.loadPulsarConfig() + gp.loadRocksMQConfig() +} + +func (gp *BaseTable) loadEtcdConfig() { + etcdEndpoints := os.Getenv("ETCD_ENDPOINTS") + if etcdEndpoints == "" { + etcdEndpoints = gp.LoadWithDefault("etcd.endpoints", DefaultEtcdEndpoints) + } + gp.Save("_EtcdEndpoints", etcdEndpoints) +} + +func (gp *BaseTable) loadMinioConfig() { + minioAddress := os.Getenv("MINIO_ADDRESS") + if minioAddress == "" { + minioHost := gp.LoadWithDefault("minio.address", DefaultMinioHost) + port := gp.LoadWithDefault("minio.port", DefaultMinioPort) + minioAddress = minioHost + ":" + port + } + gp.Save("_MinioAddress", minioAddress) + + minioAccessKey := os.Getenv("MINIO_ACCESS_KEY") + if minioAccessKey == "" { + minioAccessKey = gp.LoadWithDefault("minio.accessKeyID", DefaultMinioAccessKey) + } + gp.Save("_MinioAccessKeyID", minioAccessKey) + + minioSecretKey := os.Getenv("MINIO_SECRET_KEY") + if minioSecretKey == "" { + minioSecretKey = gp.LoadWithDefault("minio.secretAccessKey", DefaultMinioSecretAccessKey) + } + gp.Save("_MinioSecretAccessKey", minioSecretKey) + + minioUseSSL := os.Getenv("MINIO_USE_SSL") + if minioUseSSL == "" { + minioUseSSL = gp.LoadWithDefault("minio.useSSL", DefaultMinioUseSSL) + } + gp.Save("_MinioUseSSL", minioUseSSL) + + minioBucketName := os.Getenv("MINIO_BUCKET_NAME") + if minioBucketName == "" { + minioBucketName = gp.LoadWithDefault("minio.bucketName", DefaultMinioBucketName) + } + gp.Save("_MinioBucketName", minioBucketName) +} + +func (gp *BaseTable) loadDataNodeConfig() { + insertBufferFlushSize := os.Getenv("DATA_NODE_IBUFSIZE") + if insertBufferFlushSize == "" { + insertBufferFlushSize = gp.LoadWithDefault("datanode.flush.insertBufSize", DefaultInsertBufferSize) + } + gp.Save("_DATANODE_INSERTBUFSIZE", insertBufferFlushSize) +} + +func (gp *BaseTable) loadOtherEnvs() { + // try to load environment start with ENV_PREFIX + for _, e := range os.Environ() { + parts := strings.SplitN(e, "=", 2) + if strings.Contains(parts[0], DefaultEnvPrefix) { + parts := strings.SplitN(e, "=", 2) + // remove the ENV PREFIX and use the rest as key + keyParts := strings.SplitAfterN(parts[0], ".", 2) + // mem kv throw no errors + gp.Save(keyParts[1], parts[1]) + } + } +} diff --git a/scripts/run_go_unittest.sh b/scripts/run_go_unittest.sh index 97fa154c89..d13d1073eb 100755 --- a/scripts/run_go_unittest.sh +++ b/scripts/run_go_unittest.sh @@ -38,14 +38,12 @@ echo "Running go unittest under $MILVUS_DIR" go test -race -cover "${MILVUS_DIR}/allocator/..." -failfast go test -race -cover "${MILVUS_DIR}/kv/..." -failfast -go test -race -cover "${MILVUS_DIR}/msgstream/..." -failfast +go test -race -cover "${MILVUS_DIR}/mq/..." -failfast go test -race -cover "${MILVUS_DIR}/storage" -failfast go test -race -cover "${MILVUS_DIR}/tso/..." -failfast go test -race -cover "${MILVUS_DIR}/util/funcutil/..." -failfast -go test -race -cover "${MILVUS_DIR}/util/mqclient/..." -failfast go test -race -cover "${MILVUS_DIR}/util/paramtable/..." -failfast go test -race -cover "${MILVUS_DIR}/util/retry/..." -failfast -go test -race -cover "${MILVUS_DIR}/util/rocksmq/..." -failfast go test -race -cover "${MILVUS_DIR}/util/sessionutil/..." -failfast go test -race -cover "${MILVUS_DIR}/util/trace/..." -failfast go test -race -cover "${MILVUS_DIR}/util/typeutil/..." -failfast