add option WithUnitMode WithUnitName

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