mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-01 19:39:21 +08:00
add nats mq wrappers (#24445)
bug fixup, configurable natsmq, add unittest, pass e2e. move natsmq to pkg project Signed-off-by: chyezh <ye.zhen@zilliz.com> Co-authored-by: yiwangdr <yiwangdr@gmail.com>
This commit is contained in:
chyezh
GitHub
parent
5f7099a9bd
commit
f97127ae55
5
.gitignore
vendored
5
.gitignore
vendored
@ -83,6 +83,11 @@ cpp_coverage/
|
||||
|
||||
# virtualenv
|
||||
venv/
|
||||
.venv/
|
||||
|
||||
# gopls generated
|
||||
go.work
|
||||
go.work.sum
|
||||
|
||||
# docker compose volumes
|
||||
deployments/docker/*/volumes
|
||||
|
||||
@ -88,7 +88,8 @@ func runComponent[T component](ctx context.Context,
|
||||
localMsg bool,
|
||||
runWg *sync.WaitGroup,
|
||||
creator func(context.Context, dependency.Factory) (T, error),
|
||||
metricRegister func(*prometheus.Registry)) T {
|
||||
metricRegister func(*prometheus.Registry),
|
||||
) T {
|
||||
var role T
|
||||
var wg sync.WaitGroup
|
||||
|
||||
@ -247,19 +248,6 @@ func (mr *MilvusRoles) Run(local bool, alias string) {
|
||||
|
||||
paramtable.Init()
|
||||
params := paramtable.Get()
|
||||
|
||||
if params.RocksmqEnable() {
|
||||
path, err := params.Load("rocksmq.path")
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if err = rocksmqimpl.InitRocksMQ(path); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer stopRocksmq()
|
||||
}
|
||||
|
||||
if params.EtcdCfg.UseEmbedEtcd.GetAsBool() {
|
||||
// Start etcd server.
|
||||
etcd.InitEtcdServer(
|
||||
|
||||
@ -75,23 +75,28 @@ minio:
|
||||
# For more information, refer to
|
||||
# aws: https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles_use.html
|
||||
# gcp: https://cloud.google.com/storage/docs/access-control/iam
|
||||
# aliyun (ack): https://www.alibabacloud.com/help/en/container-service-for-kubernetes/latest/use-rrsa-to-enforce-access-control
|
||||
# aliyun (ack): https://www.alibabacloud.com/help/en/container-service-for-kubernetes/latest/use-rrsa-to-enforce-access-control
|
||||
# aliyun (ecs): https://www.alibabacloud.com/help/en/elastic-compute-service/latest/attach-an-instance-ram-role
|
||||
useIAM: false
|
||||
# Cloud Provider of S3. Supports: "aws", "gcp", "aliyun".
|
||||
# Cloud Provider of S3. Supports: "aws", "gcp", "aliyun".
|
||||
# You can use "aws" for other cloud provider supports S3 API with signature v4, e.g.: minio
|
||||
# You can use "gcp" for other cloud provider supports S3 API with signature v2
|
||||
# You can use "aliyun" for other cloud provider uses virtual host style bucket
|
||||
# You can use "aliyun" for other cloud provider uses virtual host style bucket
|
||||
# When useIAM enabled, only "aws", "gcp", "aliyun" is supported for now
|
||||
cloudProvider: aws
|
||||
# Custom endpoint for fetch IAM role credentials. when useIAM is true & cloudProvider is "aws".
|
||||
# Leave it empty if you want to use AWS default endpoint
|
||||
iamEndpoint:
|
||||
iamEndpoint:
|
||||
|
||||
# Milvus supports three MQ: rocksmq(based on RockDB), Pulsar and Kafka, which should be reserved in config what you use.
|
||||
# There is a note about enabling priority if we config multiple mq in this file
|
||||
# 1. standalone(local) mode: rocksmq(default) > Pulsar > Kafka
|
||||
# 2. cluster mode: Pulsar(default) > Kafka (rocksmq is unsupported)
|
||||
# Milvus supports four MQ: rocksmq(based on RockDB), natsmq(embedded nats-server), Pulsar and Kafka.
|
||||
# You can change your mq by setting mq.type field.
|
||||
# If you don't set mq.type field as default, there is a note about enabling priority if we config multiple mq in this file.
|
||||
# 1. standalone(local) mode: rocksmq(default) > natsmq > Pulsar > Kafka
|
||||
# 2. cluster mode: Pulsar(default) > Kafka (rocksmq and natsmq is unsupported in cluster mode)
|
||||
mq:
|
||||
# Default value: "default"
|
||||
# Valid values: [default, pulsar, kafka, rocksmq, natsmq]
|
||||
type: default
|
||||
|
||||
# Related configuration of pulsar, used to manage Milvus logs of recent mutation operations, output streaming log, and provide log publish-subscribe services.
|
||||
pulsar:
|
||||
@ -104,9 +109,9 @@ pulsar:
|
||||
|
||||
# If you want to enable kafka, needs to comment the pulsar configs
|
||||
# kafka:
|
||||
# brokerList:
|
||||
# saslUsername:
|
||||
# saslPassword:
|
||||
# brokerList:
|
||||
# saslUsername:
|
||||
# saslPassword:
|
||||
# saslMechanisms: PLAIN
|
||||
# securityProtocol: SASL_SSL
|
||||
|
||||
@ -120,6 +125,21 @@ rocksmq:
|
||||
retentionSizeInMB: 8192 # 8 GB, 8 * 1024 MB, The retention size of the message in rocksmq.
|
||||
compactionInterval: 86400 # 1 day, trigger rocksdb compaction every day to remove deleted data
|
||||
|
||||
# natsmq configuration.
|
||||
natsmq:
|
||||
server: # server side configuration for natsmq.
|
||||
port: 4222 # 4222 by default, Port for nats server listening.
|
||||
storeDir: /var/lib/milvus/nats # /var/lib/milvus/nats by default, directory to use for JetStream storage of nats.
|
||||
maxFileStore: 17179869184 # (B) 16GB by default, Maximum size of the 'file' storage.
|
||||
maxPayload: 8388608 # (B) 8MB by default, Maximum number of bytes in a message payload.
|
||||
maxPending: 67108864 # (B) 64MB by default, Maximum number of bytes buffered for a connection Applies to client connections.
|
||||
initializeTimeout: 4000 # (ms) 4s by default, waiting for initialization of natsmq finished.
|
||||
monitor:
|
||||
debug: false # false by default, If true enable debug log messages.
|
||||
logTime: true # true by default, If set to false, log without timestamps.
|
||||
logFile: # no log file by default, Log file path relative to.. .
|
||||
logSizeLimit: 0 # (B) 0, unlimited by default, Size in bytes after the log file rolls over to a new one.
|
||||
|
||||
# Related configuration of rootCoord, used to handle data definition language (DDL) and data control language (DCL) requests
|
||||
rootCoord:
|
||||
dmlChannelNum: 16 # The number of dml channels created at system startup
|
||||
@ -192,7 +212,7 @@ queryCoord:
|
||||
clientMaxRecvSize: 268435456
|
||||
taskMergeCap: 1
|
||||
taskExecutionCap: 256
|
||||
enableActiveStandby: false # Enable active-standby
|
||||
enableActiveStandby: false # Enable active-standby
|
||||
|
||||
# Related configuration of queryNode, used to run hybrid search between vector and scalar data.
|
||||
queryNode:
|
||||
@ -203,7 +223,7 @@ queryNode:
|
||||
stats:
|
||||
publishInterval: 1000 # Interval for querynode to report node information (milliseconds)
|
||||
segcore:
|
||||
knowhereThreadPoolNumRatio: 4
|
||||
knowhereThreadPoolNumRatio: 4
|
||||
# Use more threads to make better use of SSD throughput in disk index.
|
||||
# This parameter is only useful when enable-disk = true.
|
||||
# And this value should be a number greater than 1 and less than 32.
|
||||
@ -288,7 +308,7 @@ dataCoord:
|
||||
compactableProportion: 0.5
|
||||
# over (compactableProportion * segment max # of rows) rows.
|
||||
# MUST BE GREATER THAN OR EQUAL TO <smallProportion>!!!
|
||||
# During compaction, the size of segment # of rows is able to exceed segment max # of rows by (expansionRate-1) * 100%.
|
||||
# During compaction, the size of segment # of rows is able to exceed segment max # of rows by (expansionRate-1) * 100%.
|
||||
expansionRate: 1.25
|
||||
enableCompaction: true # Enable data segment compaction
|
||||
compaction:
|
||||
@ -331,7 +351,7 @@ dataNode:
|
||||
log:
|
||||
level: info # Only supports debug, info, warn, error, panic, or fatal. Default 'info'.
|
||||
file:
|
||||
rootPath: # root dir path to put logs, default "" means no log file will print. please adjust in embedded Milvus: /tmp/milvus/logs
|
||||
rootPath: # root dir path to put logs, default "" means no log file will print. please adjust in embedded Milvus: /tmp/milvus/logs
|
||||
maxSize: 300 # MB
|
||||
maxAge: 10 # Maximum time for log retention in day.
|
||||
maxBackups: 20
|
||||
@ -552,7 +572,7 @@ trace:
|
||||
# Fractions >= 1 will always sample. Fractions < 0 are treated as zero.
|
||||
sampleFraction: 0
|
||||
jaeger:
|
||||
url: # when exporter is jaeger should set the jaeger's URL
|
||||
url: # when exporter is jaeger should set the jaeger's URL
|
||||
|
||||
autoIndex:
|
||||
params:
|
||||
|
||||
24
go.mod
24
go.mod
@ -14,11 +14,10 @@ require (
|
||||
github.com/casbin/casbin/v2 v2.44.2
|
||||
github.com/casbin/json-adapter/v2 v2.0.0
|
||||
github.com/cockroachdb/errors v1.9.1
|
||||
github.com/confluentinc/confluent-kafka-go v1.9.1
|
||||
github.com/gin-gonic/gin v1.9.0
|
||||
github.com/gofrs/flock v0.8.1
|
||||
github.com/golang/protobuf v1.5.3
|
||||
github.com/klauspost/compress v1.14.4
|
||||
github.com/klauspost/compress v1.16.5
|
||||
github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d
|
||||
github.com/milvus-io/milvus-proto/go-api v0.0.0-20230531124827-410c849303a9
|
||||
github.com/milvus-io/milvus/pkg v0.0.0-00010101000000-000000000000
|
||||
@ -39,11 +38,11 @@ require (
|
||||
go.uber.org/atomic v1.10.0
|
||||
go.uber.org/multierr v1.6.0
|
||||
go.uber.org/zap v1.17.0
|
||||
golang.org/x/crypto v0.5.0
|
||||
golang.org/x/crypto v0.8.0
|
||||
golang.org/x/exp v0.0.0-20220303212507-bbda1eaf7a17
|
||||
golang.org/x/oauth2 v0.6.0
|
||||
golang.org/x/sync v0.1.0
|
||||
golang.org/x/text v0.8.0
|
||||
golang.org/x/text v0.9.0
|
||||
google.golang.org/grpc v1.54.0
|
||||
gorm.io/driver/mysql v1.3.5
|
||||
gorm.io/gorm v1.23.8
|
||||
@ -56,13 +55,20 @@ require (
|
||||
github.com/benesch/cgosymbolizer v0.0.0-20190515212042-bec6fe6e597b // indirect
|
||||
github.com/bytedance/sonic v1.8.0 // indirect
|
||||
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
|
||||
github.com/confluentinc/confluent-kafka-go v1.9.1 // indirect
|
||||
github.com/containerd/cgroups v1.1.0 // indirect
|
||||
github.com/facebookgo/ensure v0.0.0-20200202191622-63f1cf65ac4c // indirect
|
||||
github.com/facebookgo/stack v0.0.0-20160209184415-751773369052 // indirect
|
||||
github.com/facebookgo/subset v0.0.0-20200203212716-c811ad88dec4 // indirect
|
||||
github.com/minio/highwayhash v1.0.2 // indirect
|
||||
github.com/nats-io/jwt/v2 v2.4.1 // indirect
|
||||
github.com/nats-io/nats-server/v2 v2.9.17 // indirect
|
||||
github.com/nats-io/nats.go v1.24.0 // indirect
|
||||
github.com/nats-io/nkeys v0.4.4 // indirect
|
||||
github.com/nats-io/nuid v1.0.1 // indirect
|
||||
github.com/pelletier/go-toml/v2 v2.0.6 // indirect
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
||||
go.uber.org/automaxprocs v1.4.0 // indirect
|
||||
go.uber.org/automaxprocs v1.5.1 // indirect
|
||||
golang.org/x/arch v0.0.0-20210923205945-b76863e36670 // indirect
|
||||
)
|
||||
|
||||
@ -190,10 +196,10 @@ require (
|
||||
go.opentelemetry.io/otel/sdk v1.13.0 // indirect
|
||||
go.opentelemetry.io/proto/otlp v0.19.0 // indirect
|
||||
golang.org/x/mod v0.9.0 // indirect
|
||||
golang.org/x/net v0.8.0 // indirect
|
||||
golang.org/x/sys v0.6.0 // indirect
|
||||
golang.org/x/term v0.6.0 // indirect
|
||||
golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba // indirect
|
||||
golang.org/x/net v0.9.0 // indirect
|
||||
golang.org/x/sys v0.7.0 // indirect
|
||||
golang.org/x/term v0.7.0 // indirect
|
||||
golang.org/x/time v0.3.0 // indirect
|
||||
golang.org/x/tools v0.7.0 // indirect
|
||||
golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
|
||||
gonum.org/v1/gonum v0.9.3 // indirect
|
||||
|
||||
43
go.sum
43
go.sum
@ -511,8 +511,9 @@ github.com/klauspost/compress v1.8.2/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0
|
||||
github.com/klauspost/compress v1.9.7/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
|
||||
github.com/klauspost/compress v1.13.5/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
|
||||
github.com/klauspost/compress v1.14.2/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
|
||||
github.com/klauspost/compress v1.14.4 h1:eijASRJcobkVtSt81Olfh7JX43osYLwy5krOJo6YEu4=
|
||||
github.com/klauspost/compress v1.14.4/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
|
||||
github.com/klauspost/compress v1.16.5 h1:IFV2oUNUzZaz+XyusxpLzpzS8Pt5rh0Z16For/djlyI=
|
||||
github.com/klauspost/compress v1.16.5/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
|
||||
github.com/klauspost/cpuid v1.2.1/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
|
||||
github.com/klauspost/cpuid v1.2.3/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
|
||||
github.com/klauspost/cpuid v1.3.1 h1:5JNjFYYQrZeKRJ0734q51WCEEn2huer72Dc7K+R/b6s=
|
||||
@ -589,6 +590,8 @@ github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpsp
|
||||
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
|
||||
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI=
|
||||
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=
|
||||
github.com/minio/highwayhash v1.0.2 h1:Aak5U0nElisjDCfPSG79Tgzkn2gl66NxOMspRrKnA/g=
|
||||
github.com/minio/highwayhash v1.0.2/go.mod h1:BQskDq+xkJ12lmlUUi7U0M5Swg3EWR+dLTk+kldvVxY=
|
||||
github.com/minio/md5-simd v1.1.0 h1:QPfiOqlZH+Cj9teu0t9b1nTBfPbyTl16Of5MeuShdK4=
|
||||
github.com/minio/md5-simd v1.1.0/go.mod h1:XpBqgZULrMYD3R+M28PcmP0CkI7PEMzB3U77ZrKZ0Gw=
|
||||
github.com/minio/minio-go/v7 v7.0.17 h1:5SiS3pqiQDbNhmXMxtqn2HzAInbN5cbHT7ip9F0F07E=
|
||||
@ -619,8 +622,17 @@ github.com/mtibben/percent v0.2.1/go.mod h1:KG9uO+SZkUp+VkRHsCdYQV3XSZrrSpR3O9ib
|
||||
github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
|
||||
github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
|
||||
github.com/nats-io/jwt v0.3.0/go.mod h1:fRYCDE99xlTsqUzISS1Bi75UBJ6ljOJQOAAu5VglpSg=
|
||||
github.com/nats-io/jwt/v2 v2.4.1 h1:Y35W1dgbbz2SQUYDPCaclXcuqleVmpbRa7646Jf2EX4=
|
||||
github.com/nats-io/jwt/v2 v2.4.1/go.mod h1:24BeQtRwxRV8ruvC4CojXlx/WQ/VjuwlYiH+vu/+ibI=
|
||||
github.com/nats-io/nats-server/v2 v2.9.17 h1:gFpUQ3hqIDJrnqog+Bl5vaXg+RhhYEZIElasEuRn2tw=
|
||||
github.com/nats-io/nats-server/v2 v2.9.17/go.mod h1:eQysm3xDZmIjfkjr7DuD9DjRFpnxQc2vKVxtEg0Dp6s=
|
||||
github.com/nats-io/nats.go v1.9.1/go.mod h1:ZjDU1L/7fJ09jvUSRVBR2e7+RnLiiIQyqyzEE/Zbp4w=
|
||||
github.com/nats-io/nats.go v1.24.0 h1:CRiD8L5GOQu/DcfkmgBcTTIQORMwizF+rPk6T0RaHVQ=
|
||||
github.com/nats-io/nats.go v1.24.0/go.mod h1:dVQF+BK3SzUZpwyzHedXsvH3EO38aVKuOPkkHlv5hXA=
|
||||
github.com/nats-io/nkeys v0.1.0/go.mod h1:xpnFELMwJABBLVhffcfd1MZx6VsNRFpEugbxziKVo7w=
|
||||
github.com/nats-io/nkeys v0.4.4 h1:xvBJ8d69TznjcQl9t6//Q5xXuVhyYiSos6RPtvQNTwA=
|
||||
github.com/nats-io/nkeys v0.4.4/go.mod h1:XUkxdLPTufzlihbamfzQ7mw/VGx6ObUs+0bN5sNvt64=
|
||||
github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
|
||||
github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
|
||||
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
|
||||
github.com/nrwiersma/avro-benchmarks v0.0.0-20210913175520-21aec48c8f76/go.mod h1:iKyFMidsk/sVYONJRE372sJuX/QTRPacU7imPqqsu7g=
|
||||
@ -674,6 +686,7 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
|
||||
github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
|
||||
github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
|
||||
github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
|
||||
github.com/prometheus/client_golang v0.9.3/go.mod h1:/TN21ttK/J9q6uSwhBd54HahCDft0ttaMvbicHlPoso=
|
||||
github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
|
||||
@ -905,8 +918,8 @@ go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
|
||||
go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
|
||||
go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
|
||||
go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
|
||||
go.uber.org/automaxprocs v1.4.0 h1:CpDZl6aOlLhReez+8S3eEotD7Jx0Os++lemPlMULQP0=
|
||||
go.uber.org/automaxprocs v1.4.0/go.mod h1:/mTEdr7LvHhs0v7mjdxDreTz1OG5zdZGqgOnhWiR/+Q=
|
||||
go.uber.org/automaxprocs v1.5.1 h1:e1YG66Lrk73dn4qhg8WFSvhF0JuFQF0ERIp4rpuV8Qk=
|
||||
go.uber.org/automaxprocs v1.5.1/go.mod h1:BF4eumQw0P9GtnuxxovUd06vwm1o18oMzFtK66vU6XU=
|
||||
go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk=
|
||||
go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
|
||||
go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU=
|
||||
@ -934,8 +947,8 @@ golang.org/x/crypto v0.0.0-20201216223049-8b5274cf687f/go.mod h1:jdWPYTVW3xRLrWP
|
||||
golang.org/x/crypto v0.0.0-20210322153248-0c34fe9e7dc2/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.0.0-20220411220226-7b82a4e95df4/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
|
||||
golang.org/x/crypto v0.5.0 h1:U/0M97KRkSFvyD/3FSmdP5W5swImpNgle/EHFhOsQPE=
|
||||
golang.org/x/crypto v0.5.0/go.mod h1:NK/OQwhpMQP3MwtdjgLlYHnH9ebylxKWv3e0fK+mkQU=
|
||||
golang.org/x/crypto v0.8.0 h1:pd9TJtTueMTVQXzk8E2XESSMQDj/U7OUu0PqJqPXQjQ=
|
||||
golang.org/x/crypto v0.8.0/go.mod h1:mRqEX+O9/h5TFCrQhkgjo2yKi0yYA+9ecGkdQoHrywE=
|
||||
golang.org/x/exp v0.0.0-20180321215751-8460e604b9de/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/exp v0.0.0-20180807140117-3d87b88a115f/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
@ -1044,8 +1057,8 @@ golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f/go.mod h1:9nx3DQGgdP8bBQD5qx
|
||||
golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
|
||||
golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
|
||||
golang.org/x/net v0.0.0-20220225172249-27dd8689420f/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
|
||||
golang.org/x/net v0.8.0 h1:Zrh2ngAOFYneWTAIAPethzeaQLuHwhuBkuV6ZiRnUaQ=
|
||||
golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc=
|
||||
golang.org/x/net v0.9.0 h1:aWJ/m6xSmxWBx+V0XRHTlrYrPG56jKsLdTFmsSsCzOM=
|
||||
golang.org/x/net v0.9.0/go.mod h1:d48xBJpPfHeWQsugry2m+kC02ZBRGRgulfHnEXEuWns=
|
||||
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
|
||||
golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
|
||||
golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
|
||||
@ -1082,6 +1095,7 @@ golang.org/x/sys v0.0.0-20181026203630-95b1ffbd15a5/go.mod h1:STP8DvDyc/dI5b8T5h
|
||||
golang.org/x/sys v0.0.0-20181107165924-66b7b1311ac8/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190130150945-aca44879d564/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
@ -1153,13 +1167,13 @@ golang.org/x/sys v0.0.0-20220204135822-1c1b9b1eba6a/go.mod h1:oPkhp1MJrh7nUepCBc
|
||||
golang.org/x/sys v0.0.0-20220209214540-3681064d5158/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0 h1:MVltZSvRTcU2ljQOhs94SXPftV6DCNnZViHeQps87pQ=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.7.0 h1:3jlCCIQZPdOYu1h8BkNvLz8Kgwtae2cagcG/VamtZRU=
|
||||
golang.org/x/sys v0.7.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.6.0 h1:clScbb1cHjoCkyRbWwBEUZ5H/tIFu5TAXIqaZD0Gcjw=
|
||||
golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U=
|
||||
golang.org/x/term v0.7.0 h1:BEvjmm5fURWqcfbSKTdpkDXYBrUS1c0m8agp14W48vQ=
|
||||
golang.org/x/term v0.7.0/go.mod h1:P32HKFT3hSsZrRxla30E9HqToFYAQPCMs/zFMBUFqPY=
|
||||
golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
@ -1169,14 +1183,15 @@ golang.org/x/text v0.3.4/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
golang.org/x/text v0.8.0 h1:57P1ETyNKtuIjB4SRd15iJxuhj8Gc416Y78H3qgMh68=
|
||||
golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
|
||||
golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
|
||||
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
|
||||
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20201208040808-7e3f01d25324/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba h1:O8mE0/t419eoIwhTFpKVkHiTs/Igowgfkj25AcZrtiE=
|
||||
golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
|
||||
golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20180525024113-a5b4c53f6e8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
|
||||
@ -1,91 +1,24 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package msgstream
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
rmqimplserver "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server"
|
||||
rmqwrapper "github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/rmq"
|
||||
"github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server"
|
||||
"github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/rmq"
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RmsFactory is a rocksmq msgstream factory that implemented Factory interface(msgstream.go)
|
||||
type RmsFactory struct {
|
||||
dispatcherFactory msgstream.ProtoUDFactory
|
||||
// the following members must be public, so that mapstructure.Decode() can access them
|
||||
ReceiveBufSize int64
|
||||
RmqBufSize int64
|
||||
}
|
||||
|
||||
// NewMsgStream is used to generate a new Msgstream object
|
||||
func (f *RmsFactory) NewMsgStream(ctx context.Context) (msgstream.MsgStream, error) {
|
||||
rmqClient, err := rmqwrapper.NewClientWithDefaultOptions()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
// NewRocksmqFactory creates a new message stream factory based on rocksmq.
|
||||
func NewRocksmqFactory(path string) msgstream.Factory {
|
||||
if err := server.InitRocksMQ(path); err != nil {
|
||||
log.Fatal("fail to init rocksmq", zap.Error(err))
|
||||
}
|
||||
return msgstream.NewMqMsgStream(ctx, f.ReceiveBufSize, f.RmqBufSize, rmqClient, f.dispatcherFactory.NewUnmarshalDispatcher())
|
||||
}
|
||||
log.Info("init rocksmq msgstream success", zap.String("path", path))
|
||||
|
||||
// NewTtMsgStream is used to generate a new TtMsgstream object
|
||||
func (f *RmsFactory) NewTtMsgStream(ctx context.Context) (msgstream.MsgStream, error) {
|
||||
rmqClient, err := rmqwrapper.NewClientWithDefaultOptions()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return msgstream.NewMqTtMsgStream(ctx, f.ReceiveBufSize, f.RmqBufSize, rmqClient, f.dispatcherFactory.NewUnmarshalDispatcher())
|
||||
}
|
||||
|
||||
// NewQueryMsgStream is used to generate a new QueryMsgstream object
|
||||
func (f *RmsFactory) NewQueryMsgStream(ctx context.Context) (msgstream.MsgStream, error) {
|
||||
rmqClient, err := rmqwrapper.NewClientWithDefaultOptions()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return msgstream.NewMqMsgStream(ctx, f.ReceiveBufSize, f.RmqBufSize, rmqClient, f.dispatcherFactory.NewUnmarshalDispatcher())
|
||||
}
|
||||
|
||||
func (f *RmsFactory) NewMsgStreamDisposer(ctx context.Context) func([]string, string) error {
|
||||
return func(channels []string, subname string) error {
|
||||
msgstream, err := f.NewMsgStream(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
msgstream.AsConsumer(channels, subname, mqwrapper.SubscriptionPositionUnknown)
|
||||
msgstream.Close()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// NewRmsFactory is used to generate a new RmsFactory object
|
||||
func NewRmsFactory(path string) *RmsFactory {
|
||||
f := &RmsFactory{
|
||||
dispatcherFactory: msgstream.ProtoUDFactory{},
|
||||
return &msgstream.CommonFactory{
|
||||
Newer: rmq.NewClientWithDefaultOptions,
|
||||
DispatcherFactory: msgstream.ProtoUDFactory{},
|
||||
ReceiveBufSize: 1024,
|
||||
RmqBufSize: 1024,
|
||||
MQBufSize: 1024,
|
||||
}
|
||||
|
||||
err := rmqimplserver.InitRocksMQ(path)
|
||||
if err != nil {
|
||||
log.Error("init rmq error", zap.Error(err))
|
||||
}
|
||||
return f
|
||||
}
|
||||
|
||||
@ -32,7 +32,7 @@ func TestRmsFactory(t *testing.T) {
|
||||
|
||||
dir := t.TempDir()
|
||||
|
||||
rmsFactory := NewRmsFactory(dir)
|
||||
rmsFactory := NewRocksmqFactory(dir)
|
||||
|
||||
ctx := context.Background()
|
||||
_, err := rmsFactory.NewMsgStream(ctx)
|
||||
|
||||
@ -19,6 +19,7 @@ package rmq
|
||||
import (
|
||||
"strconv"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client"
|
||||
"github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server"
|
||||
"go.uber.org/zap"
|
||||
@ -27,12 +28,15 @@ import (
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
)
|
||||
|
||||
// nmqClient implements mqwrapper.Client.
|
||||
var _ mqwrapper.Client = &rmqClient{}
|
||||
|
||||
// rmqClient contains a rocksmq client
|
||||
type rmqClient struct {
|
||||
client client.Client
|
||||
}
|
||||
|
||||
func NewClientWithDefaultOptions() (*rmqClient, error) {
|
||||
func NewClientWithDefaultOptions() (mqwrapper.Client, error) {
|
||||
option := client.Options{Server: server.Rmq}
|
||||
return NewClient(option)
|
||||
}
|
||||
@ -60,6 +64,11 @@ func (rc *rmqClient) CreateProducer(options mqwrapper.ProducerOptions) (mqwrappe
|
||||
|
||||
// Subscribe subscribes a consumer in rmq client
|
||||
func (rc *rmqClient) Subscribe(options mqwrapper.ConsumerOptions) (mqwrapper.Consumer, error) {
|
||||
if options.BufSize == 0 {
|
||||
err := errors.New("subscription bufSize of rmq should never be zero")
|
||||
log.Warn("unexpected subscription consumer options", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
receiveChannel := make(chan client.Message, options.BufSize)
|
||||
|
||||
cli, err := rc.client.Subscribe(client.ConsumerOptions{
|
||||
|
||||
@ -36,10 +36,11 @@ import (
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
paramtable.Init()
|
||||
|
||||
rand.Seed(time.Now().UnixNano())
|
||||
path := "/tmp/milvus/rdb_data"
|
||||
defer os.RemoveAll(path)
|
||||
paramtable.Init()
|
||||
_ = rocksmqimplserver.InitRocksMQ(path)
|
||||
exitCode := m.Run()
|
||||
defer rocksmqimplserver.CloseRocksMQ()
|
||||
@ -153,13 +154,23 @@ func TestRmqClient_Subscribe(t *testing.T) {
|
||||
|
||||
subName := "subName"
|
||||
consumerOpts := mqwrapper.ConsumerOptions{
|
||||
Topic: subName,
|
||||
SubscriptionName: subName,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 0,
|
||||
}
|
||||
consumer, err := client.Subscribe(consumerOpts)
|
||||
assert.NotNil(t, err)
|
||||
assert.Nil(t, consumer)
|
||||
|
||||
consumerOpts = mqwrapper.ConsumerOptions{
|
||||
Topic: "",
|
||||
SubscriptionName: subName,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
}
|
||||
|
||||
consumer, err := client.Subscribe(consumerOpts)
|
||||
consumer, err = client.Subscribe(consumerOpts)
|
||||
assert.NotNil(t, err)
|
||||
assert.Nil(t, consumer)
|
||||
|
||||
|
||||
@ -14,7 +14,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package msgstream
|
||||
package rmq
|
||||
|
||||
import (
|
||||
"context"
|
||||
@ -25,8 +25,6 @@ import (
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/apache/pulsar-client-go/pulsar"
|
||||
"github.com/confluentinc/confluent-kafka-go/kafka"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
@ -35,48 +33,13 @@ import (
|
||||
"github.com/milvus-io/milvus/internal/allocator"
|
||||
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
|
||||
"github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server"
|
||||
"github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper/rmq"
|
||||
"github.com/milvus-io/milvus/pkg/common"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
pulsarwrapper "github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper/pulsar"
|
||||
"github.com/milvus-io/milvus/pkg/util/etcd"
|
||||
"github.com/milvus-io/milvus/pkg/util/funcutil"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultPulsarTenant = "public"
|
||||
DefaultPulsarNamespace = "default"
|
||||
)
|
||||
|
||||
var Params paramtable.ComponentParam
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
Params.Init()
|
||||
mockKafkaCluster, err := kafka.NewMockCluster(1)
|
||||
defer mockKafkaCluster.Close()
|
||||
if err != nil {
|
||||
// nolint
|
||||
fmt.Printf("Failed to create MockCluster: %s\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
broker := mockKafkaCluster.BootstrapServers()
|
||||
Params.Save("kafka.brokerList", broker)
|
||||
|
||||
exitCode := m.Run()
|
||||
os.Exit(exitCode)
|
||||
}
|
||||
|
||||
func getPulsarAddress() string {
|
||||
pulsarHost := Params.GetWithDefault("pulsar.address", "")
|
||||
port := Params.GetWithDefault("pulsar.port", "")
|
||||
if len(pulsarHost) != 0 && len(port) != 0 {
|
||||
return "pulsar://" + pulsarHost + ":" + port
|
||||
}
|
||||
panic("invalid pulsar address")
|
||||
}
|
||||
|
||||
type fixture struct {
|
||||
t *testing.T
|
||||
etcdKV *etcdkv.EtcdKV
|
||||
@ -87,10 +50,6 @@ type parameters struct {
|
||||
}
|
||||
|
||||
func (f *fixture) setup() []parameters {
|
||||
pulsarAddress := getPulsarAddress()
|
||||
pulsarClient, err := pulsarwrapper.NewClient(DefaultPulsarTenant, DefaultPulsarNamespace, pulsar.ClientOptions{URL: pulsarAddress})
|
||||
assert.Nil(f.t, err)
|
||||
|
||||
rocksdbName := "/tmp/rocksmq_unittest_" + f.t.Name()
|
||||
endpoints := os.Getenv("ETCD_ENDPOINTS")
|
||||
if endpoints == "" {
|
||||
@ -110,10 +69,9 @@ func (f *fixture) setup() []parameters {
|
||||
log.Fatalf("InitRmq error = %v", err)
|
||||
}
|
||||
|
||||
rmqClient, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient, _ := NewClientWithDefaultOptions()
|
||||
|
||||
parameters := []parameters{
|
||||
{pulsarClient},
|
||||
{rmqClient},
|
||||
}
|
||||
return parameters
|
||||
@ -437,10 +395,11 @@ func initRmqStream(ctx context.Context,
|
||||
producerChannels []string,
|
||||
consumerChannels []string,
|
||||
consumerGroupName string,
|
||||
opts ...msgstream.RepackFunc) (msgstream.MsgStream, msgstream.MsgStream) {
|
||||
opts ...msgstream.RepackFunc,
|
||||
) (msgstream.MsgStream, msgstream.MsgStream) {
|
||||
factory := msgstream.ProtoUDFactory{}
|
||||
|
||||
rmqClient, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient, _ := NewClientWithDefaultOptions()
|
||||
inputStream, _ := msgstream.NewMqMsgStream(ctx, 100, 100, rmqClient, factory.NewUnmarshalDispatcher())
|
||||
inputStream.AsProducer(producerChannels)
|
||||
for _, opt := range opts {
|
||||
@ -448,7 +407,7 @@ func initRmqStream(ctx context.Context,
|
||||
}
|
||||
var input msgstream.MsgStream = inputStream
|
||||
|
||||
rmqClient2, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient2, _ := NewClientWithDefaultOptions()
|
||||
outputStream, _ := msgstream.NewMqMsgStream(ctx, 100, 100, rmqClient2, factory.NewUnmarshalDispatcher())
|
||||
outputStream.AsConsumer(consumerChannels, consumerGroupName, mqwrapper.SubscriptionPositionEarliest)
|
||||
var output msgstream.MsgStream = outputStream
|
||||
@ -460,10 +419,11 @@ func initRmqTtStream(ctx context.Context,
|
||||
producerChannels []string,
|
||||
consumerChannels []string,
|
||||
consumerGroupName string,
|
||||
opts ...msgstream.RepackFunc) (msgstream.MsgStream, msgstream.MsgStream) {
|
||||
opts ...msgstream.RepackFunc,
|
||||
) (msgstream.MsgStream, msgstream.MsgStream) {
|
||||
factory := msgstream.ProtoUDFactory{}
|
||||
|
||||
rmqClient, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient, _ := NewClientWithDefaultOptions()
|
||||
inputStream, _ := msgstream.NewMqMsgStream(ctx, 100, 100, rmqClient, factory.NewUnmarshalDispatcher())
|
||||
inputStream.AsProducer(producerChannels)
|
||||
for _, opt := range opts {
|
||||
@ -471,7 +431,7 @@ func initRmqTtStream(ctx context.Context,
|
||||
}
|
||||
var input msgstream.MsgStream = inputStream
|
||||
|
||||
rmqClient2, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient2, _ := NewClientWithDefaultOptions()
|
||||
outputStream, _ := msgstream.NewMqTtMsgStream(ctx, 100, 100, rmqClient2, factory.NewUnmarshalDispatcher())
|
||||
outputStream.AsConsumer(consumerChannels, consumerGroupName, mqwrapper.SubscriptionPositionEarliest)
|
||||
var output msgstream.MsgStream = outputStream
|
||||
@ -578,7 +538,7 @@ func TestStream_RmqTtMsgStream_DuplicatedIDs(t *testing.T) {
|
||||
|
||||
factory := msgstream.ProtoUDFactory{}
|
||||
|
||||
rmqClient, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient, _ := NewClientWithDefaultOptions()
|
||||
outputStream, _ = msgstream.NewMqTtMsgStream(context.Background(), 100, 100, rmqClient, factory.NewUnmarshalDispatcher())
|
||||
consumerSubName = funcutil.RandomString(8)
|
||||
outputStream.AsConsumer(consumerChannels, consumerSubName, mqwrapper.SubscriptionPositionUnknown)
|
||||
@ -682,7 +642,7 @@ func TestStream_RmqTtMsgStream_Seek(t *testing.T) {
|
||||
|
||||
factory := msgstream.ProtoUDFactory{}
|
||||
|
||||
rmqClient, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient, _ := NewClientWithDefaultOptions()
|
||||
outputStream, _ = msgstream.NewMqTtMsgStream(context.Background(), 100, 100, rmqClient, factory.NewUnmarshalDispatcher())
|
||||
consumerSubName = funcutil.RandomString(8)
|
||||
outputStream.AsConsumer(consumerChannels, consumerSubName, mqwrapper.SubscriptionPositionUnknown)
|
||||
@ -731,7 +691,7 @@ func TestStream_RMqMsgStream_SeekInvalidMessage(t *testing.T) {
|
||||
outputStream.Close()
|
||||
|
||||
factory := msgstream.ProtoUDFactory{}
|
||||
rmqClient2, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient2, _ := NewClientWithDefaultOptions()
|
||||
outputStream2, _ := msgstream.NewMqMsgStream(ctx, 100, 100, rmqClient2, factory.NewUnmarshalDispatcher())
|
||||
outputStream2.AsConsumer(consumerChannels, funcutil.RandomString(8), mqwrapper.SubscriptionPositionUnknown)
|
||||
|
||||
@ -764,7 +724,6 @@ func TestStream_RMqMsgStream_SeekInvalidMessage(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestStream_RmqTtMsgStream_AsConsumerWithPosition(t *testing.T) {
|
||||
|
||||
producerChannels := []string{"insert1"}
|
||||
consumerChannels := []string{"insert1"}
|
||||
consumerSubName := "subInsert"
|
||||
@ -773,7 +732,7 @@ func TestStream_RmqTtMsgStream_AsConsumerWithPosition(t *testing.T) {
|
||||
etcdKV := initRmq(rocksdbName)
|
||||
factory := msgstream.ProtoUDFactory{}
|
||||
|
||||
rmqClient, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient, _ := NewClientWithDefaultOptions()
|
||||
|
||||
otherInputStream, _ := msgstream.NewMqMsgStream(context.Background(), 100, 100, rmqClient, factory.NewUnmarshalDispatcher())
|
||||
otherInputStream.AsProducer([]string{"root_timetick"})
|
||||
@ -786,7 +745,7 @@ func TestStream_RmqTtMsgStream_AsConsumerWithPosition(t *testing.T) {
|
||||
inputStream.Produce(getTimeTickMsgPack(int64(i)))
|
||||
}
|
||||
|
||||
rmqClient2, _ := rmq.NewClientWithDefaultOptions()
|
||||
rmqClient2, _ := NewClientWithDefaultOptions()
|
||||
outputStream, _ := msgstream.NewMqMsgStream(context.Background(), 100, 100, rmqClient2, factory.NewUnmarshalDispatcher())
|
||||
outputStream.AsConsumer(consumerChannels, consumerSubName, mqwrapper.SubscriptionPositionLatest)
|
||||
|
||||
@ -3,12 +3,31 @@ package dependency
|
||||
import (
|
||||
"context"
|
||||
|
||||
rmqstream "github.com/milvus-io/milvus/internal/mq/msgstream"
|
||||
"github.com/cockroachdb/errors"
|
||||
smsgstream "github.com/milvus-io/milvus/internal/mq/msgstream"
|
||||
"github.com/milvus-io/milvus/internal/storage"
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
mqTypeDefault = "default"
|
||||
mqTypeNatsmq = "natsmq"
|
||||
mqTypeRocksmq = "rocksmq"
|
||||
mqTypeKafka = "kafka"
|
||||
mqTypePulsar = "pulsar"
|
||||
)
|
||||
|
||||
type mqEnable struct {
|
||||
Rocksmq bool
|
||||
Natsmq bool
|
||||
Pulsar bool
|
||||
Kafka bool
|
||||
}
|
||||
|
||||
// DefaultFactory is a factory that produces instances of storage.ChunkManager and message queue.
|
||||
type DefaultFactory struct {
|
||||
standAlone bool
|
||||
chunkManagerFactory storage.Factory
|
||||
@ -19,20 +38,23 @@ type DefaultFactory struct {
|
||||
func NewDefaultFactory(standAlone bool) *DefaultFactory {
|
||||
return &DefaultFactory{
|
||||
standAlone: standAlone,
|
||||
msgStreamFactory: rmqstream.NewRmsFactory("/tmp/milvus/rocksmq/"),
|
||||
msgStreamFactory: smsgstream.NewRocksmqFactory("/tmp/milvus/rocksmq/"),
|
||||
chunkManagerFactory: storage.NewChunkManagerFactory("local",
|
||||
storage.RootPath("/tmp/milvus")),
|
||||
}
|
||||
}
|
||||
|
||||
// NewFactory creates a new instance of the DefaultFactory type.
|
||||
// If standAlone is true, the factory will operate in standalone mode.
|
||||
func NewFactory(standAlone bool) *DefaultFactory {
|
||||
return &DefaultFactory{standAlone: standAlone}
|
||||
}
|
||||
|
||||
// Init create a msg factory(TODO only support one mq at the same time.)
|
||||
// In order to guarantee backward compatibility of config file, we still support multiple mq configs.
|
||||
// 1. Rocksmq only run on local mode, and it has the highest priority
|
||||
// 2. Pulsar has higher priority than Kafka within remote msg
|
||||
// The initialization of MQ follows the following rules, if the mq.type is default.
|
||||
// 1. standalone(local) mode: rocksmq(default) > natsmq > Pulsar > Kafka
|
||||
// 2. cluster mode: Pulsar(default) > Kafka (rocksmq and natsmq is unsupported in cluster mode)
|
||||
func (f *DefaultFactory) Init(params *paramtable.ComponentParam) {
|
||||
// skip if using default factory
|
||||
if f.msgStreamFactory != nil {
|
||||
@ -41,41 +63,67 @@ func (f *DefaultFactory) Init(params *paramtable.ComponentParam) {
|
||||
|
||||
f.chunkManagerFactory = storage.NewChunkManagerFactoryWithParam(params)
|
||||
|
||||
// init mq storage
|
||||
if f.standAlone {
|
||||
f.msgStreamFactory = f.initMQLocalService(params)
|
||||
if f.msgStreamFactory != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
f.msgStreamFactory = f.initMQRemoteService(params)
|
||||
if f.msgStreamFactory == nil {
|
||||
panic("no available remote mq configuration, must config Pulsar or Kafka at least one of these!")
|
||||
// initialize mq client or embedded mq.
|
||||
if err := f.initMQ(f.standAlone, params); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *DefaultFactory) initMQLocalService(params *paramtable.ComponentParam) msgstream.Factory {
|
||||
if params.RocksmqEnable() {
|
||||
path, err := params.Load("rocksmq.path")
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return rmqstream.NewRmsFactory(path)
|
||||
func (f *DefaultFactory) initMQ(standalone bool, params *paramtable.ComponentParam) error {
|
||||
mqType := mustSelectMQType(standalone, params.MQCfg.Type.GetValue(), mqEnable{params.RocksmqEnable(), params.NatsmqEnable(), params.PulsarEnable(), params.KafkaEnable()})
|
||||
log.Info("try to init mq", zap.Bool("standalone", standalone), zap.String("mqType", mqType))
|
||||
|
||||
switch mqType {
|
||||
case mqTypeNatsmq:
|
||||
f.msgStreamFactory = msgstream.NewNatsmqFactory()
|
||||
case mqTypeRocksmq:
|
||||
f.msgStreamFactory = smsgstream.NewRocksmqFactory(params.RocksmqCfg.Path.GetValue())
|
||||
case mqTypePulsar:
|
||||
f.msgStreamFactory = msgstream.NewPmsFactory(¶ms.PulsarCfg)
|
||||
case mqTypeKafka:
|
||||
f.msgStreamFactory = msgstream.NewKmsFactory(¶ms.KafkaCfg)
|
||||
}
|
||||
if f.msgStreamFactory == nil {
|
||||
return errors.New("failed to create MQ: check the milvus log for initialization failures")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// initRemoteService Pulsar has higher priority than Kafka.
|
||||
func (f *DefaultFactory) initMQRemoteService(params *paramtable.ComponentParam) msgstream.Factory {
|
||||
if params.PulsarEnable() {
|
||||
return msgstream.NewPmsFactory(¶ms.PulsarCfg)
|
||||
// Select valid mq if mq type is default.
|
||||
func mustSelectMQType(standalone bool, mqType string, enable mqEnable) string {
|
||||
if mqType != mqTypeDefault {
|
||||
if err := validateMQType(standalone, mqType); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return mqType
|
||||
}
|
||||
|
||||
if params.KafkaEnable() {
|
||||
return msgstream.NewKmsFactory(¶ms.KafkaCfg)
|
||||
if standalone {
|
||||
if enable.Rocksmq {
|
||||
return mqTypeRocksmq
|
||||
}
|
||||
if enable.Natsmq {
|
||||
return mqTypeNatsmq
|
||||
}
|
||||
}
|
||||
if enable.Pulsar {
|
||||
return mqTypePulsar
|
||||
}
|
||||
if enable.Kafka {
|
||||
return mqTypeKafka
|
||||
}
|
||||
|
||||
panic(errors.Errorf("no available mq config found, %s, enable: %+v", mqType, enable))
|
||||
}
|
||||
|
||||
// Validate mq type.
|
||||
func validateMQType(standalone bool, mqType string) error {
|
||||
if mqType != mqTypeNatsmq && mqType != mqTypeRocksmq && mqType != mqTypeKafka && mqType != mqTypePulsar {
|
||||
return errors.Newf("mq type %s is invalid", mqType)
|
||||
}
|
||||
if !standalone && (mqType == mqTypeRocksmq || mqType == mqTypeNatsmq) {
|
||||
return errors.Newf("mq %s is only valid in standalone mode")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
35
internal/util/dependency/factory_test.go
Normal file
35
internal/util/dependency/factory_test.go
Normal file
@ -0,0 +1,35 @@
|
||||
package dependency
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestValidateMQType(t *testing.T) {
|
||||
assert.Error(t, validateMQType(true, mqTypeDefault))
|
||||
assert.Error(t, validateMQType(false, mqTypeDefault))
|
||||
assert.Error(t, validateMQType(false, mqTypeNatsmq))
|
||||
assert.Error(t, validateMQType(false, mqTypeRocksmq))
|
||||
}
|
||||
|
||||
func TestSelectMQType(t *testing.T) {
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypeDefault, mqEnable{true, true, true, true}), mqTypeRocksmq)
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypeDefault, mqEnable{false, true, true, true}), mqTypeNatsmq)
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypeDefault, mqEnable{false, false, true, true}), mqTypePulsar)
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypeDefault, mqEnable{false, false, false, true}), mqTypeKafka)
|
||||
assert.Panics(t, func() { mustSelectMQType(true, mqTypeDefault, mqEnable{false, false, false, false}) })
|
||||
assert.Equal(t, mustSelectMQType(false, mqTypeDefault, mqEnable{true, true, true, true}), mqTypePulsar)
|
||||
assert.Equal(t, mustSelectMQType(false, mqTypeDefault, mqEnable{false, true, true, true}), mqTypePulsar)
|
||||
assert.Equal(t, mustSelectMQType(false, mqTypeDefault, mqEnable{false, false, true, true}), mqTypePulsar)
|
||||
assert.Equal(t, mustSelectMQType(false, mqTypeDefault, mqEnable{false, false, false, true}), mqTypeKafka)
|
||||
assert.Panics(t, func() { mustSelectMQType(false, mqTypeDefault, mqEnable{false, false, false, false}) })
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypeRocksmq, mqEnable{true, true, true, true}), mqTypeRocksmq)
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypeNatsmq, mqEnable{true, true, true, true}), mqTypeNatsmq)
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypePulsar, mqEnable{true, true, true, true}), mqTypePulsar)
|
||||
assert.Equal(t, mustSelectMQType(true, mqTypeKafka, mqEnable{true, true, true, true}), mqTypeKafka)
|
||||
assert.Panics(t, func() { mustSelectMQType(false, mqTypeRocksmq, mqEnable{true, true, true, true}) })
|
||||
assert.Panics(t, func() { mustSelectMQType(false, mqTypeNatsmq, mqEnable{true, true, true, true}) })
|
||||
assert.Equal(t, mustSelectMQType(false, mqTypePulsar, mqEnable{true, true, true, true}), mqTypePulsar)
|
||||
assert.Equal(t, mustSelectMQType(false, mqTypeKafka, mqEnable{true, true, true, true}), mqTypeKafka)
|
||||
}
|
||||
22
pkg/go.mod
22
pkg/go.mod
@ -10,7 +10,7 @@ require (
|
||||
github.com/containerd/cgroups v1.0.4
|
||||
github.com/golang/protobuf v1.5.2
|
||||
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
|
||||
github.com/klauspost/compress v1.14.4
|
||||
github.com/klauspost/compress v1.16.5
|
||||
github.com/lingdor/stackerror v0.0.0-20191119040541-976d8885ed76
|
||||
github.com/milvus-io/milvus-proto/go-api v0.0.0-20230529034923-4579ee9d5723
|
||||
github.com/panjf2000/ants/v2 v2.7.2
|
||||
@ -25,9 +25,9 @@ require (
|
||||
go.etcd.io/etcd/client/v3 v3.5.5
|
||||
go.etcd.io/etcd/server/v3 v3.5.5
|
||||
go.uber.org/atomic v1.10.0
|
||||
go.uber.org/automaxprocs v1.4.0
|
||||
go.uber.org/automaxprocs v1.5.1
|
||||
go.uber.org/zap v1.17.0
|
||||
golang.org/x/crypto v0.0.0-20220411220226-7b82a4e95df4
|
||||
golang.org/x/crypto v0.8.0
|
||||
golang.org/x/exp v0.0.0-20220303212507-bbda1eaf7a17
|
||||
golang.org/x/sync v0.1.0
|
||||
google.golang.org/grpc v1.52.3
|
||||
@ -104,10 +104,10 @@ require (
|
||||
go.opentelemetry.io/otel/sdk v1.13.0
|
||||
go.opentelemetry.io/otel/trace v1.13.0
|
||||
go.opentelemetry.io/proto/otlp v0.19.0 // indirect
|
||||
golang.org/x/net v0.7.0 // indirect
|
||||
golang.org/x/sys v0.5.0 // indirect
|
||||
golang.org/x/term v0.5.0 // indirect
|
||||
golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba // indirect
|
||||
golang.org/x/net v0.9.0 // indirect
|
||||
golang.org/x/sys v0.7.0 // indirect
|
||||
golang.org/x/term v0.7.0 // indirect
|
||||
golang.org/x/time v0.3.0 // indirect
|
||||
google.golang.org/appengine v1.6.7 // indirect
|
||||
google.golang.org/genproto v0.0.0-20221118155620-16455021b5e6 // indirect
|
||||
gopkg.in/ini.v1 v1.62.0 // indirect
|
||||
@ -120,6 +120,8 @@ require github.com/ianlancetaylor/cgosymbolizer v0.0.0-20221217025313-27d3c9f66b
|
||||
|
||||
require (
|
||||
github.com/cockroachdb/errors v1.9.1
|
||||
github.com/nats-io/nats-server/v2 v2.9.17
|
||||
github.com/nats-io/nats.go v1.24.0
|
||||
github.com/uber/jaeger-client-go v2.30.0+incompatible
|
||||
go.opentelemetry.io/otel/exporters/jaeger v1.13.0
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.13.0
|
||||
@ -138,6 +140,10 @@ require (
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.7.0 // indirect
|
||||
github.com/kr/pretty v0.3.0 // indirect
|
||||
github.com/kr/text v0.2.0 // indirect
|
||||
github.com/minio/highwayhash v1.0.2 // indirect
|
||||
github.com/nats-io/jwt/v2 v2.4.1 // indirect
|
||||
github.com/nats-io/nkeys v0.4.4 // indirect
|
||||
github.com/nats-io/nuid v1.0.1 // indirect
|
||||
github.com/olekukonko/tablewriter v0.0.5 // indirect
|
||||
github.com/pkg/errors v0.9.1 // indirect
|
||||
github.com/rogpeppe/go-internal v1.8.1 // indirect
|
||||
@ -146,7 +152,7 @@ require (
|
||||
go.opentelemetry.io/otel/metric v0.35.0 // indirect
|
||||
go.uber.org/multierr v1.6.0 // indirect
|
||||
golang.org/x/oauth2 v0.0.0-20221014153046-6fdb5e3db783 // indirect
|
||||
golang.org/x/text v0.7.0 // indirect
|
||||
golang.org/x/text v0.9.0 // indirect
|
||||
)
|
||||
|
||||
replace (
|
||||
|
||||
42
pkg/go.sum
42
pkg/go.sum
@ -412,8 +412,9 @@ github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI
|
||||
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
|
||||
github.com/klauspost/compress v1.8.2/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
|
||||
github.com/klauspost/compress v1.9.7/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
|
||||
github.com/klauspost/compress v1.14.4 h1:eijASRJcobkVtSt81Olfh7JX43osYLwy5krOJo6YEu4=
|
||||
github.com/klauspost/compress v1.14.4/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
|
||||
github.com/klauspost/compress v1.16.5 h1:IFV2oUNUzZaz+XyusxpLzpzS8Pt5rh0Z16For/djlyI=
|
||||
github.com/klauspost/compress v1.16.5/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
|
||||
github.com/klauspost/cpuid v1.2.1/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
|
||||
github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
|
||||
github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
|
||||
@ -469,6 +470,8 @@ github.com/milvus-io/milvus-proto/go-api v0.0.0-20230529034923-4579ee9d5723 h1:V
|
||||
github.com/milvus-io/milvus-proto/go-api v0.0.0-20230529034923-4579ee9d5723/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
|
||||
github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A=
|
||||
github.com/milvus-io/pulsar-client-go v0.6.10/go.mod h1:lQqCkgwDF8YFYjKA+zOheTk1tev2B+bKj5j7+nm8M1w=
|
||||
github.com/minio/highwayhash v1.0.2 h1:Aak5U0nElisjDCfPSG79Tgzkn2gl66NxOMspRrKnA/g=
|
||||
github.com/minio/highwayhash v1.0.2/go.mod h1:BQskDq+xkJ12lmlUUi7U0M5Swg3EWR+dLTk+kldvVxY=
|
||||
github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=
|
||||
github.com/mitchellh/go-homedir v1.0.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
|
||||
github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
|
||||
@ -491,8 +494,17 @@ github.com/mtibben/percent v0.2.1/go.mod h1:KG9uO+SZkUp+VkRHsCdYQV3XSZrrSpR3O9ib
|
||||
github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
|
||||
github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
|
||||
github.com/nats-io/jwt v0.3.0/go.mod h1:fRYCDE99xlTsqUzISS1Bi75UBJ6ljOJQOAAu5VglpSg=
|
||||
github.com/nats-io/jwt/v2 v2.4.1 h1:Y35W1dgbbz2SQUYDPCaclXcuqleVmpbRa7646Jf2EX4=
|
||||
github.com/nats-io/jwt/v2 v2.4.1/go.mod h1:24BeQtRwxRV8ruvC4CojXlx/WQ/VjuwlYiH+vu/+ibI=
|
||||
github.com/nats-io/nats-server/v2 v2.9.17 h1:gFpUQ3hqIDJrnqog+Bl5vaXg+RhhYEZIElasEuRn2tw=
|
||||
github.com/nats-io/nats-server/v2 v2.9.17/go.mod h1:eQysm3xDZmIjfkjr7DuD9DjRFpnxQc2vKVxtEg0Dp6s=
|
||||
github.com/nats-io/nats.go v1.9.1/go.mod h1:ZjDU1L/7fJ09jvUSRVBR2e7+RnLiiIQyqyzEE/Zbp4w=
|
||||
github.com/nats-io/nats.go v1.24.0 h1:CRiD8L5GOQu/DcfkmgBcTTIQORMwizF+rPk6T0RaHVQ=
|
||||
github.com/nats-io/nats.go v1.24.0/go.mod h1:dVQF+BK3SzUZpwyzHedXsvH3EO38aVKuOPkkHlv5hXA=
|
||||
github.com/nats-io/nkeys v0.1.0/go.mod h1:xpnFELMwJABBLVhffcfd1MZx6VsNRFpEugbxziKVo7w=
|
||||
github.com/nats-io/nkeys v0.4.4 h1:xvBJ8d69TznjcQl9t6//Q5xXuVhyYiSos6RPtvQNTwA=
|
||||
github.com/nats-io/nkeys v0.4.4/go.mod h1:XUkxdLPTufzlihbamfzQ7mw/VGx6ObUs+0bN5sNvt64=
|
||||
github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
|
||||
github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
|
||||
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
|
||||
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
|
||||
@ -541,6 +553,7 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
|
||||
github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
|
||||
github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
|
||||
github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
|
||||
github.com/prometheus/client_golang v0.9.3/go.mod h1:/TN21ttK/J9q6uSwhBd54HahCDft0ttaMvbicHlPoso=
|
||||
github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
|
||||
@ -742,8 +755,8 @@ go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
|
||||
go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
|
||||
go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
|
||||
go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
|
||||
go.uber.org/automaxprocs v1.4.0 h1:CpDZl6aOlLhReez+8S3eEotD7Jx0Os++lemPlMULQP0=
|
||||
go.uber.org/automaxprocs v1.4.0/go.mod h1:/mTEdr7LvHhs0v7mjdxDreTz1OG5zdZGqgOnhWiR/+Q=
|
||||
go.uber.org/automaxprocs v1.5.1 h1:e1YG66Lrk73dn4qhg8WFSvhF0JuFQF0ERIp4rpuV8Qk=
|
||||
go.uber.org/automaxprocs v1.5.1/go.mod h1:BF4eumQw0P9GtnuxxovUd06vwm1o18oMzFtK66vU6XU=
|
||||
go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk=
|
||||
go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
|
||||
go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
|
||||
@ -764,8 +777,9 @@ golang.org/x/crypto v0.0.0-20191227163750-53104e6ec876/go.mod h1:LzIPMQfyMNhhGPh
|
||||
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
|
||||
golang.org/x/crypto v0.0.0-20210322153248-0c34fe9e7dc2/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.0.0-20220411220226-7b82a4e95df4 h1:kUhD7nTDoI3fVd9G4ORWrbV5NY0liEs/Jg2pv5f+bBA=
|
||||
golang.org/x/crypto v0.0.0-20220411220226-7b82a4e95df4/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
|
||||
golang.org/x/crypto v0.8.0 h1:pd9TJtTueMTVQXzk8E2XESSMQDj/U7OUu0PqJqPXQjQ=
|
||||
golang.org/x/crypto v0.8.0/go.mod h1:mRqEX+O9/h5TFCrQhkgjo2yKi0yYA+9ecGkdQoHrywE=
|
||||
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/exp v0.0.0-20190510132918-efd6b22b2522/go.mod h1:ZjyILWgesfNpC6sMxTJOJm9Kp84zZh5NQWvqDGG3Qr8=
|
||||
@ -853,8 +867,8 @@ golang.org/x/net v0.0.0-20210726213435-c6fcb2dbf985/go.mod h1:9nx3DQGgdP8bBQD5qx
|
||||
golang.org/x/net v0.0.0-20211008194852-3b03d305991f/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
|
||||
golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
|
||||
golang.org/x/net v0.0.0-20220225172249-27dd8689420f/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
|
||||
golang.org/x/net v0.7.0 h1:rJrUqqhjsgNp7KqAIc25s9pZnjU7TUcSY7HcVZjdn1g=
|
||||
golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
|
||||
golang.org/x/net v0.9.0 h1:aWJ/m6xSmxWBx+V0XRHTlrYrPG56jKsLdTFmsSsCzOM=
|
||||
golang.org/x/net v0.9.0/go.mod h1:d48xBJpPfHeWQsugry2m+kC02ZBRGRgulfHnEXEuWns=
|
||||
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
|
||||
golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
|
||||
golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
|
||||
@ -891,6 +905,7 @@ golang.org/x/sys v0.0.0-20181026203630-95b1ffbd15a5/go.mod h1:STP8DvDyc/dI5b8T5h
|
||||
golang.org/x/sys v0.0.0-20181107165924-66b7b1311ac8/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190130150945-aca44879d564/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
@ -958,12 +973,12 @@ golang.org/x/sys v0.0.0-20220128215802-99c3d69c2c27/go.mod h1:oPkhp1MJrh7nUepCBc
|
||||
golang.org/x/sys v0.0.0-20220204135822-1c1b9b1eba6a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220209214540-3681064d5158/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.5.0 h1:MUK/U/4lj1t1oPg0HfuXDN/Z1wv31ZJ/YcPiGccS4DU=
|
||||
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.7.0 h1:3jlCCIQZPdOYu1h8BkNvLz8Kgwtae2cagcG/VamtZRU=
|
||||
golang.org/x/sys v0.7.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.5.0 h1:n2a8QNdAb0sZNpU9R1ALUXBbY+w51fCQDN+7EdxNBsY=
|
||||
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
||||
golang.org/x/term v0.7.0 h1:BEvjmm5fURWqcfbSKTdpkDXYBrUS1c0m8agp14W48vQ=
|
||||
golang.org/x/term v0.7.0/go.mod h1:P32HKFT3hSsZrRxla30E9HqToFYAQPCMs/zFMBUFqPY=
|
||||
golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
@ -973,14 +988,15 @@ golang.org/x/text v0.3.4/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
golang.org/x/text v0.7.0 h1:4BRB4x83lYWy72KwLD/qYDuTu7q9PjSagHvijDw7cLo=
|
||||
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
|
||||
golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
|
||||
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
|
||||
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20201208040808-7e3f01d25324/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba h1:O8mE0/t419eoIwhTFpKVkHiTs/Igowgfkj25AcZrtiE=
|
||||
golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
|
||||
golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20181030221726-6c7e314b6563/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
|
||||
73
pkg/mq/msgstream/common_mq_factory.go
Normal file
73
pkg/mq/msgstream/common_mq_factory.go
Normal file
@ -0,0 +1,73 @@
|
||||
package msgstream
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
)
|
||||
|
||||
var _ Factory = &CommonFactory{}
|
||||
|
||||
// CommonFactory is a Factory for creating message streams with common logic.
|
||||
//
|
||||
// It contains a function field named newer, which is a function that creates
|
||||
// an mqwrapper.Client when called.
|
||||
type CommonFactory struct {
|
||||
Newer func() (mqwrapper.Client, error) // client constructor
|
||||
DispatcherFactory ProtoUDFactory
|
||||
ReceiveBufSize int64
|
||||
MQBufSize int64
|
||||
}
|
||||
|
||||
// NewMsgStream is used to generate a new Msgstream object
|
||||
func (f *CommonFactory) NewMsgStream(ctx context.Context) (ms MsgStream, err error) {
|
||||
defer wrapError(&err, "NewMsgStream")
|
||||
cli, err := f.Newer()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return NewMqMsgStream(ctx, f.ReceiveBufSize, f.MQBufSize, cli, f.DispatcherFactory.NewUnmarshalDispatcher())
|
||||
}
|
||||
|
||||
// NewTtMsgStream is used to generate a new TtMsgstream object
|
||||
func (f *CommonFactory) NewTtMsgStream(ctx context.Context) (ms MsgStream, err error) {
|
||||
defer wrapError(&err, "NewTtMsgStream")
|
||||
cli, err := f.Newer()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return NewMqTtMsgStream(ctx, f.ReceiveBufSize, f.MQBufSize, cli, f.DispatcherFactory.NewUnmarshalDispatcher())
|
||||
}
|
||||
|
||||
// NewQueryMsgStream is used to generate a new QueryMsgstream object
|
||||
func (f *CommonFactory) NewQueryMsgStream(ctx context.Context) (ms MsgStream, err error) {
|
||||
defer wrapError(&err, "NewQueryMsgStream")
|
||||
cli, err := f.Newer()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return NewMqMsgStream(ctx, f.ReceiveBufSize, f.MQBufSize, cli, f.DispatcherFactory.NewUnmarshalDispatcher())
|
||||
}
|
||||
|
||||
// NewMsgStreamDisposer returns a function that can be used to dispose of a message stream.
|
||||
// The returned function takes a slice of channel names and a subscription name, and
|
||||
// disposes of the message stream associated with those arguments.
|
||||
func (f *CommonFactory) NewMsgStreamDisposer(ctx context.Context) func([]string, string) error {
|
||||
return func(channels []string, subName string) (err error) {
|
||||
defer wrapError(&err, "NewMsgStreamDisposer")
|
||||
msgs, err := f.NewMsgStream(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
msgs.AsConsumer(channels, subName, mqwrapper.SubscriptionPositionUnknown)
|
||||
msgs.Close()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func wrapError(err *error, method string) {
|
||||
if *err != nil {
|
||||
*err = errors.Wrapf(*err, "in method: %s", method)
|
||||
}
|
||||
}
|
||||
844
pkg/mq/msgstream/factory_stream_test.go
Normal file
844
pkg/mq/msgstream/factory_stream_test.go
Normal file
@ -0,0 +1,844 @@
|
||||
package msgstream
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"testing"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/msgpb"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"github.com/milvus-io/milvus/pkg/util/funcutil"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type streamNewer func(ctx context.Context) (MsgStream, error)
|
||||
|
||||
// test all stream operation on stream factory
|
||||
func testMsgStreamOperation(t *testing.T, factories []Factory, pg positionGenerator) {
|
||||
testFuncs := []func(t *testing.T, f []Factory){
|
||||
testInsert,
|
||||
testDelete,
|
||||
testTimeTick,
|
||||
testBroadCast,
|
||||
testInsertWithRepack,
|
||||
testInsertRepackFuncWithDifferentClient,
|
||||
testDeleteRepackFuncWithDifferentClient,
|
||||
testDefaultRepack,
|
||||
testTimeTickerAndInsert,
|
||||
testTimeTickerNoSeek,
|
||||
testSeekToLast,
|
||||
testTimeTickerSeek,
|
||||
testTimeTickUnmarshalHeader,
|
||||
testTimeTickerStream1,
|
||||
testTimeTickerStream2,
|
||||
testMqMsgStreamSeek,
|
||||
func(t *testing.T, f []Factory) {
|
||||
testMqMsgStreamSeekInvalidMessage(t, f, pg)
|
||||
},
|
||||
testMqMsgStreamSeekLatest,
|
||||
testBroadcastMark,
|
||||
}
|
||||
|
||||
for _, testFunc := range testFuncs {
|
||||
t.Run(
|
||||
runtime.FuncForPC(reflect.ValueOf(testFunc).Pointer()).Name(),
|
||||
func(t *testing.T) {
|
||||
testFunc(t, factories)
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func testInsert(t *testing.T, f []Factory) {
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_Insert, 1))
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_Insert, 3))
|
||||
applyProduceAndConsume(t, &msgPack, []streamNewer{f[0].NewMsgStream, f[0].NewMsgStream}, 2)
|
||||
}
|
||||
|
||||
func testDelete(t *testing.T, f []Factory) {
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_Delete, 1))
|
||||
applyProduceAndConsume(t, &msgPack, []streamNewer{f[0].NewMsgStream, f[0].NewMsgStream}, 1)
|
||||
}
|
||||
|
||||
func testTimeTick(t *testing.T, f []Factory) {
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_TimeTick, 1))
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_TimeTick, 3))
|
||||
applyProduceAndConsume(t, &msgPack, []streamNewer{f[0].NewTtMsgStream, f[0].NewTtMsgStream}, 1)
|
||||
}
|
||||
|
||||
func testBroadCast(t *testing.T, f []Factory) {
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_TimeTick, 1))
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_TimeTick, 3))
|
||||
applyBroadCastAndConsume(t, &msgPack, []streamNewer{f[0].NewTtMsgStream, f[0].NewTtMsgStream}, 2)
|
||||
}
|
||||
|
||||
func testInsertWithRepack(t *testing.T, f []Factory) {
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_Insert, 1))
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_Insert, 3))
|
||||
applyProduceAndConsumeWithRepack(t, &msgPack, []streamNewer{f[0].NewMsgStream, f[0].NewMsgStream}, 2)
|
||||
}
|
||||
|
||||
func testInsertRepackFuncWithDifferentClient(t *testing.T, f []Factory) {
|
||||
baseMsg := BaseMsg{
|
||||
BeginTimestamp: 0,
|
||||
EndTimestamp: 0,
|
||||
HashValues: []uint32{1, 3},
|
||||
}
|
||||
|
||||
insertRequest := msgpb.InsertRequest{
|
||||
Base: &commonpb.MsgBase{
|
||||
MsgType: commonpb.MsgType_Insert,
|
||||
MsgID: 1,
|
||||
Timestamp: 1,
|
||||
SourceID: 1,
|
||||
},
|
||||
CollectionName: "Collection",
|
||||
PartitionName: "Partition",
|
||||
SegmentID: 1,
|
||||
ShardName: "1",
|
||||
Timestamps: []Timestamp{1, 1},
|
||||
RowIDs: []int64{1, 3},
|
||||
RowData: []*commonpb.Blob{{}, {}},
|
||||
}
|
||||
insertMsg := &InsertMsg{
|
||||
BaseMsg: baseMsg,
|
||||
InsertRequest: insertRequest,
|
||||
}
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, insertMsg)
|
||||
applyProduceAndConsumeWithRepack(t, &msgPack, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, 2)
|
||||
}
|
||||
|
||||
func testDeleteRepackFuncWithDifferentClient(t *testing.T, f []Factory) {
|
||||
baseMsg := BaseMsg{
|
||||
BeginTimestamp: 0,
|
||||
EndTimestamp: 0,
|
||||
HashValues: []uint32{1},
|
||||
}
|
||||
|
||||
deleteRequest := msgpb.DeleteRequest{
|
||||
Base: &commonpb.MsgBase{
|
||||
MsgType: commonpb.MsgType_Delete,
|
||||
MsgID: 1,
|
||||
Timestamp: 1,
|
||||
SourceID: 1,
|
||||
},
|
||||
CollectionName: "Collection",
|
||||
ShardName: "chan-1",
|
||||
Timestamps: []Timestamp{1},
|
||||
Int64PrimaryKeys: []int64{1},
|
||||
NumRows: 1,
|
||||
}
|
||||
deleteMsg := &DeleteMsg{
|
||||
BaseMsg: baseMsg,
|
||||
DeleteRequest: deleteRequest,
|
||||
}
|
||||
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, deleteMsg)
|
||||
applyProduceAndConsumeWithRepack(t, &msgPack, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, 2)
|
||||
}
|
||||
|
||||
func testDefaultRepack(t *testing.T, f []Factory) {
|
||||
msgPack := MsgPack{}
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_TimeTick, 1))
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_Insert, 2))
|
||||
msgPack.Msgs = append(msgPack.Msgs, getTsMsg(commonpb.MsgType_Delete, 3))
|
||||
|
||||
applyProduceAndConsumeWithRepack(t, &msgPack, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, 2)
|
||||
}
|
||||
|
||||
func testTimeTickerAndInsert(t *testing.T, f []Factory) {
|
||||
msgPack0 := MsgPack{}
|
||||
msgPack0.Msgs = append(msgPack0.Msgs, getTimeTickMsg(0))
|
||||
|
||||
msgPack1 := MsgPack{}
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 1))
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 3))
|
||||
|
||||
msgPack2 := MsgPack{}
|
||||
msgPack2.Msgs = append(msgPack2.Msgs, getTimeTickMsg(5))
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, getChannel(2))
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
var err error
|
||||
_, err = producer.Broadcast(&msgPack0)
|
||||
assert.Nil(t, err)
|
||||
err = producer.Produce(&msgPack1)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack2)
|
||||
assert.Nil(t, err)
|
||||
|
||||
receiveAndValidateMsg(ctx, consumer, len(msgPack1.Msgs))
|
||||
}
|
||||
|
||||
func testTimeTickerNoSeek(t *testing.T, f []Factory) {
|
||||
msgPack0 := MsgPack{}
|
||||
msgPack0.Msgs = append(msgPack0.Msgs, getTimeTickMsg(0))
|
||||
|
||||
msgPack1 := MsgPack{}
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 1))
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 19))
|
||||
|
||||
msgPack2 := MsgPack{}
|
||||
msgPack2.Msgs = append(msgPack2.Msgs, getTimeTickMsg(5))
|
||||
|
||||
msgPack3 := MsgPack{}
|
||||
msgPack3.Msgs = append(msgPack3.Msgs, getTsMsg(commonpb.MsgType_Insert, 14))
|
||||
msgPack3.Msgs = append(msgPack3.Msgs, getTsMsg(commonpb.MsgType_Insert, 9))
|
||||
|
||||
msgPack4 := MsgPack{}
|
||||
msgPack4.Msgs = append(msgPack4.Msgs, getTimeTickMsg(11))
|
||||
|
||||
msgPack5 := MsgPack{}
|
||||
msgPack5.Msgs = append(msgPack5.Msgs, getTimeTickMsg(15))
|
||||
|
||||
channels := getChannel(1)
|
||||
ctx := context.Background()
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewTtMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
|
||||
var err error
|
||||
_, err = producer.Broadcast(&msgPack0)
|
||||
assert.Nil(t, err)
|
||||
err = producer.Produce(&msgPack1)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack2)
|
||||
assert.Nil(t, err)
|
||||
err = producer.Produce(&msgPack3)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack4)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack5)
|
||||
assert.Nil(t, err)
|
||||
|
||||
o1 := consume(ctx, consumer)
|
||||
o2 := consume(ctx, consumer)
|
||||
o3 := consume(ctx, consumer)
|
||||
t.Log(o1.BeginTs)
|
||||
t.Log(o2.BeginTs)
|
||||
t.Log(o3.BeginTs)
|
||||
consumer.Close()
|
||||
|
||||
producer2, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewTtMsgStream}, channels)
|
||||
defer producer2.Close()
|
||||
defer consumer.Close()
|
||||
p1 := consume(ctx, consumer)
|
||||
p2 := consume(ctx, consumer)
|
||||
p3 := consume(ctx, consumer)
|
||||
t.Log(p1.BeginTs)
|
||||
t.Log(p2.BeginTs)
|
||||
t.Log(p3.BeginTs)
|
||||
|
||||
assert.Equal(t, o1.BeginTs, p1.BeginTs)
|
||||
assert.Equal(t, o2.BeginTs, p2.BeginTs)
|
||||
assert.Equal(t, o3.BeginTs, p3.BeginTs)
|
||||
}
|
||||
|
||||
func testSeekToLast(t *testing.T, f []Factory) {
|
||||
ctx := context.Background()
|
||||
channels := getChannel(1)
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
|
||||
msgPack := &MsgPack{}
|
||||
for i := 0; i < 10; i++ {
|
||||
insertMsg := getTsMsg(commonpb.MsgType_Insert, int64(i))
|
||||
msgPack.Msgs = append(msgPack.Msgs, insertMsg)
|
||||
}
|
||||
|
||||
// produce test data
|
||||
err := producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// pick a seekPosition
|
||||
var seekPosition *msgpb.MsgPosition
|
||||
for i := 0; i < 10; i++ {
|
||||
result := consume(ctx, consumer)
|
||||
assert.Equal(t, result.Msgs[0].ID(), int64(i))
|
||||
if i == 5 {
|
||||
seekPosition = result.EndPositions[0]
|
||||
}
|
||||
}
|
||||
consumer.Close()
|
||||
|
||||
// Create a unknown position consumer and seek.
|
||||
consumer = createAndSeekConsumer(ctx, t, f[1].NewMsgStream, channels, []*msgpb.MsgPosition{seekPosition})
|
||||
defer consumer.Close()
|
||||
|
||||
// Get latest MsgID.
|
||||
lastMsgID, err := consumer.GetLatestMsgID(channels[0])
|
||||
assert.Nil(t, err)
|
||||
|
||||
cnt := 0
|
||||
var value int64 = 6
|
||||
hasMore := true
|
||||
for hasMore {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
hasMore = false
|
||||
case msgPack, ok := <-consumer.Chan():
|
||||
if !ok {
|
||||
assert.Fail(t, "Should not reach here")
|
||||
}
|
||||
|
||||
assert.Equal(t, 1, len(msgPack.Msgs))
|
||||
for _, tsMsg := range msgPack.Msgs {
|
||||
assert.Equal(t, value, tsMsg.ID())
|
||||
value++
|
||||
cnt++
|
||||
|
||||
ret, err := lastMsgID.LessOrEqualThan(tsMsg.Position().MsgID)
|
||||
assert.Nil(t, err)
|
||||
if ret {
|
||||
hasMore = false
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
assert.Equal(t, 4, cnt)
|
||||
}
|
||||
|
||||
func testTimeTickerSeek(t *testing.T, f []Factory) {
|
||||
msgPack0 := MsgPack{}
|
||||
msgPack0.Msgs = append(msgPack0.Msgs, getTimeTickMsg(0))
|
||||
|
||||
msgPack1 := MsgPack{}
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 1))
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 3))
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 19))
|
||||
|
||||
msgPack2 := MsgPack{}
|
||||
msgPack2.Msgs = append(msgPack2.Msgs, getTimeTickMsg(5))
|
||||
|
||||
msgPack3 := MsgPack{}
|
||||
msgPack3.Msgs = append(msgPack3.Msgs, getTsMsg(commonpb.MsgType_Insert, 14))
|
||||
msgPack3.Msgs = append(msgPack3.Msgs, getTsMsg(commonpb.MsgType_Insert, 9))
|
||||
|
||||
msgPack4 := MsgPack{}
|
||||
msgPack4.Msgs = append(msgPack4.Msgs, getTimeTickMsg(11))
|
||||
|
||||
msgPack5 := MsgPack{}
|
||||
msgPack5.Msgs = append(msgPack5.Msgs, getTsMsg(commonpb.MsgType_Insert, 12))
|
||||
msgPack5.Msgs = append(msgPack5.Msgs, getTsMsg(commonpb.MsgType_Insert, 13))
|
||||
|
||||
msgPack6 := MsgPack{}
|
||||
msgPack6.Msgs = append(msgPack6.Msgs, getTimeTickMsg(15))
|
||||
|
||||
msgPack7 := MsgPack{}
|
||||
msgPack7.Msgs = append(msgPack7.Msgs, getTimeTickMsg(20))
|
||||
|
||||
ctx := context.Background()
|
||||
channels := getChannel(1)
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewTtMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
|
||||
// Send message
|
||||
_, err := producer.Broadcast(&msgPack0)
|
||||
assert.Nil(t, err)
|
||||
err = producer.Produce(&msgPack1)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack2)
|
||||
assert.Nil(t, err)
|
||||
err = producer.Produce(&msgPack3)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack4)
|
||||
assert.Nil(t, err)
|
||||
err = producer.Produce(&msgPack5)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack6)
|
||||
assert.Nil(t, err)
|
||||
_, err = producer.Broadcast(&msgPack7)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// Test received message
|
||||
receivedMsg := consume(ctx, consumer)
|
||||
assert.Equal(t, len(receivedMsg.Msgs), 2)
|
||||
assert.Equal(t, receivedMsg.BeginTs, uint64(0))
|
||||
assert.Equal(t, receivedMsg.EndTs, uint64(5))
|
||||
assert.Equal(t, receivedMsg.StartPositions[0].Timestamp, uint64(0))
|
||||
assert.Equal(t, receivedMsg.EndPositions[0].Timestamp, uint64(5))
|
||||
|
||||
receivedMsg2 := consume(ctx, consumer)
|
||||
assert.Equal(t, len(receivedMsg2.Msgs), 1)
|
||||
assert.Equal(t, receivedMsg2.BeginTs, uint64(5))
|
||||
assert.Equal(t, receivedMsg2.EndTs, uint64(11))
|
||||
assert.Equal(t, receivedMsg2.StartPositions[0].Timestamp, uint64(5))
|
||||
assert.Equal(t, receivedMsg2.EndPositions[0].Timestamp, uint64(11))
|
||||
|
||||
receivedMsg3 := consume(ctx, consumer)
|
||||
assert.Equal(t, len(receivedMsg3.Msgs), 3)
|
||||
assert.Equal(t, receivedMsg3.BeginTs, uint64(11))
|
||||
assert.Equal(t, receivedMsg3.EndTs, uint64(15))
|
||||
assert.Equal(t, receivedMsg3.StartPositions[0].Timestamp, uint64(11))
|
||||
assert.Equal(t, receivedMsg3.EndPositions[0].Timestamp, uint64(15))
|
||||
|
||||
receivedMsg4 := consume(ctx, consumer)
|
||||
assert.Equal(t, len(receivedMsg4.Msgs), 1)
|
||||
assert.Equal(t, receivedMsg4.BeginTs, uint64(15))
|
||||
assert.Equal(t, receivedMsg4.EndTs, uint64(20))
|
||||
assert.Equal(t, receivedMsg4.StartPositions[0].Timestamp, uint64(15))
|
||||
assert.Equal(t, receivedMsg4.EndPositions[0].Timestamp, uint64(20))
|
||||
consumer.Close()
|
||||
|
||||
consumer = createAndSeekConsumer(ctx, t, f[1].NewTtMsgStream, channels, receivedMsg3.StartPositions)
|
||||
seekMsg := consume(ctx, consumer)
|
||||
assert.Equal(t, len(seekMsg.Msgs), 3)
|
||||
result := []uint64{14, 12, 13}
|
||||
for i, msg := range seekMsg.Msgs {
|
||||
assert.Equal(t, msg.BeginTs(), result[i])
|
||||
}
|
||||
|
||||
seekMsg2 := consume(ctx, consumer)
|
||||
assert.Equal(t, len(seekMsg2.Msgs), 1)
|
||||
for _, msg := range seekMsg2.Msgs {
|
||||
assert.Equal(t, msg.BeginTs(), uint64(19))
|
||||
}
|
||||
consumer.Close()
|
||||
|
||||
consumer = createAndSeekConsumer(ctx, t, f[0].NewTtMsgStream, channels, receivedMsg3.EndPositions)
|
||||
seekMsg = consume(ctx, consumer)
|
||||
assert.Equal(t, len(seekMsg.Msgs), 1)
|
||||
for _, msg := range seekMsg.Msgs {
|
||||
assert.Equal(t, msg.BeginTs(), uint64(19))
|
||||
}
|
||||
consumer.Close()
|
||||
}
|
||||
|
||||
func testTimeTickUnmarshalHeader(t *testing.T, f []Factory) {
|
||||
msgPack0 := MsgPack{}
|
||||
msgPack0.Msgs = append(msgPack0.Msgs, getTimeTickMsg(0))
|
||||
|
||||
msgPack1 := MsgPack{}
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 1))
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 3))
|
||||
|
||||
msgPack2 := MsgPack{}
|
||||
msgPack2.Msgs = append(msgPack2.Msgs, getTimeTickMsg(5))
|
||||
|
||||
channels := getChannel(2)
|
||||
ctx := context.Background()
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewTtMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
_, err := producer.Broadcast(&msgPack0)
|
||||
require.NoErrorf(t, err, fmt.Sprintf("broadcast error = %v", err))
|
||||
|
||||
err = producer.Produce(&msgPack1)
|
||||
require.NoErrorf(t, err, fmt.Sprintf("produce error = %v", err))
|
||||
|
||||
_, err = producer.Broadcast(&msgPack2)
|
||||
require.NoErrorf(t, err, fmt.Sprintf("broadcast error = %v", err))
|
||||
|
||||
receiveAndValidateMsg(ctx, consumer, len(msgPack1.Msgs))
|
||||
}
|
||||
|
||||
func testTimeTickerStream1(t *testing.T, f []Factory) {
|
||||
consumeChannels := getChannel(2)
|
||||
pubChannel1 := []string{consumeChannels[0]}
|
||||
pubChannel2 := []string{consumeChannels[1]}
|
||||
ctx := context.Background()
|
||||
|
||||
producer1 := createProducer(ctx, t, f[0].NewMsgStream, pubChannel1)
|
||||
defer producer1.Close()
|
||||
msgPacks1 := createRandMsgPacks(3, 10, 10)
|
||||
assert.Nil(t, sendMsgPacks(producer1, msgPacks1))
|
||||
|
||||
producer2 := createProducer(ctx, t, f[0].NewMsgStream, pubChannel2)
|
||||
defer producer2.Close()
|
||||
msgPacks2 := createRandMsgPacks(5, 10, 10)
|
||||
assert.Nil(t, sendMsgPacks(producer2, msgPacks2))
|
||||
|
||||
// consume msg
|
||||
consumer := createConsumer(ctx, t, f[1].NewTtMsgStream, consumeChannels)
|
||||
defer consumer.Close()
|
||||
log.Println("===============receive msg=================")
|
||||
checkNMsgPack := func(t *testing.T, outputStream MsgStream, num int) int {
|
||||
rcvMsg := 0
|
||||
for i := 0; i < num; i++ {
|
||||
msgPack := consume(ctx, consumer)
|
||||
rcvMsg += len(msgPack.Msgs)
|
||||
if len(msgPack.Msgs) > 0 {
|
||||
for _, msg := range msgPack.Msgs {
|
||||
log.Println("msg type: ", msg.Type(), ", msg value: ", msg)
|
||||
assert.Greater(t, msg.BeginTs(), msgPack.BeginTs)
|
||||
assert.LessOrEqual(t, msg.BeginTs(), msgPack.EndTs)
|
||||
}
|
||||
log.Println("================")
|
||||
}
|
||||
}
|
||||
return rcvMsg
|
||||
}
|
||||
msgCount := checkNMsgPack(t, consumer, len(msgPacks1)/2)
|
||||
cnt1 := (len(msgPacks1)/2 - 1) * len(msgPacks1[0].Msgs)
|
||||
cnt2 := (len(msgPacks2)/2 - 1) * len(msgPacks2[0].Msgs)
|
||||
assert.Equal(t, (cnt1 + cnt2), msgCount)
|
||||
}
|
||||
|
||||
// This testcase will generate MsgPacks as following:
|
||||
//
|
||||
// Insert Insert Insert Insert Insert Insert
|
||||
//
|
||||
// c1 |----------|----------|----------|----------|----------|----------|
|
||||
//
|
||||
// ^ ^ ^ ^ ^ ^
|
||||
// TT(10) TT(20) TT(30) TT(40) TT(50) TT(100)
|
||||
//
|
||||
// Insert Insert Insert Insert Insert Insert
|
||||
//
|
||||
// c2 |----------|----------|----------|----------|----------|----------|
|
||||
//
|
||||
// ^ ^ ^ ^ ^ ^
|
||||
// TT(10) TT(20) TT(30) TT(40) TT(50) TT(100)
|
||||
//
|
||||
// Then check:
|
||||
// 1. ttMsgStream consumer can seek to the right position and resume
|
||||
// 2. The count of consumed msg should be equal to the count of produced msg
|
||||
func testTimeTickerStream2(t *testing.T, f []Factory) {
|
||||
consumeChannels := getChannel(2)
|
||||
pubChannel1 := []string{consumeChannels[0]}
|
||||
pubChannel2 := []string{consumeChannels[1]}
|
||||
ctx := context.Background()
|
||||
|
||||
producer1 := createProducer(ctx, t, f[0].NewMsgStream, pubChannel1)
|
||||
defer producer1.Close()
|
||||
msgPacks1 := createRandMsgPacks(3, 10, 10)
|
||||
assert.Nil(t, sendMsgPacks(producer1, msgPacks1))
|
||||
|
||||
producer2 := createProducer(ctx, t, f[0].NewMsgStream, pubChannel2)
|
||||
defer producer2.Close()
|
||||
msgPacks2 := createRandMsgPacks(5, 10, 10)
|
||||
assert.Nil(t, sendMsgPacks(producer2, msgPacks2))
|
||||
|
||||
// consume msg
|
||||
log.Println("=============receive msg===================")
|
||||
rcvMsgPacks := make([]*MsgPack, 0)
|
||||
|
||||
resumeMsgPack := func(t *testing.T) int {
|
||||
var consumer MsgStream
|
||||
msgCount := len(rcvMsgPacks)
|
||||
if msgCount == 0 {
|
||||
consumer = createConsumer(ctx, t, f[1].NewTtMsgStream, consumeChannels)
|
||||
} else {
|
||||
consumer = createAndSeekConsumer(ctx, t, f[1].NewTtMsgStream, consumeChannels, rcvMsgPacks[msgCount-1].EndPositions)
|
||||
}
|
||||
msgPack := consume(ctx, consumer)
|
||||
rcvMsgPacks = append(rcvMsgPacks, msgPack)
|
||||
if len(msgPack.Msgs) > 0 {
|
||||
for _, msg := range msgPack.Msgs {
|
||||
log.Println("msg type: ", msg.Type(), ", msg value: ", msg)
|
||||
assert.Greater(t, msg.BeginTs(), msgPack.BeginTs)
|
||||
assert.LessOrEqual(t, msg.BeginTs(), msgPack.EndTs)
|
||||
}
|
||||
log.Println("================")
|
||||
}
|
||||
consumer.Close()
|
||||
return len(rcvMsgPacks[msgCount].Msgs)
|
||||
}
|
||||
|
||||
msgCount := 0
|
||||
for i := 0; i < len(msgPacks1)/2; i++ {
|
||||
msgCount += resumeMsgPack(t)
|
||||
}
|
||||
cnt1 := (len(msgPacks1)/2 - 1) * len(msgPacks1[0].Msgs)
|
||||
cnt2 := (len(msgPacks2)/2 - 1) * len(msgPacks2[0].Msgs)
|
||||
assert.Equal(t, (cnt1 + cnt2), msgCount)
|
||||
}
|
||||
|
||||
func testMqMsgStreamSeek(t *testing.T, f []Factory) {
|
||||
channels := getChannel(1)
|
||||
ctx := context.Background()
|
||||
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
|
||||
msgPack := &MsgPack{}
|
||||
for i := 0; i < 10; i++ {
|
||||
insertMsg := getTsMsg(commonpb.MsgType_Insert, int64(i))
|
||||
msgPack.Msgs = append(msgPack.Msgs, insertMsg)
|
||||
}
|
||||
|
||||
err := producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
var seekPosition *msgpb.MsgPosition
|
||||
for i := 0; i < 10; i++ {
|
||||
result := consume(ctx, consumer)
|
||||
assert.Equal(t, result.Msgs[0].ID(), int64(i))
|
||||
if i == 5 {
|
||||
seekPosition = result.EndPositions[0]
|
||||
}
|
||||
}
|
||||
consumer.Close()
|
||||
|
||||
consumer = createAndSeekConsumer(ctx, t, f[0].NewMsgStream, channels, []*msgpb.MsgPosition{seekPosition})
|
||||
for i := 6; i < 10; i++ {
|
||||
result := consume(ctx, consumer)
|
||||
assert.Equal(t, result.Msgs[0].ID(), int64(i))
|
||||
}
|
||||
consumer.Close()
|
||||
}
|
||||
|
||||
func testMqMsgStreamSeekInvalidMessage(t *testing.T, f []Factory, pg positionGenerator) {
|
||||
channels := getChannel(1)
|
||||
ctx := context.Background()
|
||||
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
msgPack := &MsgPack{}
|
||||
for i := 0; i < 10; i++ {
|
||||
insertMsg := getTsMsg(commonpb.MsgType_Insert, int64(i))
|
||||
msgPack.Msgs = append(msgPack.Msgs, insertMsg)
|
||||
}
|
||||
|
||||
err := producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
var seekPosition *msgpb.MsgPosition
|
||||
for i := 0; i < 10; i++ {
|
||||
result := consume(ctx, consumer)
|
||||
assert.Equal(t, result.Msgs[0].ID(), int64(i))
|
||||
seekPosition = result.EndPositions[0]
|
||||
}
|
||||
|
||||
p := pg(seekPosition.ChannelName, seekPosition.Timestamp, seekPosition.MsgGroup, []uint64{13})
|
||||
consumer2 := createAndSeekConsumer(ctx, t, f[1].NewMsgStream, channels, p)
|
||||
defer consumer2.Close()
|
||||
|
||||
for i := 10; i < 20; i++ {
|
||||
insertMsg := getTsMsg(commonpb.MsgType_Insert, int64(i))
|
||||
msgPack.Msgs = append(msgPack.Msgs, insertMsg)
|
||||
}
|
||||
err = producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
result := consume(ctx, consumer2)
|
||||
assert.Equal(t, result.Msgs[0].ID(), int64(1))
|
||||
}
|
||||
|
||||
func testMqMsgStreamSeekLatest(t *testing.T, f []Factory) {
|
||||
channels := getChannel(1)
|
||||
ctx := context.Background()
|
||||
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
msgPack := &MsgPack{}
|
||||
for i := 0; i < 10; i++ {
|
||||
insertMsg := getTsMsg(commonpb.MsgType_Insert, int64(i))
|
||||
msgPack.Msgs = append(msgPack.Msgs, insertMsg)
|
||||
}
|
||||
|
||||
err := producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
consumer2 := createLatestConsumer(ctx, t, f[1].NewMsgStream, channels)
|
||||
defer consumer2.Close()
|
||||
|
||||
msgPack.Msgs = nil
|
||||
// produce another 10 tsMs
|
||||
for i := 10; i < 20; i++ {
|
||||
insertMsg := getTsMsg(commonpb.MsgType_Insert, int64(i))
|
||||
msgPack.Msgs = append(msgPack.Msgs, insertMsg)
|
||||
}
|
||||
err = producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
|
||||
for i := 10; i < 20; i++ {
|
||||
result := consume(ctx, consumer2)
|
||||
assert.Equal(t, result.Msgs[0].ID(), int64(i))
|
||||
}
|
||||
}
|
||||
|
||||
func testBroadcastMark(t *testing.T, f []Factory) {
|
||||
channels := getChannel(2)
|
||||
ctx := context.Background()
|
||||
|
||||
producer, consumer := createStream(ctx, t, []streamNewer{f[0].NewMsgStream, f[1].NewMsgStream}, channels)
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
msgPack0 := MsgPack{}
|
||||
msgPack0.Msgs = append(msgPack0.Msgs, getTimeTickMsg(0))
|
||||
|
||||
ids, err := producer.Broadcast(&msgPack0)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, ids)
|
||||
assert.Equal(t, len(channels), len(ids))
|
||||
for _, c := range channels {
|
||||
ids, ok := ids[c]
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, len(msgPack0.Msgs), len(ids))
|
||||
}
|
||||
|
||||
msgPack1 := MsgPack{}
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 1))
|
||||
msgPack1.Msgs = append(msgPack1.Msgs, getTsMsg(commonpb.MsgType_Insert, 3))
|
||||
|
||||
ids, err = producer.Broadcast(&msgPack1)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, ids)
|
||||
assert.Equal(t, len(channels), len(ids))
|
||||
for _, c := range channels {
|
||||
ids, ok := ids[c]
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, len(msgPack1.Msgs), len(ids))
|
||||
}
|
||||
|
||||
// edge cases
|
||||
_, err = producer.Broadcast(nil)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
msgPack2 := MsgPack{}
|
||||
msgPack2.Msgs = append(msgPack2.Msgs, &MarshalFailTsMsg{})
|
||||
_, err = producer.Broadcast(&msgPack2)
|
||||
assert.NotNil(t, err)
|
||||
}
|
||||
|
||||
func applyBroadCastAndConsume(t *testing.T, msgPack *MsgPack, newer []streamNewer, channelNum int) {
|
||||
producer, consumer := createStream(context.Background(), t, newer, getChannel(channelNum))
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
_, err := producer.Broadcast(msgPack)
|
||||
assert.Nil(t, err)
|
||||
receiveAndValidateMsg(context.Background(), consumer, len(msgPack.Msgs)*channelNum)
|
||||
}
|
||||
|
||||
func applyProduceAndConsumeWithRepack(
|
||||
t *testing.T,
|
||||
msgPack *MsgPack,
|
||||
newer []streamNewer,
|
||||
channelNum int,
|
||||
) {
|
||||
producer, consumer := createStream(context.Background(), t, newer, getChannel(channelNum))
|
||||
producer.SetRepackFunc(repackFunc)
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
err := producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
receiveAndValidateMsg(context.Background(), consumer, len(msgPack.Msgs))
|
||||
}
|
||||
|
||||
func applyProduceAndConsume(
|
||||
t *testing.T,
|
||||
msgPack *MsgPack,
|
||||
newer []streamNewer,
|
||||
channelNum int,
|
||||
) {
|
||||
producer, consumer := createStream(context.Background(), t, newer, getChannel(channelNum))
|
||||
defer producer.Close()
|
||||
defer consumer.Close()
|
||||
|
||||
err := producer.Produce(msgPack)
|
||||
assert.Nil(t, err)
|
||||
receiveAndValidateMsg(context.Background(), consumer, len(msgPack.Msgs))
|
||||
}
|
||||
|
||||
func consume(ctx context.Context, mq MsgStream) *MsgPack {
|
||||
for {
|
||||
select {
|
||||
case msgPack, ok := <-mq.Chan():
|
||||
if !ok {
|
||||
panic("Should not reach here")
|
||||
}
|
||||
return msgPack
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func createAndSeekConsumer(ctx context.Context, t *testing.T, newer streamNewer, channels []string, seekPositions []*msgpb.MsgPosition) MsgStream {
|
||||
consumer, err := newer(ctx)
|
||||
assert.Nil(t, err)
|
||||
consumer.AsConsumer(channels, funcutil.RandomString(8), mqwrapper.SubscriptionPositionUnknown)
|
||||
err = consumer.Seek(seekPositions)
|
||||
assert.Nil(t, err)
|
||||
return consumer
|
||||
}
|
||||
|
||||
func createProducer(ctx context.Context, t *testing.T, newer streamNewer, channels []string) MsgStream {
|
||||
producer, err := newer(ctx)
|
||||
assert.Nil(t, err)
|
||||
producer.AsProducer(channels)
|
||||
return producer
|
||||
}
|
||||
|
||||
func createConsumer(ctx context.Context, t *testing.T, newer streamNewer, channels []string) MsgStream {
|
||||
consumer, err := newer(ctx)
|
||||
assert.Nil(t, err)
|
||||
consumer.AsConsumer(channels, funcutil.RandomString(8), mqwrapper.SubscriptionPositionEarliest)
|
||||
return consumer
|
||||
}
|
||||
|
||||
func createLatestConsumer(ctx context.Context, t *testing.T, newer streamNewer, channels []string) MsgStream {
|
||||
consumer, err := newer(ctx)
|
||||
assert.Nil(t, err)
|
||||
consumer.AsConsumer(channels, funcutil.RandomString(8), mqwrapper.SubscriptionPositionLatest)
|
||||
return consumer
|
||||
}
|
||||
|
||||
func createStream(ctx context.Context, t *testing.T, newer []streamNewer, channels []string) (
|
||||
MsgStream, MsgStream,
|
||||
) {
|
||||
assert.NotEmpty(t, channels)
|
||||
producer, err := newer[0](ctx)
|
||||
assert.Nil(t, err)
|
||||
producer.AsProducer(channels)
|
||||
|
||||
consumer, err := newer[1](ctx)
|
||||
assert.Nil(t, err)
|
||||
consumer.AsConsumer(channels, funcutil.RandomString(8), mqwrapper.SubscriptionPositionEarliest)
|
||||
|
||||
return producer, consumer
|
||||
}
|
||||
|
||||
func getChannel(n int) []string {
|
||||
channels := make([]string, 0, n)
|
||||
for i := 0; i < n; i++ {
|
||||
channels = append(channels, funcutil.RandomString(8))
|
||||
}
|
||||
return channels
|
||||
}
|
||||
|
||||
func receiveAndValidateMsg(ctx context.Context, outputStream MsgStream, msgCount int) {
|
||||
receiveCount := 0
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case result, ok := <-outputStream.Chan():
|
||||
if !ok || result == nil || len(result.Msgs) == 0 {
|
||||
return
|
||||
}
|
||||
if len(result.Msgs) > 0 {
|
||||
msgs := result.Msgs
|
||||
for _, v := range msgs {
|
||||
receiveCount++
|
||||
log.Println("msg type: ", v.Type(), ", msg value: ", v)
|
||||
}
|
||||
log.Println("================")
|
||||
}
|
||||
if receiveCount >= msgCount {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
90
pkg/mq/msgstream/factory_test.go
Normal file
90
pkg/mq/msgstream/factory_test.go
Normal file
@ -0,0 +1,90 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package msgstream
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/msgpb"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper/nmq"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
)
|
||||
|
||||
type positionGenerator func(channelName string, timestamp uint64, msgGroup string, targetMsgIDs []uint64) []*msgpb.MsgPosition
|
||||
|
||||
func TestNmq(t *testing.T) {
|
||||
storeDir, err := os.MkdirTemp("", "milvus_mq_nmq")
|
||||
assert.Nil(t, err)
|
||||
defer os.RemoveAll(storeDir)
|
||||
|
||||
paramtable.Init()
|
||||
cfg := nmq.ParseServerOption(paramtable.Get())
|
||||
cfg.Opts.StoreDir = storeDir
|
||||
nmq.MustInitNatsMQ(cfg)
|
||||
defer nmq.CloseNatsMQ()
|
||||
|
||||
f1 := NewNatsmqFactory()
|
||||
f2 := NewNatsmqFactory()
|
||||
|
||||
client, err := nmq.NewClientWithDefaultOptions()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
testMQ(t, client, []Factory{f1, f2}, func(channelName string, timestamp uint64, msgGroup string, targetMsgIDs []uint64) []*msgpb.MsgPosition {
|
||||
result := make([]*msgpb.MsgPosition, 0, len(targetMsgIDs))
|
||||
|
||||
for _, targetMsgID := range targetMsgIDs {
|
||||
msgID := nmq.NewNmqID(targetMsgID).Serialize()
|
||||
result = append(result, &msgpb.MsgPosition{
|
||||
ChannelName: channelName,
|
||||
Timestamp: timestamp,
|
||||
MsgGroup: msgGroup,
|
||||
MsgID: msgID,
|
||||
})
|
||||
}
|
||||
return result
|
||||
})
|
||||
}
|
||||
|
||||
func testMQ(t *testing.T, client mqwrapper.Client, factories []Factory, pg positionGenerator) {
|
||||
testStreamOperation(t, client)
|
||||
testFactoryCommonOperation(t, factories[0])
|
||||
testMsgStreamOperation(t, factories, pg)
|
||||
}
|
||||
|
||||
// testFactoryOperation test common factory operation.
|
||||
func testFactoryCommonOperation(t *testing.T, f Factory) {
|
||||
var err error
|
||||
ctx := context.Background()
|
||||
_, err = f.NewMsgStream(ctx)
|
||||
assert.Nil(t, err)
|
||||
|
||||
_, err = f.NewTtMsgStream(ctx)
|
||||
assert.Nil(t, err)
|
||||
|
||||
_, err = f.NewQueryMsgStream(ctx)
|
||||
assert.Nil(t, err)
|
||||
|
||||
err = f.NewMsgStreamDisposer(ctx)([]string{"hello"}, "xx")
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
@ -29,6 +29,7 @@ import (
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
kafkawrapper "github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper/kafka"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper/nmq"
|
||||
pulsarmqwrapper "github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper/pulsar"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
"github.com/milvus-io/milvus/pkg/util/retry"
|
||||
@ -99,7 +100,6 @@ func (f *PmsFactory) NewTtMsgStream(ctx context.Context) (MsgStream, error) {
|
||||
|
||||
func (f *PmsFactory) getAuthentication() (pulsar.Authentication, error) {
|
||||
auth, err := pulsar.NewAuthentication(f.PulsarAuthPlugin, f.PulsarAuthParams)
|
||||
|
||||
if err != nil {
|
||||
log.Error("build authencation from config failed, please check it!",
|
||||
zap.String("authPlugin", f.PulsarAuthPlugin),
|
||||
@ -188,3 +188,15 @@ func NewKmsFactory(config *paramtable.KafkaConfig) Factory {
|
||||
}
|
||||
return f
|
||||
}
|
||||
|
||||
// NewNatsmqFactory create a new nats-mq factory.
|
||||
func NewNatsmqFactory() Factory {
|
||||
paramtable.Init()
|
||||
nmq.MustInitNatsMQ(nmq.ParseServerOption(paramtable.Get()))
|
||||
return &CommonFactory{
|
||||
Newer: nmq.NewClientWithDefaultOptions,
|
||||
DispatcherFactory: ProtoUDFactory{},
|
||||
ReceiveBufSize: 1024,
|
||||
MQBufSize: 1024,
|
||||
}
|
||||
}
|
||||
|
||||
148
pkg/mq/msgstream/mqwrapper/nmq/nmq_client.go
Normal file
148
pkg/mq/msgstream/mqwrapper/nmq/nmq_client.go
Normal file
@ -0,0 +1,148 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/nats-io/nats.go"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
)
|
||||
|
||||
// nmqClient implements mqwrapper.Client.
|
||||
var _ mqwrapper.Client = &nmqClient{}
|
||||
|
||||
// nmqClient contains a natsmq client
|
||||
type nmqClient struct {
|
||||
conn *nats.Conn
|
||||
}
|
||||
|
||||
// NewClientWithDefaultOptions returns a new NMQ client with default options.
|
||||
// It retrieves the NMQ client URL from the server configuration.
|
||||
func NewClientWithDefaultOptions() (mqwrapper.Client, error) {
|
||||
url := Nmq.ClientURL()
|
||||
log.Info("123123 ", zap.String("url", url))
|
||||
return NewClient(url)
|
||||
}
|
||||
|
||||
// NewClient returns a new nmqClient object
|
||||
func NewClient(url string, options ...nats.Option) (*nmqClient, error) {
|
||||
c, err := nats.Connect(url, options...)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to set nmq client")
|
||||
}
|
||||
return &nmqClient{conn: c}, nil
|
||||
}
|
||||
|
||||
// CreateProducer creates a producer for natsmq client
|
||||
func (nc *nmqClient) CreateProducer(options mqwrapper.ProducerOptions) (mqwrapper.Producer, error) {
|
||||
// TODO: inject jetstream options.
|
||||
js, err := nc.conn.JetStream()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to create jetstream context")
|
||||
}
|
||||
// TODO: (1) investigate on performance of multiple streams vs multiple topics.
|
||||
// (2) investigate if we should have topics under the same stream.
|
||||
_, err = js.AddStream(&nats.StreamConfig{
|
||||
Name: options.Topic,
|
||||
Subjects: []string{options.Topic},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to add/connect to jetstream for producer")
|
||||
}
|
||||
rp := nmqProducer{js: js, topic: options.Topic}
|
||||
return &rp, nil
|
||||
}
|
||||
|
||||
func (nc *nmqClient) Subscribe(options mqwrapper.ConsumerOptions) (mqwrapper.Consumer, error) {
|
||||
if options.Topic == "" {
|
||||
return nil, fmt.Errorf("invalid consumer config: empty topic")
|
||||
}
|
||||
|
||||
if options.SubscriptionName == "" {
|
||||
return nil, fmt.Errorf("invalid consumer config: empty subscription name")
|
||||
}
|
||||
// TODO: inject jetstream options.
|
||||
js, err := nc.conn.JetStream()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to create jetstream context")
|
||||
}
|
||||
// TODO: do we allow passing in an existing natsChan from options?
|
||||
// also, revisit the size or make it a user param
|
||||
natsChan := make(chan *nats.Msg, options.BufSize)
|
||||
// TODO: should we allow subscribe to a topic that doesn't exist yet? Current logic allows it.
|
||||
_, err = js.AddStream(&nats.StreamConfig{
|
||||
Name: options.Topic,
|
||||
Subjects: []string{options.Topic},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to add/connect to jetstream for consumer")
|
||||
}
|
||||
closeChan := make(chan struct{})
|
||||
|
||||
var sub *nats.Subscription
|
||||
position := options.SubscriptionInitialPosition
|
||||
// TODO: should we only allow exclusive subscribe? Current logic allows double subscribe.
|
||||
switch position {
|
||||
case mqwrapper.SubscriptionPositionLatest:
|
||||
sub, err = js.ChanSubscribe(options.Topic, natsChan, nats.DeliverNew())
|
||||
case mqwrapper.SubscriptionPositionEarliest:
|
||||
sub, err = js.ChanSubscribe(options.Topic, natsChan, nats.DeliverAll())
|
||||
}
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, fmt.Sprintf("failed to get consumer info, subscribe position: %d", position))
|
||||
}
|
||||
|
||||
return &Consumer{
|
||||
js: js,
|
||||
sub: sub,
|
||||
topic: options.Topic,
|
||||
groupName: options.SubscriptionName,
|
||||
options: options,
|
||||
natsChan: natsChan,
|
||||
closeChan: closeChan,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// EarliestMessageID returns the earliest message ID for nmq client
|
||||
func (nc *nmqClient) EarliestMessageID() mqwrapper.MessageID {
|
||||
return &nmqID{messageID: 1}
|
||||
}
|
||||
|
||||
// StringToMsgID converts string id to MessageID
|
||||
func (nc *nmqClient) StringToMsgID(id string) (mqwrapper.MessageID, error) {
|
||||
rID, err := strconv.ParseUint(id, 10, 64)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to parse string to MessageID")
|
||||
}
|
||||
return &nmqID{messageID: rID}, nil
|
||||
}
|
||||
|
||||
// BytesToMsgID converts a byte array to messageID
|
||||
func (nc *nmqClient) BytesToMsgID(id []byte) (mqwrapper.MessageID, error) {
|
||||
rID := DeserializeNmqID(id)
|
||||
return &nmqID{messageID: rID}, nil
|
||||
}
|
||||
|
||||
func (nc *nmqClient) Close() {
|
||||
nc.conn.Close()
|
||||
}
|
||||
237
pkg/mq/msgstream/mqwrapper/nmq/nmq_client_test.go
Normal file
237
pkg/mq/msgstream/mqwrapper/nmq/nmq_client_test.go
Normal file
@ -0,0 +1,237 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
)
|
||||
|
||||
func createNmqClient() (*nmqClient, error) {
|
||||
return NewClient(natsServerAddress)
|
||||
}
|
||||
|
||||
func Test_NewNmqClient(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, client)
|
||||
client.Close()
|
||||
}
|
||||
|
||||
func TestNmqClient_CreateProducer(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, client)
|
||||
defer client.Close()
|
||||
|
||||
topic := "TestNmqClient_CreateProducer"
|
||||
proOpts := mqwrapper.ProducerOptions{Topic: topic}
|
||||
producer, err := client.CreateProducer(proOpts)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, producer)
|
||||
defer producer.Close()
|
||||
|
||||
nmqProducer := producer.(*nmqProducer)
|
||||
assert.Equal(t, nmqProducer.Topic(), topic)
|
||||
|
||||
msg := &mqwrapper.ProducerMessage{
|
||||
Payload: []byte{},
|
||||
Properties: nil,
|
||||
}
|
||||
_, err = nmqProducer.Send(context.TODO(), msg)
|
||||
assert.Nil(t, err)
|
||||
|
||||
invalidOpts := mqwrapper.ProducerOptions{Topic: ""}
|
||||
producer, e := client.CreateProducer(invalidOpts)
|
||||
assert.Nil(t, producer)
|
||||
assert.Error(t, e)
|
||||
}
|
||||
|
||||
func TestNmqClient_GetLatestMsg(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.Nil(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := fmt.Sprintf("t2GetLatestMsg-%d", rand.Int())
|
||||
proOpts := mqwrapper.ProducerOptions{Topic: topic}
|
||||
producer, err := client.CreateProducer(proOpts)
|
||||
assert.Nil(t, err)
|
||||
defer producer.Close()
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
msg := &mqwrapper.ProducerMessage{
|
||||
Payload: []byte{byte(i)},
|
||||
Properties: nil,
|
||||
}
|
||||
_, err = producer.Send(context.TODO(), msg)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
subName := "subName"
|
||||
consumerOpts := mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: subName,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
}
|
||||
|
||||
consumer, err := client.Subscribe(consumerOpts)
|
||||
assert.Nil(t, err)
|
||||
|
||||
expectLastMsg, err := consumer.GetLatestMsgID()
|
||||
assert.Nil(t, err)
|
||||
|
||||
var actualLastMsg mqwrapper.Message
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
for i := 0; i < 10; i++ {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
fmt.Println(i)
|
||||
assert.FailNow(t, "consumer failed to yield message in 100 milliseconds")
|
||||
case msg := <-consumer.Chan():
|
||||
consumer.Ack(msg)
|
||||
actualLastMsg = msg
|
||||
}
|
||||
}
|
||||
require.NotNil(t, actualLastMsg)
|
||||
ret, err := expectLastMsg.LessOrEqualThan(actualLastMsg.ID().Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, ret)
|
||||
}
|
||||
|
||||
func TestNmqClient_IllegalSubscribe(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, client)
|
||||
defer client.Close()
|
||||
|
||||
sub, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: "",
|
||||
})
|
||||
assert.Nil(t, sub)
|
||||
assert.Error(t, err)
|
||||
|
||||
sub, err = client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: "123",
|
||||
SubscriptionName: "",
|
||||
})
|
||||
assert.Nil(t, sub)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestNmqClient_Subscribe(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, client)
|
||||
defer client.Close()
|
||||
|
||||
topic := "TestNmqClient_Subscribe"
|
||||
proOpts := mqwrapper.ProducerOptions{Topic: topic}
|
||||
producer, err := client.CreateProducer(proOpts)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, producer)
|
||||
defer producer.Close()
|
||||
|
||||
subName := "subName"
|
||||
consumerOpts := mqwrapper.ConsumerOptions{
|
||||
Topic: "",
|
||||
SubscriptionName: subName,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
}
|
||||
|
||||
consumer, err := client.Subscribe(consumerOpts)
|
||||
assert.NotNil(t, err)
|
||||
assert.Nil(t, consumer)
|
||||
|
||||
consumerOpts.Topic = topic
|
||||
consumer, err = client.Subscribe(consumerOpts)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, consumer)
|
||||
defer consumer.Close()
|
||||
assert.Equal(t, consumer.Subscription(), subName)
|
||||
|
||||
msg := &mqwrapper.ProducerMessage{
|
||||
Payload: []byte{1},
|
||||
Properties: nil,
|
||||
}
|
||||
_, err = producer.Send(context.TODO(), msg)
|
||||
assert.Nil(t, err)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
assert.FailNow(t, "consumer failed to yield message in 100 milliseconds")
|
||||
case msg := <-consumer.Chan():
|
||||
consumer.Ack(msg)
|
||||
nmqmsg := msg.(*nmqMessage)
|
||||
msgPayload := nmqmsg.Payload()
|
||||
assert.NotEmpty(t, msgPayload)
|
||||
msgTopic := nmqmsg.Topic()
|
||||
assert.Equal(t, msgTopic, topic)
|
||||
msgProp := nmqmsg.Properties()
|
||||
assert.Empty(t, msgProp)
|
||||
msgID := nmqmsg.ID()
|
||||
rID := msgID.(*nmqID)
|
||||
assert.Equal(t, rID.messageID, MessageIDType(1))
|
||||
}
|
||||
}
|
||||
|
||||
func TestNmqClient_EarliestMessageID(t *testing.T) {
|
||||
client, _ := createNmqClient()
|
||||
defer client.Close()
|
||||
|
||||
mid := client.EarliestMessageID()
|
||||
assert.NotNil(t, mid)
|
||||
nmqmsg := mid.(*nmqID)
|
||||
assert.Equal(t, nmqmsg.messageID, MessageIDType(1))
|
||||
}
|
||||
|
||||
func TestNmqClient_StringToMsgID(t *testing.T) {
|
||||
client, _ := createNmqClient()
|
||||
defer client.Close()
|
||||
|
||||
str := "5"
|
||||
res, err := client.StringToMsgID(str)
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, res)
|
||||
|
||||
str = "X"
|
||||
res, err = client.StringToMsgID(str)
|
||||
assert.Nil(t, res)
|
||||
assert.NotNil(t, err)
|
||||
}
|
||||
|
||||
func TestNmqClient_BytesToMsgID(t *testing.T) {
|
||||
client, _ := createNmqClient()
|
||||
defer client.Close()
|
||||
|
||||
mid := client.EarliestMessageID()
|
||||
res, err := client.BytesToMsgID(mid.Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, res)
|
||||
}
|
||||
185
pkg/mq/msgstream/mqwrapper/nmq/nmq_consumer.go
Normal file
185
pkg/mq/msgstream/mqwrapper/nmq/nmq_consumer.go
Normal file
@ -0,0 +1,185 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/nats-io/nats.go"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"github.com/milvus-io/milvus/pkg/util/merr"
|
||||
)
|
||||
|
||||
// Consumer is a client that used to consume messages from natsmq
|
||||
type Consumer struct {
|
||||
options mqwrapper.ConsumerOptions
|
||||
js nats.JetStreamContext
|
||||
sub *nats.Subscription
|
||||
topic string
|
||||
groupName string
|
||||
natsChan chan *nats.Msg
|
||||
msgChan chan mqwrapper.Message
|
||||
closeChan chan struct{}
|
||||
once sync.Once
|
||||
closeOnce sync.Once
|
||||
skip bool
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// Subscription returns the subscription name of this consumer
|
||||
func (nc *Consumer) Subscription() string {
|
||||
return nc.groupName
|
||||
}
|
||||
|
||||
// Chan returns a channel to read messages from natsmq
|
||||
func (nc *Consumer) Chan() <-chan mqwrapper.Message {
|
||||
if err := nc.closed(); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if nc.sub == nil {
|
||||
log.Error("accessing Chan of an uninitialized subscription.", zap.String("topic", nc.topic), zap.String("groupName", nc.groupName))
|
||||
panic("failed to chan a consumer without assign")
|
||||
}
|
||||
if nc.msgChan == nil {
|
||||
nc.once.Do(func() {
|
||||
nc.msgChan = make(chan mqwrapper.Message, 256)
|
||||
nc.wg.Add(1)
|
||||
go func() {
|
||||
defer nc.wg.Done()
|
||||
for {
|
||||
select {
|
||||
case msg := <-nc.natsChan:
|
||||
if nc.skip {
|
||||
nc.skip = false
|
||||
continue
|
||||
}
|
||||
nc.msgChan <- &nmqMessage{
|
||||
raw: msg,
|
||||
}
|
||||
case <-nc.closeChan:
|
||||
log.Info("close nmq consumer ", zap.String("topic", nc.topic), zap.String("groupName", nc.groupName))
|
||||
close(nc.msgChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
})
|
||||
}
|
||||
return nc.msgChan
|
||||
}
|
||||
|
||||
// Seek is used to seek the position in natsmq topic
|
||||
func (nc *Consumer) Seek(id mqwrapper.MessageID, inclusive bool) error {
|
||||
if err := nc.closed(); err != nil {
|
||||
return err
|
||||
}
|
||||
if nc.sub != nil {
|
||||
return errors.New("can not seek() on an initilized consumer")
|
||||
}
|
||||
if nc.msgChan != nil {
|
||||
return errors.New("Seek should be called before Chan")
|
||||
}
|
||||
log.Info("Seek is called", zap.String("topic", nc.topic), zap.Any("id", id))
|
||||
msgID := id.(*nmqID).messageID
|
||||
// skip the first message when consume
|
||||
nc.skip = !inclusive
|
||||
var err error
|
||||
nc.sub, err = nc.js.ChanSubscribe(nc.topic, nc.natsChan, nats.StartSequence(msgID))
|
||||
if err != nil {
|
||||
log.Warn("fail to Seek", zap.Error(err))
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// Ack is used to ask a natsmq message
|
||||
func (nc *Consumer) Ack(message mqwrapper.Message) {
|
||||
if err := message.(*nmqMessage).raw.Ack(); err != nil {
|
||||
log.Warn("failed to ack message of nmq", zap.String("topic", message.Topic()), zap.Reflect("msgID", message.ID()))
|
||||
}
|
||||
}
|
||||
|
||||
// Close is used to free the resources of this consumer
|
||||
func (nc *Consumer) Close() {
|
||||
nc.closeOnce.Do(func() {
|
||||
if err := nc.sub.Unsubscribe(); err != nil {
|
||||
log.Warn("failed to unsubscribe subscription of nmq", zap.String("topic", nc.topic))
|
||||
}
|
||||
close(nc.closeChan)
|
||||
nc.wg.Wait()
|
||||
})
|
||||
}
|
||||
|
||||
// GetLatestMsgID returns the ID of the most recent message processed by the consumer.
|
||||
func (nc *Consumer) GetLatestMsgID() (mqwrapper.MessageID, error) {
|
||||
if err := nc.closed(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
info, err := nc.js.StreamInfo(nc.topic)
|
||||
if err != nil {
|
||||
log.Warn("fail to get stream info of nats", zap.String("topic", nc.topic), zap.Error(err))
|
||||
return nil, errors.Wrap(err, "failed to get stream info of nats jetstream")
|
||||
}
|
||||
msgID := info.State.LastSeq
|
||||
return &nmqID{messageID: msgID}, nil
|
||||
}
|
||||
|
||||
// CheckTopicValid verifies if the given topic is valid for this consumer.
|
||||
// 1. topic should exist.
|
||||
// 2. topic should be empty.
|
||||
func (nc *Consumer) CheckTopicValid(topic string) error {
|
||||
if err := nc.closed(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// A consumer is tied to a topic. In a multi-tenant situation,
|
||||
// a consumer is not supposed to check on other topics.
|
||||
if topic != nc.topic {
|
||||
return fmt.Errorf("consumer of topic %s checking validness of topic %s", nc.topic, topic)
|
||||
}
|
||||
|
||||
// check if topic valid or exist.
|
||||
streamInfo, err := nc.js.StreamInfo(topic)
|
||||
if errors.Is(err, nats.ErrStreamNotFound) {
|
||||
return merr.WrapErrTopicNotFound(topic, err.Error())
|
||||
} else if err != nil {
|
||||
log.Warn("fail to get stream info of nats", zap.String("topic", nc.topic), zap.Error(err))
|
||||
return errors.Wrap(err, "failed to get stream info of nats jetstream")
|
||||
}
|
||||
|
||||
// check if topic stream is empty.
|
||||
if streamInfo.State.Msgs > 0 {
|
||||
return merr.WrapErrTopicNotEmpty(topic, "stream in nats is not empty")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Closed check if Consumer is closed.
|
||||
func (nc *Consumer) closed() error {
|
||||
select {
|
||||
case <-nc.closeChan:
|
||||
return errors.Newf("Closed Nmq Consumer, topic: %s, subscription name:", nc.topic, nc.groupName)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
470
pkg/mq/msgstream/mqwrapper/nmq/nmq_consumer_test.go
Normal file
470
pkg/mq/msgstream/mqwrapper/nmq/nmq_consumer_test.go
Normal file
@ -0,0 +1,470 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"github.com/milvus-io/milvus/pkg/util/merr"
|
||||
)
|
||||
|
||||
func TestNatsConsumer_Subscription(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := t.Name()
|
||||
proOpts := mqwrapper.ProducerOptions{Topic: topic}
|
||||
_, err = client.CreateProducer(proOpts)
|
||||
assert.NoError(t, err)
|
||||
|
||||
consumer, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, consumer)
|
||||
defer consumer.Close()
|
||||
|
||||
str := consumer.Subscription()
|
||||
assert.NotNil(t, str)
|
||||
}
|
||||
|
||||
func Test_GetEarliestMessageID(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
mid := client.EarliestMessageID()
|
||||
|
||||
assert.NotNil(t, mid)
|
||||
assert.Equal(t, mid.(*nmqID).messageID, MessageIDType(1))
|
||||
}
|
||||
|
||||
func Test_BadLatestMessageID(t *testing.T) {
|
||||
topic := t.Name()
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
consumer, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
consumer.Close()
|
||||
id, err := consumer.GetLatestMsgID()
|
||||
assert.Nil(t, id)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestComsumeMessage(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := t.Name()
|
||||
p, err := client.CreateProducer(mqwrapper.ProducerOptions{Topic: topic})
|
||||
assert.NoError(t, err)
|
||||
|
||||
c, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
defer c.Close()
|
||||
|
||||
msg := []byte("test the first message")
|
||||
prop := map[string]string{"k1": "v1", "k2": "v2"}
|
||||
_, err = p.Send(context.Background(), &mqwrapper.ProducerMessage{
|
||||
Payload: msg,
|
||||
Properties: prop,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
recvMsg, err := c.(*Consumer).sub.NextMsg(1 * time.Second)
|
||||
assert.NoError(t, err)
|
||||
recvMsg.Ack()
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, msg, recvMsg.Data)
|
||||
properties := make(map[string]string)
|
||||
for k, vs := range recvMsg.Header {
|
||||
if len(vs) > 0 {
|
||||
properties[k] = vs[0]
|
||||
}
|
||||
}
|
||||
assert.True(t, reflect.DeepEqual(prop, properties))
|
||||
|
||||
msg2 := []byte("test the second message")
|
||||
prop2 := map[string]string{"k1": "v3", "k4": "v4"}
|
||||
_, err = p.Send(context.Background(), &mqwrapper.ProducerMessage{
|
||||
Payload: msg2,
|
||||
Properties: prop2,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
recvMsg, err = c.(*Consumer).sub.NextMsg(1 * time.Second)
|
||||
assert.NoError(t, err)
|
||||
recvMsg.Ack()
|
||||
assert.Equal(t, msg2, recvMsg.Data)
|
||||
properties = make(map[string]string)
|
||||
for k, vs := range recvMsg.Header {
|
||||
if len(vs) > 0 {
|
||||
properties[k] = vs[0]
|
||||
}
|
||||
}
|
||||
assert.True(t, reflect.DeepEqual(prop2, properties))
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, c)
|
||||
}
|
||||
|
||||
func TestNatsConsumer_Close(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := t.Name()
|
||||
c, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, c)
|
||||
|
||||
str := c.Subscription()
|
||||
assert.NotNil(t, str)
|
||||
|
||||
c.Close()
|
||||
|
||||
// Disallow double close.
|
||||
assert.Panics(t, func() { c.Chan() })
|
||||
assert.Error(t, c.Seek(NewNmqID(1), false))
|
||||
assert.Error(t, func() error { _, err := c.GetLatestMsgID(); return err }())
|
||||
|
||||
c.Close() // Allow double close, nothing happened.
|
||||
}
|
||||
|
||||
func TestNatsClientErrorOnUnsubscribeTwice(t *testing.T) {
|
||||
topic := t.Name()
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
consumer, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
defer consumer.Close()
|
||||
|
||||
err = consumer.(*Consumer).sub.Unsubscribe()
|
||||
assert.NoError(t, err)
|
||||
err = consumer.(*Consumer).sub.Unsubscribe()
|
||||
assert.True(t, strings.Contains(err.Error(), "invalid subscription"))
|
||||
t.Log(err)
|
||||
}
|
||||
|
||||
func TestCheckTopicValid(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := t.Name()
|
||||
consumer, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, consumer)
|
||||
|
||||
str := consumer.Subscription()
|
||||
assert.NotNil(t, str)
|
||||
|
||||
// empty topic should pass
|
||||
err = consumer.CheckTopicValid(topic)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Not allowed to check other topic's validness.
|
||||
err = consumer.CheckTopicValid("BadTopic")
|
||||
assert.Error(t, err)
|
||||
|
||||
// non empty topic should fail
|
||||
pub, err := client.CreateProducer(mqwrapper.ProducerOptions{
|
||||
Topic: topic,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
_, err = pub.Send(context.TODO(), &mqwrapper.ProducerMessage{
|
||||
Payload: []byte("123123123"),
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = consumer.CheckTopicValid(topic)
|
||||
assert.ErrorIs(t, err, merr.ErrTopicNotEmpty)
|
||||
|
||||
consumer.Close()
|
||||
err = consumer.CheckTopicValid(topic)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func newTestConsumer(topic string, fromEarliest bool, fromLatest bool) (mqwrapper.Consumer, error) {
|
||||
conn, err := nats.Connect(natsServerAddress)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
js, err := conn.JetStream()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
groupName := topic
|
||||
_, err = js.AddStream(&nats.StreamConfig{
|
||||
Name: groupName,
|
||||
Subjects: []string{topic},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
natsChan := make(chan *nats.Msg, 8192)
|
||||
closeChan := make(chan struct{})
|
||||
return &Consumer{
|
||||
js: js,
|
||||
topic: topic,
|
||||
groupName: groupName,
|
||||
natsChan: natsChan,
|
||||
closeChan: closeChan,
|
||||
skip: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func newProducer(t *testing.T, topic string) (*nmqClient, mqwrapper.Producer) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
producer, err := client.CreateProducer(mqwrapper.ProducerOptions{Topic: topic})
|
||||
assert.NoError(t, err)
|
||||
return client, producer
|
||||
}
|
||||
|
||||
func process(t *testing.T, msgs []string, p mqwrapper.Producer) {
|
||||
for _, msg := range msgs {
|
||||
_, err := p.Send(context.Background(), &mqwrapper.ProducerMessage{
|
||||
Payload: []byte(msg),
|
||||
Properties: map[string]string{},
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNmqConsumer_GetLatestMsgID(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := t.Name()
|
||||
p, err := client.CreateProducer(mqwrapper.ProducerOptions{Topic: topic})
|
||||
assert.NoError(t, err)
|
||||
|
||||
c, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
defer c.Close()
|
||||
|
||||
latestMsgID, err := c.GetLatestMsgID()
|
||||
assert.NoError(t, err)
|
||||
|
||||
msgs := []string{"111", "222", "333", "444", "555"}
|
||||
newMessageID := latestMsgID.(*nmqID).messageID + uint64(len(msgs))
|
||||
process(t, msgs, p)
|
||||
latestMsgID, err = c.GetLatestMsgID()
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, newMessageID, latestMsgID.(*nmqID).messageID)
|
||||
}
|
||||
|
||||
func TestNmqConsumer_ConsumeFromLatest(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := t.Name()
|
||||
p, err := client.CreateProducer(mqwrapper.ProducerOptions{Topic: topic})
|
||||
assert.NoError(t, err)
|
||||
|
||||
msgs := []string{"111", "222", "333"}
|
||||
process(t, msgs, p)
|
||||
|
||||
c, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionLatest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
defer c.Close()
|
||||
|
||||
msgs = []string{"444", "555"}
|
||||
process(t, msgs, p)
|
||||
|
||||
msg := <-c.Chan()
|
||||
assert.Equal(t, "444", string(msg.Payload()))
|
||||
msg = <-c.Chan()
|
||||
assert.Equal(t, "555", string(msg.Payload()))
|
||||
}
|
||||
|
||||
func TestNmqConsumer_ConsumeFromEarliest(t *testing.T) {
|
||||
client, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
topic := t.Name()
|
||||
p, err := client.CreateProducer(mqwrapper.ProducerOptions{Topic: topic})
|
||||
assert.NoError(t, err)
|
||||
|
||||
msgs := []string{"111", "222"}
|
||||
process(t, msgs, p)
|
||||
|
||||
c, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
defer c.Close()
|
||||
|
||||
msgs = []string{"333", "444", "555"}
|
||||
process(t, msgs, p)
|
||||
|
||||
msg := <-c.Chan()
|
||||
assert.Equal(t, "111", string(msg.Payload()))
|
||||
msg = <-c.Chan()
|
||||
assert.Equal(t, "222", string(msg.Payload()))
|
||||
|
||||
c2, err := client.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topic,
|
||||
SubscriptionName: topic,
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
defer c2.Close()
|
||||
|
||||
msgs = []string{"777"}
|
||||
process(t, msgs, p)
|
||||
|
||||
msg = <-c2.Chan()
|
||||
assert.Equal(t, "111", string(msg.Payload()))
|
||||
msg = <-c2.Chan()
|
||||
assert.Equal(t, "222", string(msg.Payload()))
|
||||
}
|
||||
|
||||
func TestNatsConsumer_SeekExclusive(t *testing.T) {
|
||||
topic := t.Name()
|
||||
c, p := newProducer(t, topic)
|
||||
defer c.Close()
|
||||
defer p.Close()
|
||||
|
||||
msgs := []string{"111", "222", "333", "444", "555"}
|
||||
process(t, msgs, p)
|
||||
|
||||
msgID := &nmqID{messageID: 2}
|
||||
consumer, err := newTestConsumer(topic, false, false)
|
||||
assert.NoError(t, err)
|
||||
defer consumer.Close()
|
||||
err = consumer.Seek(msgID, false)
|
||||
assert.NoError(t, err)
|
||||
|
||||
msg := <-consumer.Chan()
|
||||
assert.Equal(t, "333", string(msg.Payload()))
|
||||
msg = <-consumer.Chan()
|
||||
assert.Equal(t, "444", string(msg.Payload()))
|
||||
}
|
||||
|
||||
func TestNatsConsumer_SeekInclusive(t *testing.T) {
|
||||
topic := t.Name()
|
||||
c, p := newProducer(t, topic)
|
||||
defer c.Close()
|
||||
defer p.Close()
|
||||
|
||||
msgs := []string{"111", "222", "333", "444", "555"}
|
||||
process(t, msgs, p)
|
||||
|
||||
msgID := &nmqID{messageID: 2}
|
||||
consumer, err := newTestConsumer(topic, false, false)
|
||||
assert.NoError(t, err)
|
||||
defer consumer.Close()
|
||||
err = consumer.Seek(msgID, true)
|
||||
assert.NoError(t, err)
|
||||
|
||||
msg := <-consumer.Chan()
|
||||
assert.Equal(t, "222", string(msg.Payload()))
|
||||
msg = <-consumer.Chan()
|
||||
assert.Equal(t, "333", string(msg.Payload()))
|
||||
}
|
||||
|
||||
func TestNatsConsumer_NoDoubleSeek(t *testing.T) {
|
||||
topic := t.Name()
|
||||
c, p := newProducer(t, topic)
|
||||
defer c.Close()
|
||||
defer p.Close()
|
||||
|
||||
msgID := &nmqID{messageID: 2}
|
||||
consumer, err := newTestConsumer(topic, false, false)
|
||||
assert.NoError(t, err)
|
||||
defer consumer.Close()
|
||||
err = consumer.Seek(msgID, true)
|
||||
assert.NoError(t, err)
|
||||
err = consumer.Seek(msgID, true)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestNatsConsumer_ChanWithNoAssign(t *testing.T) {
|
||||
topic := t.Name()
|
||||
c, p := newProducer(t, topic)
|
||||
defer c.Close()
|
||||
defer p.Close()
|
||||
|
||||
msgs := []string{"111", "222", "333", "444", "555"}
|
||||
process(t, msgs, p)
|
||||
|
||||
consumer, err := newTestConsumer(topic, false, false)
|
||||
assert.NoError(t, err)
|
||||
defer consumer.Close()
|
||||
|
||||
assert.Panics(t, func() {
|
||||
<-consumer.Chan()
|
||||
})
|
||||
}
|
||||
69
pkg/mq/msgstream/mqwrapper/nmq/nmq_id.go
Normal file
69
pkg/mq/msgstream/mqwrapper/nmq/nmq_id.go
Normal file
@ -0,0 +1,69 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"github.com/milvus-io/milvus/pkg/common"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
)
|
||||
|
||||
// MessageIDType is a type alias for server.UniqueID that represents the ID of a Nmq message.
|
||||
type MessageIDType = uint64
|
||||
|
||||
// nmqID wraps message ID for natsmq
|
||||
type nmqID struct {
|
||||
messageID MessageIDType
|
||||
}
|
||||
|
||||
// Check if nmqID implements MessageID interface
|
||||
var _ mqwrapper.MessageID = &nmqID{}
|
||||
|
||||
// NewNmqID creates and returns a new instance of the nmqID struct with the given MessageID.
|
||||
func NewNmqID(id MessageIDType) mqwrapper.MessageID {
|
||||
return &nmqID{
|
||||
messageID: id,
|
||||
}
|
||||
}
|
||||
|
||||
// Serialize convert nmq message id to []byte
|
||||
func (nid *nmqID) Serialize() []byte {
|
||||
return SerializeNmqID(nid.messageID)
|
||||
}
|
||||
|
||||
func (nid *nmqID) AtEarliestPosition() bool {
|
||||
return nid.messageID <= 1
|
||||
}
|
||||
|
||||
func (nid *nmqID) LessOrEqualThan(msgID []byte) (bool, error) {
|
||||
return nid.messageID <= DeserializeNmqID(msgID), nil
|
||||
}
|
||||
|
||||
func (nid *nmqID) Equal(msgID []byte) (bool, error) {
|
||||
return nid.messageID == DeserializeNmqID(msgID), nil
|
||||
}
|
||||
|
||||
// SerializeNmqID is used to serialize a message ID to byte array
|
||||
func SerializeNmqID(messageID MessageIDType) []byte {
|
||||
b := make([]byte, 8)
|
||||
common.Endian.PutUint64(b, messageID)
|
||||
return b
|
||||
}
|
||||
|
||||
// DeserializeNmqID is used to deserialize a message ID from byte array
|
||||
func DeserializeNmqID(messageID []byte) MessageIDType {
|
||||
return common.Endian.Uint64(messageID)
|
||||
}
|
||||
102
pkg/mq/msgstream/mqwrapper/nmq/nmq_id_test.go
Normal file
102
pkg/mq/msgstream/mqwrapper/nmq/nmq_id_test.go
Normal file
@ -0,0 +1,102 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestNmqID_Serialize(t *testing.T) {
|
||||
rid := &nmqID{
|
||||
messageID: 8,
|
||||
}
|
||||
|
||||
bin := rid.Serialize()
|
||||
assert.NotNil(t, bin)
|
||||
assert.NotZero(t, len(bin))
|
||||
}
|
||||
|
||||
func Test_AtEarliestPosition(t *testing.T) {
|
||||
rid := &nmqID{
|
||||
messageID: 0,
|
||||
}
|
||||
assert.True(t, rid.AtEarliestPosition())
|
||||
|
||||
rid = &nmqID{
|
||||
messageID: math.MaxInt64,
|
||||
}
|
||||
assert.False(t, rid.AtEarliestPosition())
|
||||
}
|
||||
|
||||
func TestLessOrEqualThan(t *testing.T) {
|
||||
rid1 := &nmqID{
|
||||
messageID: 0,
|
||||
}
|
||||
rid2 := &nmqID{
|
||||
messageID: math.MaxInt64,
|
||||
}
|
||||
|
||||
ret, err := rid1.LessOrEqualThan(rid2.Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, ret)
|
||||
|
||||
ret, err = rid2.LessOrEqualThan(rid1.Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, ret)
|
||||
|
||||
ret, err = rid1.LessOrEqualThan(rid1.Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, ret)
|
||||
}
|
||||
|
||||
func Test_Equal(t *testing.T) {
|
||||
rid1 := &nmqID{
|
||||
messageID: 0,
|
||||
}
|
||||
|
||||
rid2 := &nmqID{
|
||||
messageID: math.MaxInt64,
|
||||
}
|
||||
|
||||
{
|
||||
ret, err := rid1.Equal(rid1.Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, ret)
|
||||
|
||||
}
|
||||
|
||||
{
|
||||
ret, err := rid1.Equal(rid2.Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, ret)
|
||||
}
|
||||
}
|
||||
|
||||
func Test_SerializeNmqID(t *testing.T) {
|
||||
bin := SerializeNmqID(10)
|
||||
assert.NotNil(t, bin)
|
||||
assert.NotZero(t, len(bin))
|
||||
}
|
||||
|
||||
func Test_DeserializeNmqID(t *testing.T) {
|
||||
bin := SerializeNmqID(5)
|
||||
id := DeserializeNmqID(bin)
|
||||
assert.Equal(t, id, MessageIDType(5))
|
||||
}
|
||||
75
pkg/mq/msgstream/mqwrapper/nmq/nmq_message.go
Normal file
75
pkg/mq/msgstream/mqwrapper/nmq/nmq_message.go
Normal file
@ -0,0 +1,75 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"log"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"github.com/nats-io/nats.go"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Check nmqMessage implements ConsumerMessage
|
||||
var (
|
||||
_ mqwrapper.Message = (*nmqMessage)(nil)
|
||||
)
|
||||
|
||||
// nmqMessage wraps the message for natsmq
|
||||
type nmqMessage struct {
|
||||
raw *nats.Msg
|
||||
|
||||
// lazy initialized field.
|
||||
meta *nats.MsgMetadata
|
||||
}
|
||||
|
||||
// Topic returns the topic name of natsmq message
|
||||
func (nm *nmqMessage) Topic() string {
|
||||
// TODO: Dependency: implement of subscription logic of nmq.
|
||||
// 1:1 Subject:Topic model is appied on this implementation.
|
||||
// M:N model should be a optimize option in future.
|
||||
return nm.raw.Subject
|
||||
}
|
||||
|
||||
// Properties returns the properties of natsmq message
|
||||
func (nm *nmqMessage) Properties() map[string]string {
|
||||
properties := make(map[string]string, len(nm.raw.Header))
|
||||
for k, vs := range nm.raw.Header {
|
||||
if len(vs) > 0 {
|
||||
properties[k] = vs[0]
|
||||
}
|
||||
}
|
||||
return properties
|
||||
}
|
||||
|
||||
// Payload returns the payload of natsmq message
|
||||
func (nm *nmqMessage) Payload() []byte {
|
||||
return nm.raw.Data
|
||||
}
|
||||
|
||||
// ID returns the id of natsmq message
|
||||
func (nm *nmqMessage) ID() mqwrapper.MessageID {
|
||||
if nm.meta == nil {
|
||||
var err error
|
||||
// raw is always a jetstream message, should never fail.
|
||||
nm.meta, err = nm.raw.Metadata()
|
||||
if err != nil {
|
||||
log.Fatal("raw is always a jetstream message, should never fail", zap.Error(err))
|
||||
}
|
||||
}
|
||||
return &nmqID{messageID: nm.meta.Sequence.Stream}
|
||||
}
|
||||
44
pkg/mq/msgstream/mqwrapper/nmq/nmq_message_test.go
Normal file
44
pkg/mq/msgstream/mqwrapper/nmq/nmq_message_test.go
Normal file
@ -0,0 +1,44 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestNmqMessage_All(t *testing.T) {
|
||||
topic := t.Name()
|
||||
raw := nats.NewMsg(topic)
|
||||
raw.Data = []byte(`test payload`)
|
||||
raw.Header.Add("test", "test")
|
||||
nm := nmqMessage{
|
||||
meta: &nats.MsgMetadata{
|
||||
Sequence: nats.SequencePair{
|
||||
Stream: 12,
|
||||
},
|
||||
},
|
||||
raw: raw,
|
||||
}
|
||||
payload := []byte("test payload")
|
||||
assert.Equal(t, topic, nm.Topic())
|
||||
assert.Equal(t, MessageIDType(12), nm.ID().(*nmqID).messageID)
|
||||
assert.Equal(t, payload, nm.Payload())
|
||||
assert.Equal(t, nm.Properties()["test"], "test")
|
||||
}
|
||||
66
pkg/mq/msgstream/mqwrapper/nmq/nmq_producer.go
Normal file
66
pkg/mq/msgstream/mqwrapper/nmq/nmq_producer.go
Normal file
@ -0,0 +1,66 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"github.com/nats-io/nats.go"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
var _ mqwrapper.Producer = (*nmqProducer)(nil)
|
||||
|
||||
// nmqProducer contains a natsmq producer
|
||||
type nmqProducer struct {
|
||||
js nats.JetStreamContext
|
||||
topic string
|
||||
}
|
||||
|
||||
// Topic returns the topic of nmq producer
|
||||
func (np *nmqProducer) Topic() string {
|
||||
return np.topic
|
||||
}
|
||||
|
||||
// Send send the producer messages to natsmq
|
||||
func (np *nmqProducer) Send(ctx context.Context, message *mqwrapper.ProducerMessage) (mqwrapper.MessageID, error) {
|
||||
// Encode message
|
||||
msg := &nats.Msg{
|
||||
Subject: np.topic,
|
||||
Header: make(nats.Header, len(message.Properties)),
|
||||
Data: message.Payload,
|
||||
}
|
||||
for k, v := range message.Properties {
|
||||
msg.Header.Add(k, v)
|
||||
}
|
||||
|
||||
// publish to nats-server
|
||||
pa, err := np.js.PublishMsg(msg)
|
||||
if err != nil {
|
||||
log.Warn("failed to publish message by nmq", zap.String("topic", np.topic), zap.Error(err), zap.Int("payload_size", len(message.Payload)))
|
||||
return nil, err
|
||||
}
|
||||
return &nmqID{messageID: pa.Sequence}, err
|
||||
}
|
||||
|
||||
// Close does nothing currently
|
||||
func (np *nmqProducer) Close() {
|
||||
// No specific producer to be closed.
|
||||
// stream doesn't close here.
|
||||
}
|
||||
47
pkg/mq/msgstream/mqwrapper/nmq/nmq_producer_test.go
Normal file
47
pkg/mq/msgstream/mqwrapper/nmq/nmq_producer_test.go
Normal file
@ -0,0 +1,47 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
)
|
||||
|
||||
func TestNatsMQProducer(t *testing.T) {
|
||||
c, err := createNmqClient()
|
||||
assert.NoError(t, err)
|
||||
defer c.Close()
|
||||
topic := t.Name()
|
||||
pOpts := mqwrapper.ProducerOptions{Topic: topic}
|
||||
|
||||
// Check Topic()
|
||||
p, err := c.CreateProducer(pOpts)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, p.(*nmqProducer).Topic(), topic)
|
||||
|
||||
// Check Send()
|
||||
msg := &mqwrapper.ProducerMessage{
|
||||
Payload: []byte{},
|
||||
Properties: map[string]string{},
|
||||
}
|
||||
_, err = p.Send(context.TODO(), msg)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
93
pkg/mq/msgstream/mqwrapper/nmq/nmq_server.go
Normal file
93
pkg/mq/msgstream/mqwrapper/nmq/nmq_server.go
Normal file
@ -0,0 +1,93 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/nats-io/nats-server/v2/server"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
)
|
||||
|
||||
// Nmq is global natsmq instance that will be initialized only once
|
||||
var Nmq *server.Server
|
||||
|
||||
// once is used to init global natsmq
|
||||
var once sync.Once
|
||||
|
||||
// NatsMQConfig is used to initialize NatsMQ.
|
||||
type NatsMQConfig struct {
|
||||
Opts server.Options
|
||||
InitializeTimeout time.Duration
|
||||
}
|
||||
|
||||
// MustInitNatsMQ init global local natsmq instance.
|
||||
// Panic if initailizing operation failed.
|
||||
func MustInitNatsMQ(cfg *NatsMQConfig) {
|
||||
once.Do(func() {
|
||||
log.Info("try to initialize global nmq", zap.Any("config", cfg))
|
||||
var err error
|
||||
Nmq, err = server.NewServer(&cfg.Opts)
|
||||
if err != nil {
|
||||
log.Fatal("fail to initailize nmq", zap.Error(err))
|
||||
}
|
||||
|
||||
log.Info("initialize nmq finished", zap.Error(err))
|
||||
// Start Nmq in background and wait until it's ready for connection.
|
||||
go Nmq.Start()
|
||||
// Wait for server to be ready for connections
|
||||
if !Nmq.ReadyForConnections(cfg.InitializeTimeout) {
|
||||
log.Fatal("nmq is not ready")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ParseServerOption get nats server option from paramstable.
|
||||
func ParseServerOption(params *paramtable.ComponentParam) *NatsMQConfig {
|
||||
return &NatsMQConfig{
|
||||
Opts: server.Options{
|
||||
Host: "127.0.0.1", // Force to use loopback address.
|
||||
Port: params.NatsmqCfg.ServerPort.GetAsInt(),
|
||||
MaxPayload: params.NatsmqCfg.ServerMaxPayload.GetAsInt32(),
|
||||
MaxPending: params.NatsmqCfg.ServerMaxPending.GetAsInt64(),
|
||||
JetStream: true,
|
||||
JetStreamMaxStore: params.NatsmqCfg.ServerMaxFileStore.GetAsInt64(),
|
||||
StoreDir: params.NatsmqCfg.ServerStoreDir.GetValue(),
|
||||
Debug: params.NatsmqCfg.ServerMonitorDebug.GetAsBool(),
|
||||
Logtime: params.NatsmqCfg.ServerMonitorLogTime.GetAsBool(),
|
||||
LogFile: params.NatsmqCfg.ServerMonitorLogFile.GetValue(),
|
||||
LogSizeLimit: params.NatsmqCfg.ServerMonitorLogSizeLimit.GetAsInt64(),
|
||||
},
|
||||
InitializeTimeout: time.Duration(params.NatsmqCfg.ServerInitializeTimeout.GetAsInt()) * time.Millisecond,
|
||||
}
|
||||
}
|
||||
|
||||
// CloseNatsMQ is used to close global natsmq
|
||||
func CloseNatsMQ() {
|
||||
log.Debug("Closing Natsmq!")
|
||||
if Nmq != nil {
|
||||
// Shut down the server.
|
||||
Nmq.Shutdown()
|
||||
// Wait for server shutdown.
|
||||
Nmq.WaitForShutdown()
|
||||
Nmq = nil
|
||||
}
|
||||
}
|
||||
60
pkg/mq/msgstream/mqwrapper/nmq/nmq_server_test.go
Normal file
60
pkg/mq/msgstream/mqwrapper/nmq/nmq_server_test.go
Normal file
@ -0,0 +1,60 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package nmq
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
var natsServerAddress string
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
paramtable.Init()
|
||||
|
||||
storeDir, _ := os.MkdirTemp("", "milvus_mq_nmq")
|
||||
defer os.RemoveAll(storeDir)
|
||||
|
||||
cfg := ParseServerOption(paramtable.Get())
|
||||
cfg.Opts.StoreDir = storeDir
|
||||
MustInitNatsMQ(cfg)
|
||||
defer CloseNatsMQ()
|
||||
|
||||
natsServerAddress = Nmq.ClientURL()
|
||||
exitCode := m.Run()
|
||||
os.Exit(exitCode)
|
||||
}
|
||||
|
||||
func TestGetServerOptionDefault(t *testing.T) {
|
||||
cfg := ParseServerOption(paramtable.Get())
|
||||
assert.Equal(t, "127.0.0.1", cfg.Opts.Host)
|
||||
assert.Equal(t, 4222, cfg.Opts.Port)
|
||||
assert.Equal(t, true, cfg.Opts.JetStream)
|
||||
assert.Equal(t, "/var/lib/milvus/nats", cfg.Opts.StoreDir)
|
||||
assert.Equal(t, int64(17179869184), cfg.Opts.JetStreamMaxStore)
|
||||
assert.Equal(t, int32(8388608), cfg.Opts.MaxPayload)
|
||||
assert.Equal(t, int64(67108864), cfg.Opts.MaxPending)
|
||||
assert.Equal(t, 4000*time.Millisecond, cfg.InitializeTimeout)
|
||||
assert.Equal(t, false, cfg.Opts.Debug)
|
||||
assert.Equal(t, true, cfg.Opts.Logtime)
|
||||
assert.Equal(t, "", cfg.Opts.LogFile)
|
||||
assert.Equal(t, int64(0), cfg.Opts.LogSizeLimit)
|
||||
}
|
||||
261
pkg/mq/msgstream/stream_test.go
Normal file
261
pkg/mq/msgstream/stream_test.go
Normal file
@ -0,0 +1,261 @@
|
||||
package msgstream
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math/rand"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
||||
"github.com/milvus-io/milvus/pkg/util/funcutil"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func testStreamOperation(t *testing.T, mqClient mqwrapper.Client) {
|
||||
testFuncs := []func(t *testing.T, c mqwrapper.Client){
|
||||
testConcurrentStream,
|
||||
testConcurrentStreamAndSubscribeLast,
|
||||
testConcurrentStreamAndSeekInclusive,
|
||||
testConcurrentStreamAndSeekNoInclusive,
|
||||
testConcurrentStreamAndSeekToLast,
|
||||
}
|
||||
|
||||
for _, testFunc := range testFuncs {
|
||||
t.Run(
|
||||
runtime.FuncForPC(reflect.ValueOf(testFunc).Pointer()).Name(),
|
||||
func(t *testing.T) {
|
||||
testFunc(t, mqClient)
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func testConcurrentStream(t *testing.T, mqClient mqwrapper.Client) {
|
||||
topics := getChannel(2)
|
||||
|
||||
producer, err := mqClient.CreateProducer(mqwrapper.ProducerOptions{
|
||||
Topic: topics[0],
|
||||
})
|
||||
defer producer.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
consumer, err := mqClient.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topics[0],
|
||||
SubscriptionName: funcutil.RandomString(8),
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
defer consumer.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
testSendAndRecv(t, producer, consumer)
|
||||
}
|
||||
|
||||
func testConcurrentStreamAndSubscribeLast(t *testing.T, mqClient mqwrapper.Client) {
|
||||
topics := getChannel(2)
|
||||
|
||||
producer, err := mqClient.CreateProducer(mqwrapper.ProducerOptions{
|
||||
Topic: topics[0],
|
||||
})
|
||||
defer producer.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
ids := sendMessages(context.Background(), t, producer, generateRandMessage(1024, 1000))
|
||||
|
||||
consumer, err := mqClient.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topics[0],
|
||||
SubscriptionName: funcutil.RandomString(8),
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionLatest,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
latestID, err := consumer.GetLatestMsgID()
|
||||
assert.Nil(t, err)
|
||||
|
||||
compare, err := ids[len(ids)-1].Equal(latestID.Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, compare)
|
||||
|
||||
defer consumer.Close()
|
||||
testSendAndRecv(t, producer, consumer)
|
||||
}
|
||||
|
||||
func testConcurrentStreamAndSeekInclusive(t *testing.T, mqClient mqwrapper.Client) {
|
||||
topics := getChannel(2)
|
||||
|
||||
producer, err := mqClient.CreateProducer(mqwrapper.ProducerOptions{
|
||||
Topic: topics[0],
|
||||
})
|
||||
defer producer.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
cases := generateRandMessage(1024, 1000)
|
||||
ids := sendMessages(context.Background(), t, producer, cases)
|
||||
|
||||
consumer, err := mqClient.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topics[0],
|
||||
SubscriptionName: funcutil.RandomString(8),
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionUnknown,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
defer consumer.Close()
|
||||
|
||||
// seek half and inclusive.
|
||||
// consume all and compare.
|
||||
half := len(ids) / 2
|
||||
ids = ids[half:]
|
||||
err = consumer.Seek(ids[0], true)
|
||||
assert.Nil(t, err)
|
||||
consumerIDs := recvMessages(context.Background(), t, consumer, cases[half:], time.Minute)
|
||||
compareMultiIDs(t, ids, consumerIDs)
|
||||
assert.Empty(t, recvMessages(context.Background(), t, consumer, cases, 5*time.Second))
|
||||
|
||||
testSendAndRecv(t, producer, consumer)
|
||||
}
|
||||
|
||||
func testConcurrentStreamAndSeekNoInclusive(t *testing.T, mqClient mqwrapper.Client) {
|
||||
topics := getChannel(2)
|
||||
|
||||
producer, err := mqClient.CreateProducer(mqwrapper.ProducerOptions{
|
||||
Topic: topics[0],
|
||||
})
|
||||
defer producer.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
cases := generateRandMessage(1024, 1000)
|
||||
ids := sendMessages(context.Background(), t, producer, cases)
|
||||
|
||||
consumer, err := mqClient.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topics[0],
|
||||
SubscriptionName: funcutil.RandomString(8),
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionUnknown,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
defer consumer.Close()
|
||||
|
||||
// seek half and inclusive.
|
||||
// consume all and compare.
|
||||
half := len(ids) / 2
|
||||
ids = ids[half:]
|
||||
err = consumer.Seek(ids[0], false)
|
||||
assert.Nil(t, err)
|
||||
consumerIDs := recvMessages(context.Background(), t, consumer, cases[half+1:], time.Minute)
|
||||
compareMultiIDs(t, ids[1:], consumerIDs)
|
||||
assert.Empty(t, recvMessages(context.Background(), t, consumer, cases, 5*time.Second))
|
||||
|
||||
testSendAndRecv(t, producer, consumer)
|
||||
}
|
||||
|
||||
func testConcurrentStreamAndSeekToLast(t *testing.T, mqClient mqwrapper.Client) {
|
||||
topics := getChannel(2)
|
||||
|
||||
producer, err := mqClient.CreateProducer(mqwrapper.ProducerOptions{
|
||||
Topic: topics[0],
|
||||
})
|
||||
defer producer.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
cases := generateRandMessage(1024, 1000)
|
||||
sendMessages(context.Background(), t, producer, cases)
|
||||
|
||||
consumer, err := mqClient.Subscribe(mqwrapper.ConsumerOptions{
|
||||
Topic: topics[0],
|
||||
SubscriptionName: funcutil.RandomString(8),
|
||||
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionUnknown,
|
||||
BufSize: 1024,
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
defer consumer.Close()
|
||||
latestID, err := consumer.GetLatestMsgID()
|
||||
assert.Nil(t, err)
|
||||
|
||||
// seek half and inclusive.
|
||||
// consume all and compare.
|
||||
err = consumer.Seek(latestID, false)
|
||||
assert.Nil(t, err)
|
||||
testSendAndRecv(t, producer, consumer)
|
||||
}
|
||||
|
||||
func testSendAndRecv(t *testing.T, p mqwrapper.Producer, c mqwrapper.Consumer) {
|
||||
ctx := context.Background()
|
||||
msg := generateRandMessage(1024*5, 10)
|
||||
|
||||
var (
|
||||
producerIDs []mqwrapper.MessageID
|
||||
consumerIDs []mqwrapper.MessageID
|
||||
)
|
||||
|
||||
wg := sync.WaitGroup{}
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
producerIDs = sendMessages(ctx, t, p, msg)
|
||||
log.Debug("producing finished", zap.Any("id", producerIDs[0].Serialize()), zap.Any("ids", producerIDs))
|
||||
}()
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
consumerIDs = recvMessages(ctx, t, c, msg, 10*time.Second)
|
||||
}()
|
||||
wg.Wait()
|
||||
|
||||
compareMultiIDs(t, producerIDs, consumerIDs)
|
||||
// should be empty.
|
||||
assert.Empty(t, recvMessages(context.Background(), t, c, msg, time.Second))
|
||||
}
|
||||
|
||||
func compareMultiIDs(t *testing.T, producerIDs []mqwrapper.MessageID, consumerIDs []mqwrapper.MessageID) {
|
||||
assert.Equal(t, len(producerIDs), len(consumerIDs))
|
||||
for i := range producerIDs {
|
||||
compare, err := producerIDs[i].Equal(consumerIDs[i].Serialize())
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, compare)
|
||||
}
|
||||
}
|
||||
|
||||
func generateRandMessage(m int, n int) []string {
|
||||
cases := make([]string, 0, n)
|
||||
for i := 0; i < n; i++ {
|
||||
l := rand.Intn(m)
|
||||
cases = append(cases, funcutil.RandomString(l))
|
||||
}
|
||||
return cases
|
||||
}
|
||||
|
||||
func sendMessages(ctx context.Context, t *testing.T, p mqwrapper.Producer, testCase []string) []mqwrapper.MessageID {
|
||||
ids := make([]mqwrapper.MessageID, 0, len(testCase))
|
||||
for _, s := range testCase {
|
||||
id, err := p.Send(ctx, &mqwrapper.ProducerMessage{
|
||||
Payload: []byte(s),
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func recvMessages(ctx context.Context, t *testing.T, c mqwrapper.Consumer, testCase []string, timeout time.Duration) []mqwrapper.MessageID {
|
||||
ids := make([]mqwrapper.MessageID, 0, len(testCase))
|
||||
timeoutTicker := time.NewTicker(timeout)
|
||||
defer timeoutTicker.Stop()
|
||||
for {
|
||||
select {
|
||||
case msg := <-c.Chan():
|
||||
ids = append(ids, msg.ID())
|
||||
c.Ack(msg)
|
||||
case <-timeoutTicker.C:
|
||||
return ids
|
||||
}
|
||||
if len(ids) >= len(testCase) {
|
||||
break
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
@ -120,7 +120,7 @@ func TestGlobalMethods(t *testing.T) {
|
||||
err = ZstdCompress(strings.NewReader(data), compressed, zstd.WithWindowSize(3))
|
||||
assert.Error(t, err)
|
||||
|
||||
err = ZstdDecompress(compressed, origin, zstd.WithDecoderConcurrency(0))
|
||||
err = ZstdDecompress(compressed, origin, zstd.WithDecoderConcurrency(-1))
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
|
||||
@ -48,9 +48,11 @@ type ServiceParam struct {
|
||||
MetaStoreCfg MetaStoreConfig
|
||||
EtcdCfg EtcdConfig
|
||||
DBCfg MetaDBConfig
|
||||
MQCfg MQConfig
|
||||
PulsarCfg PulsarConfig
|
||||
KafkaCfg KafkaConfig
|
||||
RocksmqCfg RocksmqConfig
|
||||
NatsmqCfg NatsmqConfig
|
||||
MinioCfg MinioConfig
|
||||
}
|
||||
|
||||
@ -61,9 +63,11 @@ func (p *ServiceParam) init() {
|
||||
p.MetaStoreCfg.Init(&p.BaseTable)
|
||||
p.EtcdCfg.Init(&p.BaseTable)
|
||||
p.DBCfg.Init(&p.BaseTable)
|
||||
p.MQCfg.Init(&p.BaseTable)
|
||||
p.PulsarCfg.Init(&p.BaseTable)
|
||||
p.KafkaCfg.Init(&p.BaseTable)
|
||||
p.RocksmqCfg.Init(&p.BaseTable)
|
||||
p.NatsmqCfg.Init(&p.BaseTable)
|
||||
p.MinioCfg.Init(&p.BaseTable)
|
||||
}
|
||||
|
||||
@ -71,6 +75,11 @@ func (p *ServiceParam) RocksmqEnable() bool {
|
||||
return p.RocksmqCfg.Path.GetValue() != ""
|
||||
}
|
||||
|
||||
// NatsmqEnable checks if NATS messaging queue is enabled.
|
||||
func (p *ServiceParam) NatsmqEnable() bool {
|
||||
return p.NatsmqCfg.ServerStoreDir.GetValue() != ""
|
||||
}
|
||||
|
||||
func (p *ServiceParam) PulsarEnable() bool {
|
||||
return p.PulsarCfg.Address.GetValue() != ""
|
||||
}
|
||||
@ -353,7 +362,27 @@ func (p *MetaDBConfig) Init(base *BaseTable) {
|
||||
Export: true,
|
||||
}
|
||||
p.MaxIdleConns.Init(base.mgr)
|
||||
}
|
||||
|
||||
// /////////////////////////////////////////////////////////////////////////////
|
||||
// --- mq ---
|
||||
|
||||
// MQConfig represents the configuration settings for the message queue.
|
||||
type MQConfig struct {
|
||||
Type ParamItem `refreshable:"false"`
|
||||
}
|
||||
|
||||
// Init initializes the MQConfig object with a BaseTable.
|
||||
func (p *MQConfig) Init(base *BaseTable) {
|
||||
p.Type = ParamItem{
|
||||
Key: "mq.type",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "",
|
||||
Doc: `Default value: "default"
|
||||
Valid values: [default, pulsar, kafka, rocksmq, natsmq]`,
|
||||
Export: true,
|
||||
}
|
||||
p.Type.Init(base.mgr)
|
||||
}
|
||||
|
||||
// /////////////////////////////////////////////////////////////////////////////
|
||||
@ -477,7 +506,6 @@ func (p *PulsarConfig) Init(base *BaseTable) {
|
||||
},
|
||||
}
|
||||
p.AuthParams.Init(base.mgr)
|
||||
|
||||
}
|
||||
|
||||
// --- kafka ---
|
||||
@ -625,6 +653,104 @@ please adjust in embedded Milvus: /tmp/milvus/rdb_data`,
|
||||
r.TickerTimeInSeconds.Init(base.mgr)
|
||||
}
|
||||
|
||||
// NatsmqConfig describes the configuration options for the Nats message queue
|
||||
type NatsmqConfig struct {
|
||||
ServerPort ParamItem `refreshable:"false"`
|
||||
ServerStoreDir ParamItem `refreshable:"false"`
|
||||
ServerMaxFileStore ParamItem `refreshable:"false"`
|
||||
ServerMaxPayload ParamItem `refreshable:"false"`
|
||||
ServerMaxPending ParamItem `refreshable:"false"`
|
||||
ServerInitializeTimeout ParamItem `refreshable:"false"`
|
||||
ServerMonitorDebug ParamItem `refreshable:"false"`
|
||||
ServerMonitorLogTime ParamItem `refreshable:"false"`
|
||||
ServerMonitorLogFile ParamItem `refreshable:"false"`
|
||||
ServerMonitorLogSizeLimit ParamItem `refreshable:"false"`
|
||||
}
|
||||
|
||||
// Init sets up a new NatsmqConfig instance using the provided BaseTable
|
||||
func (r *NatsmqConfig) Init(base *BaseTable) {
|
||||
r.ServerPort = ParamItem{
|
||||
Key: "natsmq.server.port",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "4222",
|
||||
Doc: `Port for nats server listening`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerPort.Init(base.mgr)
|
||||
r.ServerStoreDir = ParamItem{
|
||||
Key: "natsmq.server.storeDir",
|
||||
Version: "2.3.0",
|
||||
Doc: `Directory to use for JetStream storage of nats`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerStoreDir.Init(base.mgr)
|
||||
r.ServerMaxFileStore = ParamItem{
|
||||
Key: "natsmq.server.maxFileStore",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "17179869184",
|
||||
Doc: `Maximum size of the 'file' storage`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerMaxFileStore.Init(base.mgr)
|
||||
r.ServerMaxPayload = ParamItem{
|
||||
Key: "natsmq.server.maxPayload",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "8388608",
|
||||
Doc: `Maximum number of bytes in a message payload`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerMaxPayload.Init(base.mgr)
|
||||
r.ServerMaxPending = ParamItem{
|
||||
Key: "natsmq.server.maxPending",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "67108864",
|
||||
Doc: `Maximum number of bytes buffered for a connection Applies to client connections`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerMaxPending.Init(base.mgr)
|
||||
r.ServerInitializeTimeout = ParamItem{
|
||||
Key: "natsmq.server.initializeTimeout",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "4000",
|
||||
Doc: `waiting for initialization of natsmq finished`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerInitializeTimeout.Init(base.mgr)
|
||||
|
||||
r.ServerMonitorDebug = ParamItem{
|
||||
Key: "natsmq.server.monitor.debug",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "false",
|
||||
Doc: `If true enable debug log messages`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerMonitorDebug.Init(base.mgr)
|
||||
r.ServerMonitorLogTime = ParamItem{
|
||||
Key: "natsmq.server.monitor.logTime",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "true",
|
||||
Doc: `If set to false, log without timestamps.`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerMonitorLogTime.Init(base.mgr)
|
||||
r.ServerMonitorLogFile = ParamItem{
|
||||
Key: "natsmq.server.monitor.logFile",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "",
|
||||
Doc: `Log file path relative to..`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerMonitorLogFile.Init(base.mgr)
|
||||
r.ServerMonitorLogSizeLimit = ParamItem{
|
||||
Key: "natsmq.server.monitor.logSizeLimit",
|
||||
Version: "2.3.0",
|
||||
DefaultValue: "0",
|
||||
Doc: `Size in bytes after the log file rolls over to a new one`,
|
||||
Export: true,
|
||||
}
|
||||
r.ServerMonitorLogSizeLimit.Init(base.mgr)
|
||||
}
|
||||
|
||||
// /////////////////////////////////////////////////////////////////////////////
|
||||
// --- minio ---
|
||||
type MinioConfig struct {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user