From 3a927989d36d48b0b35a4caa4553dd4e4c248a65 Mon Sep 17 00:00:00 2001 From: WeizhongTu Date: Tue, 31 Jan 2023 16:51:47 +0800 Subject: [PATCH] add option WithUnitMode WithUnitName --- producer/option.go | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/producer/option.go b/producer/option.go index 9fd8374d..082000fe 100644 --- a/producer/option.go +++ b/producer/option.go @@ -18,6 +18,7 @@ limitations under the License. package producer import ( + "strings" "time" "github.com/apache/rocketmq-client-go/v2/internal" @@ -144,7 +145,28 @@ func WithNameServer(nameServers primitive.NamesrvAddr) Option { // WithNameServerDomain set NameServer domain func WithNameServerDomain(nameServerUrl string) Option { return func(opts *producerOptions) { - opts.Resolver = primitive.NewHttpResolver("DEFAULT", nameServerUrl) + h := primitive.NewHttpResolver("DEFAULT", nameServerUrl) + if opts.UnitName != "" { + h.SetUnitName(opts.UnitName) + } + opts.Resolver = h + } +} + +// WithUnitMode set the unit mode +func WithUnitMode(unitMode bool) Option { + return func(opts *producerOptions) { + opts.UnitMode = unitMode + } +} + +// WithUnitName set the name of specified unit +func WithUnitName(unitName string) Option { + return func(opts *producerOptions) { + opts.UnitName = strings.TrimSpace(unitName) + if ns, ok := opts.Resolver.(*primitive.HttpResolver); ok { + ns.SetUnitName(opts.UnitName) + } } }