fix: Fix shared trace dispatcher lifecycle and discovery

client_lifecycle_test.go

@@ -0,0 +1,97 @@
+/*
+Licensed to the Apache Software Foundation (ASF) 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 rocketmq_test
+
+import (
+	"testing"
+
+	rocketmq "github.com/apache/rocketmq-client-go/v2"
+	"github.com/apache/rocketmq-client-go/v2/consumer"
+	"github.com/apache/rocketmq-client-go/v2/internal"
+	"github.com/apache/rocketmq-client-go/v2/producer"
+	"github.com/stretchr/testify/require"
+)
+
+type clientOwner interface{ Shutdown() error }
+
+func TestClientOwnershipBeforeStart(t *testing.T) {
+	constructors := map[string]func(string, string) (clientOwner, error){
+		"producer": func(instance, address string) (clientOwner, error) {
+			return rocketmq.NewProducer(producer.WithInstanceName(instance), producer.WithNameServer([]string{address}))
+		},
+		"push": func(instance, address string) (clientOwner, error) {
+			return rocketmq.NewPushConsumer(consumer.WithInstance(instance), consumer.WithNameServer([]string{address}))
+		},
+		"pull": func(instance, address string) (clientOwner, error) {
+			return rocketmq.NewPullConsumer(consumer.WithInstance(instance), consumer.WithNameServer([]string{address}))
+		},
+	}
+	for name, create := range constructors {
+		t.Run(name, func(t *testing.T) {
+			instance := t.Name()
+			clientID := internal.DefaultClientOptions().ClientIP + "@" + instance
+			first, err := create(instance, "127.0.0.1:9876")
+			require.NoError(t, err)
+			defer first.Shutdown()
+			second, err := create(instance, "127.0.0.1:9876")
+			require.NoError(t, err)
+			defer second.Shutdown()
+			require.NoError(t, first.Shutdown())
+			require.NoError(t, first.Shutdown())
+			_, err = internal.GetNamesrv(clientID)
+			require.NoError(t, err, "another unstarted owner still needs the shared client")
+			_, err = create(instance, "127.0.0.2:9876")
+			require.Error(t, err, "a live owner's NameServer compatibility check must remain in force")
+			require.NoError(t, second.Shutdown())
+			_, err = internal.GetNamesrv(clientID)
+			require.Error(t, err, "the final owner must release the registry entry, including after a rejected acquisition")
+			replacement, err := create(instance, "127.0.0.2:9876")
+			require.NoError(t, err, "the same instance must be reusable with new addresses")
+			require.NoError(t, replacement.Shutdown())
+		})
+	}
+}
+
+func TestClientOwnershipProtectsRunningProducer(t *testing.T) {
+	instance := t.Name()
+	clientID := internal.DefaultClientOptions().ClientIP + "@" + instance
+	create := func(group string) (rocketmq.Producer, error) {
+		return rocketmq.NewProducer(producer.WithInstanceName(instance), producer.WithGroupName(group), producer.WithNameServer([]string{"127.0.0.1:9876"}))
+	}
+	active, err := create("active-group")
+	require.NoError(t, err)
+	defer active.Shutdown()
+	require.NoError(t, active.Start())
+	require.NoError(t, active.Start())
+	for i := 0; i < 2; i++ {
+		duplicate, err := create("active-group")
+		require.NoError(t, err)
+		require.Error(t, duplicate.Start(), "failed owners must not unregister the active producer")
+		require.NoError(t, duplicate.Shutdown())
+	}
+	pending, err := create("pending-group")
+	require.NoError(t, err)
+	defer pending.Shutdown()
+	require.NoError(t, active.Shutdown())
+	_, err = internal.GetNamesrv(clientID)
+	require.NoError(t, err, "the unstarted owner must retain its client after the active owner exits")
+	require.NoError(t, pending.Start())
+	require.NoError(t, pending.Shutdown())
+	_, err = internal.GetNamesrv(clientID)
+	require.Error(t, err)
+}

consumer/consumer.go

@@ -247,13 +247,14 @@ type defaultConsumer struct {
 	fromWhere              ConsumeFromWhere
 	consumerStartTimestamp int64
 
-	cType     ConsumeType
-	client    internal.RMQClient
-	mqChanged func(topic string, mqAll, mqDivided []*primitive.MessageQueue)
-	state     *atomic.Int32
-	pause     *atomic.Bool
-	once      sync.Once
-	option    consumerOptions
+	cType       ConsumeType
+	client      internal.RMQClient
+	mqChanged   func(topic string, mqAll, mqDivided []*primitive.MessageQueue)
+	state       *atomic.Int32
+	pause       *atomic.Bool
+	lifecycleMu sync.Mutex
+	once        sync.Once
+	option      consumerOptions
 	// key: primitive.MessageQueue
 	// value: *processQueue
 	processQueueTable sync.Map

consumer/interceptor.go

@@ -19,8 +19,6 @@ package consumer
 
 import (
 	"context"
-	"fmt"
-	"reflect"
 	"time"
 
 	"github.com/apache/rocketmq-client-go/v2/internal"
@@ -31,24 +29,42 @@ import (
 // WithTrace support rocketmq trace: https://github.com/apache/rocketmq/wiki/RIP-6-Message-Trace.
 func WithTrace(traceCfg *primitive.TraceConfig) Option {
 	return func(options *consumerOptions) {
-		dispatcher := internal.NewTraceDispatcher(traceCfg)
-		options.TraceDispatcher = dispatcher
-		ori := options.Interceptors
-		options.Interceptors = make([]primitive.Interceptor, 0)
-		options.Interceptors = append(options.Interceptors, newTraceInterceptor(dispatcher))
-		options.Interceptors = append(options.Interceptors, ori...)
+		installTraceInterceptor(options, internal.NewTraceDispatcher(traceCfg))
 	}
 }
 
+// WithSharedTrace shares trace connections by logical cluster identity. Discovery
+// addresses may change without recreating the client. See SharedTraceClientConfig
+// for resolver ownership and key requirements.
+func WithSharedTrace(traceCfg *primitive.TraceConfig, shared primitive.SharedTraceClientConfig) Option {
+	return func(options *consumerOptions) {
+		installTraceInterceptor(options, internal.NewSharedTraceDispatcher(traceCfg, shared))
+	}
+}
+
+func installTraceInterceptor(options *consumerOptions, dispatcher internal.TraceDispatcher) {
+	if internal.IsNilTraceDispatcher(dispatcher) {
+		return
+	}
+	if !internal.IsNilTraceDispatcher(options.TraceDispatcher) {
+		options.TraceDispatcher.Close()
+		// Trace is always prepended; WithInterceptor appends user interceptors.
+		// Remove the old wrapper before installing its replacement.
+		options.Interceptors = options.Interceptors[1:]
+	}
+	options.TraceDispatcher = dispatcher
+	options.Interceptors = append([]primitive.Interceptor{newTraceInterceptor(dispatcher)}, options.Interceptors...)
+}
+
 func newTraceInterceptor(dispatcher internal.TraceDispatcher) primitive.Interceptor {
-	if dispatcher != nil && !reflect.ValueOf(dispatcher).IsNil() {
-		dispatcher.Start()
+	if internal.IsNilTraceDispatcher(dispatcher) {
+		return func(ctx context.Context, req, reply interface{}, next primitive.Invoker) error {
+			return next(ctx, req, reply)
+		}
 	}
+	dispatcher.Start()
 
 	return func(ctx context.Context, req, reply interface{}, next primitive.Invoker) error {
-		if dispatcher == nil {
-			return fmt.Errorf("GetOrNewRocketMQClient faild")
-		}
 		consumerCtx, exist := primitive.GetConsumerCtx(ctx)
 		if !exist || len(consumerCtx.Msgs) == 0 {
 			return next(ctx, req, reply)

consumer/pull_consumer.go

@@ -95,6 +95,12 @@ type defaultPullConsumer struct {
 
 func NewPullConsumer(options ...Option) (*defaultPullConsumer, error) {
 	defaultOpts := defaultPullConsumerOptions()
+	constructed := false
+	defer func() {
+		if !constructed && !internal.IsNilTraceDispatcher(defaultOpts.TraceDispatcher) {
+			defaultOpts.TraceDispatcher.Close()
+		}
+	}()
 	for _, apply := range options {
 		apply(&defaultOpts)
 	}
@@ -133,6 +139,7 @@ func NewPullConsumer(options ...Option) (*defaultPullConsumer, error) {
 	dc.mqChanged = c.messageQueueChanged
 	c.submitToConsume = c.consumeMessageConcurrently
 	c.interceptor = primitive.ChainInterceptors(c.option.Interceptors...)
+	constructed = true
 	return c, nil
 }
 
@@ -226,7 +233,17 @@ func (pc *defaultPullConsumer) nextPullOffset(mq *primitive.MessageQueue, origin
 }
 
 func (pc *defaultPullConsumer) Start() error {
+	pc.lifecycleMu.Lock()
+	defer pc.lifecycleMu.Unlock()
+	if pc.state.Load() == int32(internal.StateShutdown) {
+		return internal.ErrServiceState
+	}
 	var err error
+	defer func() {
+		if err != nil && !internal.IsNilTraceDispatcher(pc.option.TraceDispatcher) {
+			pc.option.TraceDispatcher.Close()
+		}
+	}()
 	pc.once.Do(func() {
 		err = pc.validate()
 		if err != nil {
@@ -272,7 +289,7 @@ func (pc *defaultPullConsumer) Start() error {
 	pc.client.UpdateTopicRouteInfo()
 	_, exist := pc.topicSubscribeInfoTable.Load(pc.topic)
 	if !exist {
-		err = pc.Shutdown()
+		err = pc.shutdownLocked()
 		if err != nil {
 			rlog.Error("defaultPullConsumer.Shutdown . route info not found, it may not exist", map[string]interface{}{
 				rlog.LogKeyTopic:         pc.topic,
@@ -570,12 +587,25 @@ func (pc *defaultPullConsumer) CurrentOffset(queue *primitive.MessageQueue) (int
 
 // Shutdown close defaultConsumer, refuse new request.
 func (pc *defaultPullConsumer) Shutdown() error {
+	pc.lifecycleMu.Lock()
+	defer pc.lifecycleMu.Unlock()
+	return pc.shutdownLocked()
+}
+
+func (pc *defaultPullConsumer) shutdownLocked() error {
 	var err error
 	pc.closeOnce.Do(func() {
-		if pc.option.TraceDispatcher != nil {
+		if !internal.IsNilTraceDispatcher(pc.option.TraceDispatcher) {
 			pc.option.TraceDispatcher.Close()
 		}
 		close(pc.done)
+		// Release the construction reference without unregistering another
+		// owner's group when this consumer never started successfully.
+		if pc.state.Load() != int32(internal.StateRunning) {
+			pc.state.Store(int32(internal.StateShutdown))
+			pc.client.Shutdown()
+			return
+		}
 
 		pc.client.UnregisterConsumer(pc.consumerGroup)
 		err = pc.defaultConsumer.shutdown()

consumer/push_consumer.go

@@ -77,6 +77,12 @@ type pushConsumer struct {
 
 func NewPushConsumer(opts ...Option) (*pushConsumer, error) {
 	defaultOpts := defaultPushConsumerOptions()
+	constructed := false
+	defer func() {
+		if !constructed && !internal.IsNilTraceDispatcher(defaultOpts.TraceDispatcher) {
+			defaultOpts.TraceDispatcher.Close()
+		}
+	}()
 	for _, apply := range opts {
 		apply(&defaultOpts)
 	}
@@ -126,11 +132,22 @@ func NewPushConsumer(opts ...Option) (*pushConsumer, error) {
 
 	p.interceptor = primitive.ChainInterceptors(p.option.Interceptors...)
 
+	constructed = true
 	return p, nil
 }
 
 func (pc *pushConsumer) Start() error {
+	pc.lifecycleMu.Lock()
+	defer pc.lifecycleMu.Unlock()
+	if pc.state.Load() == int32(internal.StateShutdown) {
+		return internal.ErrServiceState
+	}
 	var err error
+	defer func() {
+		if err != nil && !internal.IsNilTraceDispatcher(pc.option.TraceDispatcher) {
+			pc.option.TraceDispatcher.Close()
+		}
+	}()
 	pc.once.Do(func() {
 		rlog.Info("the consumer start beginning", map[string]interface{}{
 			rlog.LogKeyConsumerGroup: pc.consumerGroup,
@@ -237,7 +254,7 @@ func (pc *pushConsumer) Start() error {
 	pc.subscribedTopic.Range(func(k, v interface{}) bool {
 		_, exist := pc.topicSubscribeInfoTable.Load(k)
 		if !exist {
-			pc.Shutdown()
+			pc.shutdownLocked()
 			err = fmt.Errorf("the topic=%s route info not found, it may not exist", k)
 			return false
 		}
@@ -277,12 +294,25 @@ func (pc *pushConsumer) GetOffsetDiffMap() map[string]int64 {
 }
 
 func (pc *pushConsumer) Shutdown() error {
+	pc.lifecycleMu.Lock()
+	defer pc.lifecycleMu.Unlock()
+	return pc.shutdownLocked()
+}
+
+func (pc *pushConsumer) shutdownLocked() error {
 	var err error
 	pc.closeOnce.Do(func() {
-		if pc.option.TraceDispatcher != nil {
+		if !internal.IsNilTraceDispatcher(pc.option.TraceDispatcher) {
 			pc.option.TraceDispatcher.Close()
 		}
 		close(pc.done)
+		// Release the construction reference without unregistering another
+		// owner's group when this consumer never started successfully.
+		if pc.state.Load() != int32(internal.StateRunning) {
+			pc.state.Store(int32(internal.StateShutdown))
+			pc.client.Shutdown()
+			return
+		}
 		if pc.consumeOrderly && pc.model == Clustering {
 			pc.unlockAll(false)
 		}