milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Reconstruct mqstream module (#15784) 

Signed-off-by: yun.zhang <yun.zhang@zilliz.com>
pull/15843/head
jaime 2022-03-03 21:57:56 +08:00 committed by GitHub
parent 54862a31b0
commit 7a3a721380
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
145 changed files with 582 additions and 508 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
cmd/components/data_coord.go
View File

@ -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

2
cmd/components/data_node.go
View File

@ -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

2
cmd/components/proxy.go
View File

@ -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

2
cmd/components/query_coord.go
View File

@ -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

2
cmd/components/query_node.go
View File

@ -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

2
cmd/components/root_coord.go
View File

@ -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"
)

9
cmd/roles/roles.go
View File

@ -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.

2
internal/datacoord/channel_manager.go
View File

@ -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"

6
internal/datacoord/server.go
View File

@ -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))

2
internal/datacoord/server_test.go
View File

@ -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"

2
internal/datanode/data_node.go
View File

@ -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"

2
internal/datanode/data_node_test.go
View File

@ -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"

2
internal/datanode/data_sync_service.go
View File

@ -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"

2
internal/datanode/data_sync_service_test.go
View File

@ -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"

2
internal/datanode/flow_graph_dd_node.go
View File

@ -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"

2
internal/datanode/flow_graph_dd_node_test.go
View File

@ -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"

2
internal/datanode/flow_graph_delete_node.go
View File

@ -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"

2
internal/datanode/flow_graph_delete_node_test.go
View File

@ -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"

6
internal/datanode/flow_graph_dmstream_input_node_test.go
View File

@ -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 {

2
internal/datanode/flow_graph_insert_buffer_node.go
View File

@ -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"

2
internal/datanode/flow_graph_insert_buffer_node_test.go
View File

@ -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"

2
internal/datanode/flow_graph_message.go
View File

@ -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"
)

2
internal/datanode/mock_test.go
View File

@ -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"

2
internal/distributed/datacoord/service.go
View File

@ -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"

2
internal/distributed/datanode/service.go
View File

@ -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"

2
internal/distributed/proxy/service.go
View File

@ -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"

2
internal/distributed/querycoord/service.go
View File

@ -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"

2
internal/distributed/querynode/service.go
View File

@ -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"

2
internal/distributed/rootcoord/service.go
View File

@ -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"

2
internal/distributed/rootcoord/service_test.go
View File

@ -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"

8
internal/util/rocksmq/client/rocksmq/client.go → internal/mq/mqimpl/rocksmq/client/client.go
View File

@ -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
}

10
internal/util/rocksmq/client/rocksmq/client_impl.go → internal/mq/mqimpl/rocksmq/client/client_impl.go
View File

@ -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
}

16
internal/util/rocksmq/client/rocksmq/client_impl_test.go → internal/mq/mqimpl/rocksmq/client/client_impl_test.go
View File

@ -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)

6
internal/util/rocksmq/client/rocksmq/consumer.go → internal/mq/mqimpl/rocksmq/client/consumer.go
View File

@ -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

2
internal/util/rocksmq/client/rocksmq/consumer_impl.go → internal/mq/mqimpl/rocksmq/client/consumer_impl.go
View File

@ -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"

6
internal/util/rocksmq/client/rocksmq/consumer_impl_test.go → internal/mq/mqimpl/rocksmq/client/consumer_impl_test.go
View File

@ -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)

2
internal/util/rocksmq/client/rocksmq/error.go → internal/mq/mqimpl/rocksmq/client/error.go
View File

@ -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"

2
internal/util/rocksmq/client/rocksmq/producer.go → internal/mq/mqimpl/rocksmq/client/producer.go
View File

@ -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 {

5
internal/util/rocksmq/client/rocksmq/producer_impl.go → internal/mq/mqimpl/rocksmq/client/producer_impl.go
View File

@ -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"
)

2
internal/util/rocksmq/client/rocksmq/producer_impl_test.go → internal/mq/mqimpl/rocksmq/client/producer_impl_test.go
View File

@ -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"

2
internal/util/rocksmq/client/rocksmq/reader.go → internal/mq/mqimpl/rocksmq/client/reader.go
View File

@ -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"

2
internal/util/rocksmq/client/rocksmq/reader_impl.go → internal/mq/mqimpl/rocksmq/client/reader_impl.go
View File

@ -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"

4
internal/util/rocksmq/client/rocksmq/reader_impl_test.go → internal/mq/mqimpl/rocksmq/client/reader_impl_test.go
View File

@ -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)

10
internal/util/rocksmq/client/rocksmq/test_helper.go → internal/mq/mqimpl/rocksmq/client/test_helper.go
View File

@ -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
}

2
internal/util/rocksmq/server/rocksmq/global_rmq.go → internal/mq/mqimpl/rocksmq/server/global_rmq.go
View File

@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package rocksmq
package server
import (
"errors"

2
internal/util/rocksmq/server/rocksmq/global_rmq_test.go → internal/mq/mqimpl/rocksmq/server/global_rmq_test.go
View File

@ -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"

2
internal/util/rocksmq/server/rocksmq/rocksmq.go → internal/mq/mqimpl/rocksmq/server/rocksmq.go
View File

@ -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"

2
internal/util/rocksmq/server/rocksmq/rocksmq_impl.go → internal/mq/mqimpl/rocksmq/server/rocksmq_impl.go
View File

@ -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"

2
internal/util/rocksmq/server/rocksmq/rocksmq_impl_test.go → internal/mq/mqimpl/rocksmq/server/rocksmq_impl_test.go
View File

@ -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"

2
internal/util/rocksmq/server/rocksmq/rocksmq_reader.go → internal/mq/mqimpl/rocksmq/server/rocksmq_reader.go
View File

@ -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"

2
internal/util/rocksmq/server/rocksmq/rocksmq_retention.go → internal/mq/mqimpl/rocksmq/server/rocksmq_retention.go
View File

@ -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"

2
internal/util/rocksmq/server/rocksmq/rocksmq_retention_test.go → internal/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go
View File

@ -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"

0
internal/msgstream/OWNERS → internal/mq/msgstream/OWNERS
View File

22
internal/msgstream/mq_factory.go → internal/mq/msgstream/mq_factory.go
View File

@ -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
}

0
internal/msgstream/mq_factory_test.go → internal/mq/msgstream/mq_factory_test.go
View File

78
internal/msgstream/mq_msgstream.go → internal/mq/msgstream/mq_msgstream.go
View File

@ -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 {

104
internal/msgstream/mq_msgstream_test.go → internal/mq/msgstream/mq_msgstream_test.go
View File

@ -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))

2
internal/util/mqclient/client.go → internal/mq/msgstream/mqwrapper/client.go
View File

@ -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 {

2
internal/util/mqclient/consumer.go → internal/mq/msgstream/mqwrapper/consumer.go
View File

@ -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

2
internal/util/mqclient/id.go → internal/mq/msgstream/mqwrapper/id.go
View File

@ -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 {

2
internal/util/mqclient/message.go → internal/mq/msgstream/mqwrapper/message.go
View File

@ -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 {

2
internal/util/mqclient/producer.go → internal/mq/msgstream/mqwrapper/producer.go
View File

@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package mqclient
package mqwrapper
import "context"

26
internal/util/mqclient/pulsar_client.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_client.go
View File

@ -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

84
internal/util/mqclient/pulsar_client_test.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_client_test.go
View File

@ -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))
})

22
internal/util/mqclient/pulsar_consumer.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer.go
View File

@ -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

18
internal/util/mqclient/pulsar_consumer_test.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_consumer_test.go
View File

@ -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

14
internal/util/mqclient/pulsar_id.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_id.go
View File

@ -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))
}

8
internal/util/mqclient/pulsar_id_test.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_id_test.go
View File

@ -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)
}

7
internal/util/mqclient/pulsar_message.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_message.go
View File

@ -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

8
internal/util/mqclient/pulsar_producer.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer.go
View File

@ -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

10
internal/util/mqclient/pulsar_producer_test.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_producer_test.go
View File

@ -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{},
}

8
internal/util/mqclient/pulsar_reader.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader.go
View File

@ -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 {

18
internal/util/mqclient/pulsar_reader_test.go → internal/mq/msgstream/mqwrapper/pulsar/pulsar_reader_test.go
View File

@ -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(),
})

2
internal/util/mqclient/reader.go → internal/mq/msgstream/mqwrapper/reader.go
View File

@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package mqclient
package mqwrapper
import (
"context"

46
internal/util/mqclient/rmq_client.go → internal/mq/msgstream/mqwrapper/rmq/rmq_client.go
View File

@ -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

57
internal/util/mqclient/rmq_client_test.go → internal/mq/msgstream/mqwrapper/rmq/rmq_client_test.go
View File

@ -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)
}

26
internal/util/mqclient/rmq_consumer.go → internal/mq/msgstream/mqwrapper/rmq/rmq_consumer.go
View File

@ -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()
}

9
internal/util/mqclient/rmq_id.go → internal/mq/msgstream/mqwrapper/rmq/rmq_id.go
View File

@ -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 {

2
internal/util/mqclient/rmq_id_test.go → internal/mq/msgstream/mqwrapper/rmq/rmq_id_test.go
View File

@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package mqclient
package rmq
import (
"testing"

11
internal/util/mqclient/rmq_message.go → internal/mq/msgstream/mqwrapper/rmq/rmq_message.go
View File

@ -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}
}

14
internal/util/mqclient/rmq_producer.go → internal/mq/msgstream/mqwrapper/rmq/rmq_producer.go
View File

@ -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
}

14
internal/util/mqclient/rmq_reader.go → internal/mq/msgstream/mqwrapper/rmq/rmq_reader.go
View File

@ -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)
}

0
internal/msgstream/msg.go → internal/mq/msgstream/msg.go
View File

0
internal/msgstream/msg_test.go → internal/mq/msgstream/msg_test.go
View File

6
internal/msgstream/msgstream.go → internal/mq/msgstream/msgstream.go
View File

@ -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

0
internal/msgstream/repack_func.go → internal/mq/msgstream/repack_func.go
View File

0
internal/msgstream/trace.go → internal/mq/msgstream/trace.go
View File

0
internal/msgstream/unmarshal.go → internal/mq/msgstream/unmarshal.go
View File

0
internal/msgstream/unmarshal_test.go → internal/mq/msgstream/unmarshal_test.go
View File

2
internal/proxy/channels_mgr.go
View File

@ -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"

2
internal/proxy/impl.go
View File

@ -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"

6
internal/proxy/mock_test.go
View File

@ -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 {

2
internal/proxy/proxy.go
View File

@ -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"

2
internal/proxy/proxy_test.go
View File

@ -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"

2
internal/proxy/repack_func.go
View File

@ -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.

2
internal/proxy/repack_func_test.go
View File

@ -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"
)

2
internal/proxy/task.go
View File

@ -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"

6
internal/proxy/task_scheduler.go
View File

@ -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))

2
internal/proxy/task_test.go
View File

@ -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"

2
internal/querycoord/channel_unsubscribe.go
View File

@ -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"
)

Some files were not shown because too many files have changed in this diff Show More