This is an automated email from the ASF dual-hosted git repository.
ShannonDing pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/rocketmq-client-go.git
The following commit(s) were added to refs/heads/master by this push:
new 461524a fix: Fix resource leaks and shutdown handling in shared
clients (#1244)
461524a is described below
commit 461524ac40fd5e3b3f47735c9b73721374f43504
Author: wenxuwan <[email protected]>
AuthorDate: Sat Oct 10 15:03:58 2026 +0800
fix: Fix resource leaks and shutdown handling in shared clients (#1244)
* seperate interface and implement
* fix panic when close tracedispatcher
* Restore rlog/log.go
* Delete default.go
* fix: release client resources and unblock workers on shutdown
* fix: isolate NameServer transports when rebuilding clients
* fix: copy NameServer snapshots outside the client registry lock
* fix: avoid retaining retired clients in transaction handlers
---------
Co-authored-by: 文徐 <[email protected]>
---
consumer/consumer.go | 41 ++++-
consumer/pull_consumer.go | 14 ++
consumer/push_consumer.go | 58 ++++--
consumer/push_consumer_test.go | 15 +-
consumer/resource_lifecycle_test.go | 59 ++++++
consumer/shutdown_wait_test.go | 160 ++++++++++++++++
consumer/statistics.go | 13 +-
internal/client.go | 87 +++++++--
internal/client_lifecycle_test.go | 282 ++++++++++++++++++++++++++++-
internal/namesrv.go | 30 ++-
internal/remote/future.go | 22 ++-
internal/remote/remote_client.go | 72 ++++++--
internal/remote/remote_client_test.go | 5 +-
internal/remote/resource_lifecycle_test.go | 212 ++++++++++++++++++++++
internal/remote/tcp_conn.go | 48 ++++-
internal/resource_lifecycle_test.go | 141 +++++++++++++++
internal/utils/wait.go | 43 +++++
producer/producer.go | 50 +++--
producer/resource_lifecycle_test.go | 86 +++++++++
19 files changed, 1351 insertions(+), 87 deletions(-)
diff --git a/consumer/consumer.go b/consumer/consumer.go
index 34f1cb5..c749d49 100644
--- a/consumer/consumer.go
+++ b/consumer/consumer.go
@@ -269,6 +269,8 @@ type defaultConsumer struct {
storage OffsetStore
// chan for push consumer
prCh chan PullRequest
+ // Shares the shutdown signal owned by the push or pull consumer.
+ done <-chan struct{}
pullFromWhichNodeTable sync.Map
@@ -702,6 +704,11 @@ func (dc *defaultConsumer)
buildProcessQueueTableByBrokerName() map[string][]*pr
func (dc *defaultConsumer) updateProcessQueueTable(topic string, mqs
[]*primitive.MessageQueue) bool {
var changed bool
+ select {
+ case <-dc.done:
+ return false
+ default:
+ }
mqSet := make(map[primitive.MessageQueue]bool)
for idx := range mqs {
mqSet[*mqs[idx]] = true
@@ -736,6 +743,11 @@ func (dc *defaultConsumer) updateProcessQueueTable(topic
string, mqs []*primitiv
})
for item := range mqSet {
+ select {
+ case <-dc.done:
+ return changed
+ default:
+ }
// BUG: the mq will send to channel, if not copy once, the next
iter will modify the mq in the channel.
mq := item
_, exist := dc.processQueueTable.Load(mq)
@@ -772,7 +784,13 @@ func (dc *defaultConsumer) updateProcessQueueTable(topic
string, mqs []*primitiv
pq: pq,
nextOffset: nextOffset,
}
- dc.prCh <- pr
+ if !dc.dispatchPullRequest(pr) {
+ // Shutdown may have finished scanning
the table while offset lookup
+ // was still in flight. Dispose the
newly published queue ourselves.
+ pq.WithDropped(true)
+ dc.processQueueTable.Delete(mq)
+ return changed
+ }
changed = true
}
} else {
@@ -786,6 +804,27 @@ func (dc *defaultConsumer) updateProcessQueueTable(topic
string, mqs []*primitiv
return changed
}
+// dispatchPullRequest does not retain a rebalance worker after its receiver
stops.
+func (dc *defaultConsumer) dispatchPullRequest(pr PullRequest) bool {
+ select {
+ case <-dc.done:
+ return false
+ default:
+ }
+ select {
+ case <-dc.done:
+ return false
+ case dc.prCh <- pr:
+ }
+ // A buffered send can win a select at the same time as shutdown.
+ select {
+ case <-dc.done:
+ return false
+ default:
+ return true
+ }
+}
+
func (dc *defaultConsumer) removeUnnecessaryMessageQueue(mq
*primitive.MessageQueue, pq *processQueue) bool {
dc.storage.persist([]*primitive.MessageQueue{mq})
dc.storage.remove(mq)
diff --git a/consumer/pull_consumer.go b/consumer/pull_consumer.go
index 8c38305..76cdc9e 100644
--- a/consumer/pull_consumer.go
+++ b/consumer/pull_consumer.go
@@ -136,6 +136,7 @@ func NewPullConsumer(options ...Option)
(*defaultPullConsumer, error) {
consumeRequestCache: make(chan *ConsumeRequest, 4),
GroupName: dc.option.GroupName,
}
+ dc.done = c.done
dc.mqChanged = c.messageQueueChanged
c.submitToConsume = c.consumeMessageConcurrently
c.interceptor = primitive.ChainInterceptors(c.option.Interceptors...)
@@ -305,13 +306,26 @@ func (pc *defaultPullConsumer) Start() error {
return err
}
+// Poll returns ErrService from the errors package when the consumer is shut
down.
func (pc *defaultPullConsumer) Poll(ctx context.Context, timeout
time.Duration) (*ConsumeRequest, error) {
+ select {
+ case <-pc.done:
+ return nil, errors2.ErrService
+ default:
+ }
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
select {
case <-ctx.Done():
return nil, ErrNoNewMsg
+ case <-pc.done:
+ return nil, errors2.ErrService
case cr := <-pc.consumeRequestCache:
+ select {
+ case <-pc.done:
+ return nil, errors2.ErrService
+ default:
+ }
if cr.processQueue.IsDroppd() {
rlog.Info("defaultPullConsumer poll the message queue
not be able to consume, because it was dropped", map[string]interface{}{
rlog.LogKeyMessageQueue:
cr.messageQueue.String(),
diff --git a/consumer/push_consumer.go b/consumer/push_consumer.go
index 2d640bc..9ca2f55 100644
--- a/consumer/push_consumer.go
+++ b/consumer/push_consumer.go
@@ -123,6 +123,7 @@ func NewPushConsumer(opts ...Option) (*pushConsumer, error)
{
done: make(chan struct{}, 1),
consumeFunc: utils.NewSet(),
}
+ dc.done = p.done
dc.mqChanged = p.messageQueueChanged
if p.consumeOrderly {
p.submitToConsume = p.consumeMessageOrderly
@@ -203,7 +204,9 @@ func (pc *pushConsumer) Start() error {
if pc.consumeOrderly {
return
}
- time.Sleep(pc.option.ConsumeTimeout)
+ if !utils.WaitFor(pc.done, pc.option.ConsumeTimeout) {
+ return
+ }
pc.cleanExpiredMsg()
ticker := time.NewTicker(pc.option.ConsumeTimeout)
@@ -227,7 +230,9 @@ func (pc *pushConsumer) Start() error {
return
}
- time.Sleep(1000 * time.Millisecond)
+ if !utils.WaitFor(pc.done, 1000*time.Millisecond) {
+ return
+ }
pc.lockAll()
lockTicker :=
time.NewTicker(pc.option.RebalanceLockInterval)
@@ -709,7 +714,9 @@ func (pc *pushConsumer) pullMessage(request *PullRequest) {
}
if sleepTime > 0 {
rlog.Debug(fmt.Sprintf("pull MessageQueue: %d sleep %d
ms for mq: %v", request.mq.QueueId, sleepTime/time.Millisecond, request.mq),
nil)
- time.Sleep(sleepTime)
+ if !utils.WaitFor(pc.done, sleepTime) {
+ return
+ }
}
// reset time
sleepTime = pc.option.PullInterval.Load()
@@ -935,7 +942,9 @@ func (pc *pushConsumer) pullMessage(request *PullRequest) {
})
request.nextOffset = result.NextBeginOffset
pq.WithDropped(true)
- time.Sleep(10 * time.Second)
+ if !utils.WaitFor(pc.done, 10*time.Second) {
+ return
+ }
pc.storage.update(request.mq, request.nextOffset, false)
pc.storage.persist([]*primitive.MessageQueue{request.mq})
pc.processQueueTable.Delete(*request.mq)
@@ -1035,7 +1044,9 @@ func (pc *pushConsumer) ResetOffset(topic string, table
map[primitive.MessageQue
copyPc.Store(&mq, pq)
return true
})
- time.Sleep(10 * time.Second)
+ if !utils.WaitFor(pc.done, 10*time.Second) {
+ return
+ }
for _, mq := range mqs {
if _, ok := table[*mq]; ok {
pc.storage.update(mq, table[*mq], false)
@@ -1144,7 +1155,13 @@ func (pc *pushConsumer) consumeMessageConcurrently(pq
*processQueue, mq *primiti
}
ch, _ := pc.crCh.Load(mq.Topic)
if channel, ok := ch.(chan struct{}); ok {
- channel <- struct{}{}
+ select {
+ case <-pc.done:
+ return
+ case <-pq.closeChan:
+ return
+ case channel <- struct{}{}:
+ }
}
go primitive.WithRecover(func() {
defer func() {
@@ -1164,6 +1181,11 @@ func (pc *pushConsumer) consumeMessageConcurrently(pq
*processQueue, mq *primiti
}
}()
RETRY:
+ select {
+ case <-pc.done:
+ return
+ default:
+ }
if pq.IsDroppd() {
rlog.Info("the message queue not be able to
consume, because it was dropped", map[string]interface{}{
rlog.LogKeyMessageQueue: mq.String(),
@@ -1250,7 +1272,9 @@ func (pc *pushConsumer) consumeMessageConcurrently(pq
*processQueue, mq *primiti
}
if len(msgBackFailed) > 0 {
subMsgs = msgBackFailed
- time.Sleep(5 * time.Second)
+ if !utils.WaitFor(pc.done,
5*time.Second) {
+ return
+ }
goto RETRY
}
} else {
@@ -1303,7 +1327,7 @@ func (pc *pushConsumer) consumeMessageOrderly(pq
*processQueue, mq *primitive.Me
}
interval := time.Now().Sub(beginTime)
if interval > pc.option.MaxTimeConsumeContinuously {
- time.Sleep(10 * time.Millisecond)
+ utils.WaitFor(pc.done, 10*time.Millisecond)
return
}
batchSize := pc.option.ConsumeMessageBatchMaxSize
@@ -1377,7 +1401,9 @@ func (pc *pushConsumer) consumeMessageOrderly(pq
*processQueue, mq *primitive.Me
case SuspendCurrentQueueAMoment:
if pc.checkReconsumeTimes(msgs) {
pq.makeMessageToCosumeAgain(msgs...)
-
time.Sleep(time.Duration(orderlyCtx.SuspendCurrentQueueTimeMillis) *
time.Millisecond)
+ if !utils.WaitFor(pc.done,
time.Duration(orderlyCtx.SuspendCurrentQueueTimeMillis)*time.Millisecond) {
+ return
+ }
continueConsume = false
} else {
commitOffset = pq.commit()
@@ -1391,11 +1417,15 @@ func (pc *pushConsumer) consumeMessageOrderly(pq
*processQueue, mq *primitive.Me
commitOffset = pq.commit()
case Rollback:
// pq.rollback
-
time.Sleep(time.Duration(orderlyCtx.SuspendCurrentQueueTimeMillis) *
time.Millisecond)
+ if !utils.WaitFor(pc.done,
time.Duration(orderlyCtx.SuspendCurrentQueueTimeMillis)*time.Millisecond) {
+ return
+ }
continueConsume = false
case SuspendCurrentQueueAMoment:
if pc.checkReconsumeTimes(msgs) {
-
time.Sleep(time.Duration(orderlyCtx.SuspendCurrentQueueTimeMillis) *
time.Millisecond)
+ if !utils.WaitFor(pc.done,
time.Duration(orderlyCtx.SuspendCurrentQueueTimeMillis)*time.Millisecond) {
+ return
+ }
continueConsume = false
}
default:
@@ -1453,7 +1483,9 @@ func (pc *pushConsumer) getMaxReconsumeTimes() int32 {
}
func (pc *pushConsumer) tryLockLaterAndReconsume(mq *primitive.MessageQueue,
delay int64) {
- time.Sleep(time.Duration(delay) * time.Millisecond)
+ if !utils.WaitFor(pc.done, time.Duration(delay)*time.Millisecond) {
+ return
+ }
if pc.lock(mq) == true {
pc.submitConsumeRequestLater(10)
} else {
@@ -1470,7 +1502,7 @@ func (pc *pushConsumer)
submitConsumeRequestLater(suspendTimeMillis int64) {
} else if suspendTimeMillis > 30000 {
suspendTimeMillis = 30000
}
- time.Sleep(time.Duration(suspendTimeMillis) * time.Millisecond)
+ utils.WaitFor(pc.done,
time.Duration(suspendTimeMillis)*time.Millisecond)
}
func (pc *pushConsumer) cleanExpiredMsg() {
diff --git a/consumer/push_consumer_test.go b/consumer/push_consumer_test.go
index 70779e5..89b823f 100644
--- a/consumer/push_consumer_test.go
+++ b/consumer/push_consumer_test.go
@@ -20,6 +20,7 @@ package consumer
import (
"context"
"testing"
+ "time"
"github.com/apache/rocketmq-client-go/v2/rlog"
@@ -41,11 +42,15 @@ func TestStart(t *testing.T) {
WithConsumerModel(BroadCasting),
)
+ defer c.client.Shutdown() // Release the real construction
reference replaced below.
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := internal.NewMockRMQClient(ctrl)
c.client = client
+ defer c.Shutdown()
+ client.EXPECT().Shutdown().Return()
+ client.EXPECT().UnregisterConsumer(gomock.Any()).Return()
err := c.Subscribe("TopicTest", MessageSelector{}, func(ctx
context.Context,
msgs ...*primitive.MessageExt) (ConsumeResult, error) {
@@ -80,19 +85,23 @@ func TestStart(t *testing.T) {
client.EXPECT().UpdateTopicRouteInfo().AnyTimes().Return()
Convey("test topic route info not found", func() {
- client.EXPECT().Shutdown().Return()
-
client.EXPECT().UnregisterConsumer(gomock.Any()).Return()
err = c.Start()
So(err.Error(), ShouldContainSubstring, "route info not
found")
})
Convey("test topic route info found", func() {
- client.EXPECT().RebalanceImmediately().Return()
+ rebalanced := make(chan struct{})
+ client.EXPECT().RebalanceImmediately().Do(func() {
close(rebalanced) })
client.EXPECT().CheckClientInBroker().Return()
client.EXPECT().SendHeartbeatToAllBrokerWithLock().Return()
mockB4Start(c)
err = c.Start()
So(err, ShouldBeNil)
+ select {
+ case <-rebalanced:
+ case <-time.After(time.Second):
+ t.Fatal("startup rebalance did not complete")
+ }
})
})
}
diff --git a/consumer/resource_lifecycle_test.go
b/consumer/resource_lifecycle_test.go
new file mode 100644
index 0000000..98e79f2
--- /dev/null
+++ b/consumer/resource_lifecycle_test.go
@@ -0,0 +1,59 @@
+/*
+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 consumer
+
+import (
+ "bytes"
+ "runtime/pprof"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/apache/rocketmq-client-go/v2/primitive"
+ "github.com/stretchr/testify/require"
+)
+
+func consumerBackgroundWorkers() int {
+ var b bytes.Buffer
+ pprof.Lookup("goroutine").WriteTo(&b, 2)
+ return strings.Count(b.String(), "(*pushConsumer).Start.func") +
strings.Count(b.String(), "(*statsItemSet).init.func")
+}
+
+func TestConsumerRepeatedLifecycleStopsBackgroundWorkers(t *testing.T) {
+ baseline := consumerBackgroundWorkers()
+ for i := 0; i < 10; i++ {
+ pc, err := NewPushConsumer(WithInstance(t.Name()),
WithGroupName(t.Name()), WithConsumerModel(BroadCasting),
WithConsumeTimeout(15*time.Minute),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer pc.Shutdown()
+ require.NoError(t, pc.Start())
+ require.NoError(t, pc.Shutdown())
+ require.Eventually(t, func() bool { return
consumerBackgroundWorkers() <= baseline }, time.Second, time.Millisecond,
"shutdown left cleanup/statistics workers waiting")
+ }
+}
+
+func TestShutdownInterruptsConsumerRetryDelay(t *testing.T) {
+ pc := &pushConsumer{done: make(chan struct{})}
+ exited := make(chan struct{})
+ go func() { pc.tryLockLaterAndReconsume(nil, 15*60*1000); close(exited)
}()
+ close(pc.done)
+ select {
+ case <-exited:
+ case <-time.After(time.Second):
+ t.Fatal("retry delay ignored shutdown")
+ }
+}
diff --git a/consumer/shutdown_wait_test.go b/consumer/shutdown_wait_test.go
new file mode 100644
index 0000000..06ca95f
--- /dev/null
+++ b/consumer/shutdown_wait_test.go
@@ -0,0 +1,160 @@
+/*
+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 consumer
+
+import (
+ "context"
+ errors2 "github.com/apache/rocketmq-client-go/v2/errors"
+ "github.com/apache/rocketmq-client-go/v2/primitive"
+ "github.com/golang/mock/gomock"
+ "github.com/stretchr/testify/require"
+ "testing"
+ "time"
+)
+
+func TestShutdownReleasesPendingRebalance(t *testing.T) {
+ pc, err := NewPushConsumer(WithInstance(t.Name()),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer pc.Shutdown()
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+ storage := NewMockOffsetStore(ctrl)
+ pc.storage = storage
+ storage.EXPECT().remove(gomock.Any()).AnyTimes()
+ // Model an in-flight rebalance whose offset lookup completes after
shutdown.
+ entered, resume := make(chan struct{}), make(chan struct{})
+ first := true
+ storage.EXPECT().readWithException(gomock.Any(),
gomock.Any()).DoAndReturn(func(_ *primitive.MessageQueue, _ readType) (int64,
error) {
+ if first {
+ first = false
+ close(entered)
+ <-resume
+ }
+ return 0, nil
+ }).AnyTimes()
+ queues := make([]*primitive.MessageQueue, cap(pc.prCh)+1)
+ for i := range queues {
+ queues[i] = &primitive.MessageQueue{Topic: "test", BrokerName:
"broker", QueueId: i}
+ }
+ exited := make(chan struct{})
+ go func() { defer close(exited); pc.updateProcessQueueTable("test",
queues) }()
+ <-entered
+ require.NoError(t, pc.Shutdown())
+ close(resume)
+ select {
+ case <-exited:
+ case <-time.After(time.Second):
+ t.Error("rebalance remains blocked sending to prCh after
shutdown")
+ }
+ // Drain the orphaned channel so the reproduction itself does not leak.
+ for {
+ select {
+ case <-exited:
+ pc.processQueueTable.Range(func(_, _ interface{}) bool
{ t.Error("shutdown retained a newly published queue"); return true })
+ return
+ case <-pc.prCh:
+ }
+ }
+}
+
+func TestShutdownReleasesPushConcurrencyWait(t *testing.T) {
+ pc, err := NewPushConsumer(WithInstance(t.Name()),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer pc.Shutdown()
+ pc.option.ConsumeMessageBatchMaxSize = 1
+ pq := newProcessQueue(false)
+ mq := &primitive.MessageQueue{Topic: "test"}
+ pq.msgCh <- []*primitive.MessageExt{{Message: primitive.Message{Topic:
"test"}}}
+ slots := make(chan struct{}, 1)
+ slots <- struct{}{}
+ pc.crCh.Store("test", slots)
+ exited := make(chan struct{})
+ go func() { defer close(exited); pc.consumeMessageConcurrently(pq, mq)
}()
+ time.Sleep(30 * time.Millisecond)
+ require.NoError(t, pc.Shutdown())
+ pq.WithDropped(true)
+ select {
+ case <-exited:
+ case <-time.After(time.Second):
+ t.Error("consumer dispatcher remains blocked on concurrency
slot after shutdown")
+ }
+ <-slots
+ <-exited
+}
+
+func TestShutdownWakesPoll(t *testing.T) {
+ pc, err := NewPullConsumer(WithInstance(t.Name()),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer pc.Shutdown()
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ exited := make(chan struct{})
+ go func() {
+ defer close(exited)
+ _, err := pc.Poll(ctx, time.Hour)
+ require.Equal(t, errors2.ErrService, err)
+ }()
+ require.NoError(t, pc.Shutdown())
+ select {
+ case <-exited:
+ case <-time.After(time.Second):
+ t.Error("Poll retains caller until its timeout even after
Shutdown")
+ }
+ cancel()
+ <-exited
+}
+
+func TestPollBeforeAndAfterShutdown(t *testing.T) {
+ pc, err := NewPullConsumer(WithInstance(t.Name()),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer pc.Shutdown()
+ _, err = pc.Poll(context.Background(), time.Millisecond)
+ require.Equal(t, ErrNoNewMsg, err)
+ cr := &ConsumeRequest{processQueue: newProcessQueue(false), msgList:
[]*primitive.MessageExt{{}}}
+ pc.consumeRequestCache <- cr
+ result, err := pc.Poll(context.Background(), time.Second)
+ require.NoError(t, err)
+ require.Same(t, cr, result)
+ pc.consumeRequestCache <- cr
+ require.NoError(t, pc.Shutdown())
+ _, err = pc.Poll(context.Background(), time.Hour)
+ require.Equal(t, errors2.ErrService, err, "closed Poll must not return
a cached request")
+}
+
+func TestRebalanceDispatchesWhileOpen(t *testing.T) {
+ pc, err := NewPushConsumer(WithInstance(t.Name()),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer pc.Shutdown()
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+ storage := NewMockOffsetStore(ctrl)
+ pc.storage = storage
+ mq := &primitive.MessageQueue{Topic: "test", BrokerName: "broker"}
+ storage.EXPECT().remove(mq)
+ storage.EXPECT().readWithException(mq, gomock.Any()).Return(int64(42),
nil)
+ require.True(t, pc.updateProcessQueueTable("test",
[]*primitive.MessageQueue{mq}))
+ select {
+ case pr := <-pc.prCh:
+ require.Equal(t, int64(42), pr.nextOffset)
+ require.Equal(t, *mq, *pr.mq)
+ require.False(t, pr.pq.IsDroppd())
+ pr.pq.WithDropped(true)
+ case <-time.After(time.Second):
+ t.Fatal("rebalance did not dispatch its queue")
+ }
+}
diff --git a/consumer/statistics.go b/consumer/statistics.go
index 19e0f3b..1ca2ce9 100644
--- a/consumer/statistics.go
+++ b/consumer/statistics.go
@@ -24,6 +24,7 @@ import (
"sync/atomic"
"time"
+ "github.com/apache/rocketmq-client-go/v2/internal/utils"
"github.com/apache/rocketmq-client-go/v2/primitive"
"github.com/apache/rocketmq-client-go/v2/rlog"
)
@@ -211,7 +212,9 @@ func (sis *statsItemSet) init() {
})
go primitive.WithRecover(func() {
- time.Sleep(nextMinutesTime().Sub(time.Now()))
+ if !utils.WaitFor(sis.closed,
nextMinutesTime().Sub(time.Now())) {
+ return
+ }
ticker := time.NewTicker(time.Minute)
defer ticker.Stop()
for {
@@ -225,7 +228,9 @@ func (sis *statsItemSet) init() {
})
go primitive.WithRecover(func() {
- time.Sleep(nextHourTime().Sub(time.Now()))
+ if !utils.WaitFor(sis.closed, nextHourTime().Sub(time.Now())) {
+ return
+ }
ticker := time.NewTicker(time.Hour)
defer ticker.Stop()
for {
@@ -239,7 +244,9 @@ func (sis *statsItemSet) init() {
})
go primitive.WithRecover(func() {
- time.Sleep(nextMonthTime().Sub(time.Now()))
+ if !utils.WaitFor(sis.closed, nextMonthTime().Sub(time.Now())) {
+ return
+ }
ticker := time.NewTicker(24 * time.Hour)
defer ticker.Stop()
for {
diff --git a/internal/client.go b/internal/client.go
index c302cf7..3b17632 100644
--- a/internal/client.go
+++ b/internal/client.go
@@ -177,7 +177,6 @@ type rmqClient struct {
remoteClient remote.RemotingClient
hbMutex sync.Mutex
- close bool
rbMutex sync.Mutex
done chan struct{}
shutdownOnce sync.Once
@@ -197,20 +196,53 @@ var clientMapMu sync.Mutex
// GetOrNewRocketMQClient acquires a reference that must be released by
Shutdown,
// whether or not the caller starts the client.
-func GetOrNewRocketMQClient(option ClientOptions, callbackCh chan interface{})
RMQClient {
+// callbackStops optionally cancels callback delivery when its producer stops,
+// even if other owners keep the shared client alive.
+func GetOrNewRocketMQClient(option ClientOptions, callbackCh chan interface{},
callbackStops ...<-chan struct{}) RMQClient {
+ var callbackStop <-chan struct{}
+ if len(callbackStops) > 0 {
+ callbackStop = callbackStops[0]
+ }
client := &rmqClient{
option: option,
remoteClient:
remote.NewRemotingClient(option.RemotingClientConfig),
done: make(chan struct{}),
ready: make(chan struct{}),
}
+ clientID := client.ClientID()
+ var cloned *namesrvs
clientMapMu.Lock()
- actual, loaded := clientMap.LoadOrStore(client.ClientID(), client)
- client = actual.(*rmqClient)
+ actual, loaded := clientMap.Load(clientID)
+ if !loaded && option.Namesrv.(*namesrvs).bundleClient != nil {
+ // Snapshot copying can wait for discovery to release the
NameServer
+ // lock. Never hold the global client registry lock while
waiting.
+ clientMapMu.Unlock()
+ cloned = option.Namesrv.(*namesrvs).cloneForClient()
+ clientMapMu.Lock()
+ // Another caller may have published this client ID in the
meantime.
+ actual, loaded = clientMap.Load(clientID)
+ }
+ if loaded {
+ client = actual.(*rmqClient)
+ } else {
+ srvs := option.Namesrv.(*namesrvs)
+ if cloned != nil {
+ srvs = cloned
+ }
+ // Ownership and publication stay atomic. Never rebind an old
instance's
+ // NameServer, including while its shutdown is still in flight.
+ srvs.bundleClient = client
+ client.option.Namesrv = srvs
+ clientMap.Store(clientID, client)
+ }
client.refs++
clientMapMu.Unlock()
if loaded {
+ if cloned != nil {
+ // A concurrent creator won; dispose the unused private
transport.
+ cloned.nameSrvClient.ShutDown()
+ }
// Do not expose a shared client before its request handlers
are ready.
<-client.ready
// compare namesrv address
@@ -265,7 +297,9 @@ func GetOrNewRocketMQClient(option ClientOptions,
callbackCh chan interface{}) R
rlog.Warning("checkTransactionState, pick
producer group failed", nil)
return nil
}
- if option.GroupName != group {
+ // Capture only the current client. The input options
may reference a
+ // retired NameServer and keep every previous
generation reachable.
+ if client.option.GroupName != group {
rlog.Warning("producer group is not equal", nil)
return nil
}
@@ -274,7 +308,11 @@ func GetOrNewRocketMQClient(option ClientOptions,
callbackCh chan interface{}) R
Msg: msgExt,
Header: *header,
}
- callbackCh <- callback
+ select {
+ case callbackCh <- callback:
+ case <-callbackStop:
+ case <-client.done:
+ }
return nil
})
@@ -419,8 +457,6 @@ func GetOrNewRocketMQClient(option ClientOptions,
callbackCh chan interface{}) R
}
return res
})
- // Only the creator binds the NameServer; reuse must not race
with readers.
- client.GetNameSrv().(*namesrvs).bundleClient = client
close(client.ready)
}
return client
@@ -435,7 +471,9 @@ func (c *rmqClient) Start() {
op := func() {
c.GetNameSrv().UpdateNameServerAddress()
}
- time.Sleep(10 * time.Second)
+ if !utils.WaitFor(c.done, 10*time.Second) {
+ return
+ }
op()
ticker := time.NewTicker(2 * time.Minute)
@@ -459,7 +497,9 @@ func (c *rmqClient) Start() {
op := func() {
c.UpdateTopicRouteInfo()
}
- time.Sleep(10 * time.Millisecond)
+ if !utils.WaitFor(c.done, 10*time.Millisecond) {
+ return
+ }
op()
ticker := time.NewTicker(_PullNameServerInterval)
@@ -483,7 +523,9 @@ func (c *rmqClient) Start() {
c.SendHeartbeatToAllBrokerWithLock()
}
- time.Sleep(time.Second)
+ if !utils.WaitFor(c.done, time.Second) {
+ return
+ }
op()
ticker := time.NewTicker(_HeartbeatBrokerInterval)
@@ -515,7 +557,9 @@ func (c *rmqClient) Start() {
return true
})
}
- time.Sleep(10 * time.Second)
+ if !utils.WaitFor(c.done, 10*time.Second) {
+ return
+ }
op()
ticker := time.NewTicker(_PersistOffsetInterval)
@@ -572,11 +616,22 @@ func (c *rmqClient) Shutdown() {
"clientID": c.ClientID(),
})
close(c.done)
- c.close = true
c.remoteClient.ShutDown()
+ if srvs, ok := c.GetNameSrv().(*namesrvs); ok && srvs != nil {
+ srvs.nameSrvClient.ShutDown()
+ }
})
}
+func (c *rmqClient) isClosed() bool {
+ select {
+ case <-c.done:
+ return true
+ default:
+ return false
+ }
+}
+
func (c *rmqClient) ClientID() string {
id := c.option.ClientIP + "@"
if c.option.InstanceName == "DEFAULT" {
@@ -592,7 +647,7 @@ func (c *rmqClient) ClientID() string {
func (c *rmqClient) InvokeSync(ctx context.Context, addr string, request
*remote.RemotingCommand,
timeoutMillis time.Duration) (*remote.RemotingCommand, error) {
- if c.close {
+ if c.isClosed() {
return nil, ErrServiceState
}
var cancel context.CancelFunc
@@ -603,7 +658,7 @@ func (c *rmqClient) InvokeSync(ctx context.Context, addr
string, request *remote
func (c *rmqClient) InvokeAsync(ctx context.Context, addr string, request
*remote.RemotingCommand,
f func(*remote.RemotingCommand, error)) error {
- if c.close {
+ if c.isClosed() {
return ErrServiceState
}
return c.remoteClient.InvokeAsync(ctx, addr, request, func(future
*remote.ResponseFuture) {
@@ -614,7 +669,7 @@ func (c *rmqClient) InvokeAsync(ctx context.Context, addr
string, request *remot
func (c *rmqClient) InvokeOneWay(ctx context.Context, addr string, request
*remote.RemotingCommand,
timeoutMillis time.Duration) error {
- if c.close {
+ if c.isClosed() {
return ErrServiceState
}
return c.remoteClient.InvokeOneWay(ctx, addr, request)
diff --git a/internal/client_lifecycle_test.go
b/internal/client_lifecycle_test.go
index a2d7502..bbd5d7f 100644
--- a/internal/client_lifecycle_test.go
+++ b/internal/client_lifecycle_test.go
@@ -18,8 +18,18 @@ limitations under the License.
package internal
import (
+ "context"
+ "encoding/binary"
+ "encoding/json"
+ "io"
+ "net"
+ "runtime"
"sync"
"testing"
+ "time"
+
+ "github.com/apache/rocketmq-client-go/v2/internal/remote"
+ "github.com/golang/mock/gomock"
"github.com/apache/rocketmq-client-go/v2/primitive"
"github.com/stretchr/testify/require"
@@ -67,8 +77,11 @@ func TestClientOwnershipConcurrentAcquireAndRelease(t
*testing.T) {
}
func TestClientOwnershipOldCleanupKeepsReplacement(t *testing.T) {
- options := ownershipTestOptions(t)
+ options, stopServer := namesrvRebuildTestOptions(t)
+ defer stopServer()
old := GetOrNewRocketMQClient(options, nil).(*rmqClient)
+ defer old.Shutdown()
+ assertNamesrvRebuildRoute(t, old)
old.Shutdown()
replacement := GetOrNewRocketMQClient(options, nil).(*rmqClient)
defer replacement.Shutdown()
@@ -77,6 +90,15 @@ func TestClientOwnershipOldCleanupKeepsReplacement(t
*testing.T) {
stored, exists := clientMap.Load(replacement.ClientID())
require.True(t, exists)
require.Same(t, replacement, stored)
+ assertNamesrvRebuildRoute(t, replacement)
+ require.NotSame(t, old.GetNameSrv(), replacement.GetNameSrv())
+ require.Same(t, old.GetNameSrv().(*namesrvs).remotingConfig,
replacement.GetNameSrv().(*namesrvs).remotingConfig)
+ require.Same(t, old, old.GetNameSrv().(*namesrvs).bundleClient)
+ // Reusing the options retained inside a replacement must work as well.
+ replacement.Shutdown()
+ third := GetOrNewRocketMQClient(replacement.option, nil).(*rmqClient)
+ defer third.Shutdown()
+ assertNamesrvRebuildRoute(t, third)
}
func TestClientOwnershipKeepsResolverSnapshot(t *testing.T) {
@@ -94,3 +116,261 @@ func TestClientOwnershipKeepsResolverSnapshot(t
*testing.T) {
require.Equal(t, []string{"127.0.0.2:9876", "127.0.0.1:9876"},
addresses,
"compatibility checks must not sort a resolver-owned snapshot
in place")
}
+
+func TestClientOwnershipReplacementDuringOldShutdown(t *testing.T) {
+ options, stopServer := namesrvRebuildTestOptions(t)
+ defer stopServer()
+ old := GetOrNewRocketMQClient(options, nil).(*rmqClient)
+ defer old.Shutdown()
+ assertNamesrvRebuildRoute(t, old)
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+ transport := remote.NewMockRemotingClient(ctrl)
+ entered, resume, stopped := make(chan struct{}), make(chan struct{}),
make(chan struct{})
+ transport.EXPECT().ShutDown().Do(func() { close(entered); <-resume })
+ old.remoteClient = transport
+ go func() { old.Shutdown(); close(stopped) }()
+ var releaseOnce sync.Once
+ release := func() { releaseOnce.Do(func() { close(resume) }); <-stopped
}
+ defer release()
+ <-entered
+ replacement := GetOrNewRocketMQClient(options, nil).(*rmqClient)
+ defer replacement.Shutdown()
+ assertNamesrvRebuildRoute(t, replacement)
+ release()
+ assertNamesrvRebuildRoute(t, replacement)
+ ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+ defer cancel()
+ _, err :=
old.GetNameSrv().(*namesrvs).queryTopicRouteInfoWithContext(ctx, "test")
+ require.Error(t, err, "the old transport must remain closed")
+}
+
+func assertNamesrvRebuildRoute(t *testing.T, client *rmqClient) {
+ t.Helper()
+ ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+ defer cancel()
+ route, _, err :=
client.GetNameSrv().(*namesrvs).updateTopicRouteInfoWithContext(ctx, "test",
"", 0)
+ require.NoError(t, err)
+ require.Equal(t, "broker:1", route.OrderTopicConf)
+}
+
+// Reply to real framed route requests, including checking that ACL survives a
rebuild.
+func namesrvRebuildTestOptions(t *testing.T) (ClientOptions, func()) {
+ t.Helper()
+ listener, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ var wg sync.WaitGroup
+ var conns sync.Map
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ for {
+ conn, err := listener.Accept()
+ if err != nil {
+ return
+ }
+ conns.Store(conn, true)
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ defer conn.Close()
+ defer conns.Delete(conn)
+ _ = conn.SetDeadline(time.Now().Add(5 *
time.Second))
+ for {
+ var size uint32
+ if err := binary.Read(conn,
binary.BigEndian, &size); err != nil {
+ return
+ }
+ if size < 4 || size > 1024*1024 {
+ t.Error("invalid request size")
+ return
+ }
+ frame := make([]byte, size)
+ if _, err := io.ReadFull(conn, frame);
err != nil {
+ return
+ }
+ headerSize :=
binary.BigEndian.Uint32(frame[:4])
+ if headerSize > size-4 {
+ t.Error("invalid JSON header
size")
+ return
+ }
+ var request remote.RemotingCommand
+ if err :=
json.Unmarshal(frame[4:4+headerSize], &request); err != nil {
+ t.Error(err)
+ return
+ }
+ if request.Code !=
ReqGetRouteInfoByTopic || request.ExtFields["topic"] != "test" ||
+ request.ExtFields["AccessKey"]
!= "test-key" || request.ExtFields["SecurityToken"] != "test-token" ||
request.ExtFields["Signature"] == "" {
+ t.Error("route request lost its
topic or ACL credentials")
+ return
+ }
+ response :=
remote.NewRemotingCommand(ResSuccess, nil,
[]byte(`{"orderTopicConf":"broker:1","queueDatas":[],"brokerDatas":[]}`))
+ response.Opaque = request.Opaque
+ response.Flag = remote.ResponseType
+ if err := response.WriteTo(conn); err
!= nil {
+ return
+ }
+ }
+ }()
+ }
+ }()
+ stop := func() {
+ listener.Close()
+ conns.Range(func(key, _ interface{}) bool {
key.(net.Conn).Close(); return true })
+ wg.Wait()
+ }
+ options := DefaultClientOptions()
+ options.InstanceName = t.Name()
+ config := remote.DefaultRemotingClientConfig
+ config.ConnectionTimeout = time.Second
+ ns, err :=
NewNamesrv(primitive.NewPassthroughResolver([]string{listener.Addr().String()}),
&config)
+ if err != nil {
+ stop()
+ t.Fatal(err)
+ }
+ ns.SetCredentials(primitive.Credentials{AccessKey: "test-key",
SecretKey: "test-secret", SecurityToken: "test-token"})
+ options.Namesrv = ns
+ return options, stop
+}
+
+// Observe entry to Lock without relying on sleeps to schedule the clone
attempt.
+type snapshotAttemptLocker struct {
+ sync.Locker
+ attempts chan struct{}
+}
+
+func (l *snapshotAttemptLocker) Lock() {
+ select {
+ case l.attempts <- struct{}{}:
+ default:
+ }
+ l.Locker.Lock()
+}
+
+type blockedDiscoveryResolver struct {
+ primitive.NsResolver
+ entered chan struct{}
+ resume chan struct{}
+ once sync.Once
+}
+
+func (r *blockedDiscoveryResolver) Resolve() []string {
+ r.once.Do(func() { close(r.entered) })
+ <-r.resume
+ return r.NsResolver.Resolve()
+}
+
+func TestClientOwnershipSlowSnapshotDoesNotBlockOtherClients(t *testing.T) {
+ options := ownershipTestOptions(t)
+ old := GetOrNewRocketMQClient(options, nil).(*rmqClient)
+ old.Shutdown()
+ ns := options.Namesrv.(*namesrvs)
+ resolver := &blockedDiscoveryResolver{
+ NsResolver: ns.resolver, entered: make(chan struct{}), resume:
make(chan struct{}),
+ }
+ ns.resolver = resolver
+ lock := &snapshotAttemptLocker{Locker: ns.lock, attempts: make(chan
struct{}, 4)}
+ ns.lock = lock
+
+ closingOptions := ownershipTestOptions(t)
+ closingOptions.InstanceName += "-closing"
+ closing := GetOrNewRocketMQClient(closingOptions, nil)
+ defer closing.Shutdown()
+ creatingOptions := ownershipTestOptions(t)
+ creatingOptions.InstanceName += "-creating"
+ // A competing creator uses fresh discovery state for the same client
ID.
+ winnerOptions := ownershipTestOptions(t)
+ var workers sync.WaitGroup
+ var resumeOnce sync.Once
+ resume := func() { resumeOnce.Do(func() { close(resolver.resume) }) }
+ defer func() { resume(); workers.Wait() }()
+ workers.Add(1)
+ go func() { defer workers.Done(); ns.UpdateNameServerAddress() }()
+ select {
+ case <-resolver.entered:
+ case <-time.After(time.Second):
+ t.Fatal("discovery did not start")
+ }
+ <-lock.attempts // The refresh has acquired the NameServer lock.
+
+ rebuilt := make(chan *rmqClient, 1)
+ workers.Add(1)
+ go func() {
+ defer workers.Done()
+ client := GetOrNewRocketMQClient(options, nil).(*rmqClient)
+ defer client.Shutdown()
+ rebuilt <- client
+ }()
+ select {
+ case <-lock.attempts: // The rebuilding client is now trying to copy
the snapshot.
+ case <-time.After(time.Second):
+ t.Fatal("rebuild did not attempt to acquire the NameServer
lock")
+ }
+
+ closed, created := make(chan struct{}), make(chan struct{})
+ winnerReady, releaseWinner := make(chan *rmqClient, 1), make(chan
struct{})
+ defer close(releaseWinner)
+ workers.Add(3)
+ go func() { defer workers.Done(); closing.Shutdown(); close(closed) }()
+ go func() {
+ defer workers.Done()
+ client := GetOrNewRocketMQClient(creatingOptions, nil)
+ client.Shutdown()
+ close(created)
+ }()
+ go func() {
+ defer workers.Done()
+ client := GetOrNewRocketMQClient(winnerOptions,
nil).(*rmqClient)
+ defer client.Shutdown()
+ winnerReady <- client
+ <-releaseWinner
+ }()
+ for name, done := range map[string]<-chan struct{}{"shutdown": closed,
"creation": created} {
+ select {
+ case <-done:
+ case <-time.After(time.Second):
+ t.Errorf("unrelated client %s blocked on slow
discovery", name)
+ }
+ }
+ var winner *rmqClient
+ select {
+ case winner = <-winnerReady:
+ case <-time.After(time.Second):
+ t.Error("same-ID creator with fresh discovery state was
blocked")
+ }
+ resume()
+ replacement := <-rebuilt
+ if winner != nil {
+ require.Same(t, winner, replacement, "recheck the client table
after copying the snapshot")
+ require.Same(t, winner,
winner.GetNameSrv().(*namesrvs).bundleClient)
+ }
+ require.NotSame(t, ns, replacement.GetNameSrv())
+ require.Same(t, old, ns.bundleClient, "the old NameServer binding must
remain unchanged")
+}
+
+func TestClientOwnershipRebuildDoesNotRetainRetiredClients(t *testing.T) {
+ client, collected := rebuiltClientWithRetiredCaches(t)
+ defer client.Shutdown()
+ require.Eventually(t, func() bool {
+ runtime.GC()
+ return len(collected) == 3
+ }, 3*time.Second, 10*time.Millisecond, "the live replacement must not
retain retired clients through its request handlers")
+ runtime.KeepAlive(client)
+}
+
+func rebuiltClientWithRetiredCaches(t *testing.T) (*rmqClient, <-chan
struct{}) {
+ t.Helper()
+ collected := make(chan struct{}, 3)
+ client := GetOrNewRocketMQClient(ownershipTestOptions(t),
nil).(*rmqClient)
+ for generation := 0; generation < 3; generation++ {
+ // Use a separately allocated leaf object: finalizers on the
client itself
+ // would interact with its intentional internal reference
cycles.
+ marker := new([1024]byte)
+ runtime.SetFinalizer(marker, func(*[1024]byte) { collected <-
struct{}{} })
+
client.GetNameSrv().(*namesrvs).routeDataMap.Store("retired-cache", marker)
+ options := client.option
+ client.Shutdown()
+ client = GetOrNewRocketMQClient(options, nil).(*rmqClient)
+ }
+ return client, collected
+}
diff --git a/internal/namesrv.go b/internal/namesrv.go
index 84042cd..91e0696 100644
--- a/internal/namesrv.go
+++ b/internal/namesrv.go
@@ -93,7 +93,9 @@ type namesrvs struct {
lockNamesrv sync.Mutex
- nameSrvClient remote.RemotingClient
+ nameSrvClient remote.RemotingClient
+ remotingConfig *remote.RemotingClientConfig
+ credentials *primitive.Credentials // protected by lock
resolver primitive.NsResolver
}
@@ -119,15 +121,32 @@ func NewNamesrv(resolver primitive.NsResolver, config
*remote.RemotingClientConf
if err := primitive.NamesrvAddr(addr).Check(); err != nil {
return nil, err
}
- nameSrvClient := remote.NewRemotingClient(config)
+ return newNamesrvWithAddresses(addr, resolver, config), nil
+}
+
+func newNamesrvWithAddresses(addr []string, resolver primitive.NsResolver,
config *remote.RemotingClientConfig) *namesrvs {
return &namesrvs{
srvs: append([]string(nil), addr...),
lock: new(sync.Mutex),
- nameSrvClient: nameSrvClient,
+ nameSrvClient: remote.NewRemotingClient(config),
+ remotingConfig: config,
brokerVersionMap: make(map[string]map[string]int32, 0),
brokerLock: new(sync.RWMutex),
resolver: resolver,
- }, nil
+ }
+}
+
+// cloneForClient retains discovery configuration but gives a new client its
own
+// transport and route caches. Do not resolve again: discovery can temporarily
be
+// unavailable, while the last known address snapshot is still usable.
+func (s *namesrvs) cloneForClient() *namesrvs {
+ s.lock.Lock()
+ defer s.lock.Unlock()
+ cloned := newNamesrvWithAddresses(s.srvs, s.resolver, s.remotingConfig)
+ if s.credentials != nil {
+ cloned.SetCredentials(*s.credentials)
+ }
+ return cloned
}
// getNameServerAddress return namesrv using round-robin strategy.
@@ -158,6 +177,9 @@ func (s *namesrvs) String() string {
return strings.Join(s.AddrList(), ";")
}
func (s *namesrvs) SetCredentials(credentials primitive.Credentials) {
+ s.lock.Lock()
+ defer s.lock.Unlock()
+ s.credentials = &credentials
s.nameSrvClient.RegisterInterceptor(remote.ACLInterceptor(credentials))
}
diff --git a/internal/remote/future.go b/internal/remote/future.go
index a7d3268..a74b314 100644
--- a/internal/remote/future.go
+++ b/internal/remote/future.go
@@ -31,6 +31,7 @@ type ResponseFuture struct {
callback func(*ResponseFuture)
Done chan bool
callbackOnce sync.Once
+ completeOnce sync.Once
ctx context.Context
}
@@ -52,17 +53,22 @@ func (r *ResponseFuture) executeInvokeCallback() {
})
}
+// complete publishes exactly one result before waking waiters. Callbacks run
+// outside transport locks, in the goroutine owning the asynchronous request.
+func (r *ResponseFuture) complete(cmd *RemotingCommand, err error) {
+ r.completeOnce.Do(func() {
+ r.ResponseCommand, r.Err = cmd, err
+ if r.Done != nil {
+ close(r.Done)
+ }
+ })
+}
+
func (r *ResponseFuture) waitResponse() (*RemotingCommand, error) {
- var (
- cmd *RemotingCommand
- err error
- )
select {
case <-r.Done:
- cmd, err = r.ResponseCommand, r.Err
case <-r.ctx.Done():
- err = errors.ErrRequestTimeout
- r.Err = err
+ r.complete(nil, errors.ErrRequestTimeout)
}
- return cmd, err
+ return r.ResponseCommand, r.Err
}
diff --git a/internal/remote/remote_client.go b/internal/remote/remote_client.go
index 55a2646..74a736a 100644
--- a/internal/remote/remote_client.go
+++ b/internal/remote/remote_client.go
@@ -21,6 +21,7 @@ import (
"bytes"
"context"
"encoding/binary"
+ "errors"
"github.com/apache/rocketmq-client-go/v2/internal/utils"
"io"
"net"
@@ -54,12 +55,17 @@ type RemotingClient interface {
var _ RemotingClient = &remotingClient{}
+var errClientClosed = errors.New("remoting client is shut down")
+
type remotingClient struct {
responseTable sync.Map
connectionTable sync.Map
config *RemotingClientConfig
processors map[int16]ClientRequestFunc
connectionLocker sync.Mutex
+ lifecycleMu sync.Mutex
+ closed bool // protected by lifecycleMu
+ done chan struct{}
interceptor primitive.Interceptor
}
@@ -84,6 +90,7 @@ func NewRemotingClient(config *RemotingClientConfig)
*remotingClient {
return &remotingClient{
processors: make(map[int16]ClientRequestFunc),
config: config,
+ done: make(chan struct{}),
}
}
@@ -100,7 +107,9 @@ func (c *remotingClient) InvokeSync(ctx context.Context,
addr string, request *R
resp := NewResponseFuture(ctx, request.Opaque, nil)
- c.responseTable.Store(resp.Opaque, resp)
+ if err := c.registerResponse(resp); err != nil {
+ return nil, err
+ }
defer c.responseTable.Delete(request.Opaque)
err = c.sendRequest(ctx, conn, request)
@@ -114,7 +123,9 @@ func (c *remotingClient) InvokeSync(ctx context.Context,
addr string, request *R
// InvokeAsync send request without blocking, just return immediately.
func (c *remotingClient) InvokeAsync(ctx context.Context, addr string, request
*RemotingCommand, callback func(*ResponseFuture)) error {
resp := NewResponseFuture(ctx, request.Opaque, callback)
- c.responseTable.Store(resp.Opaque, resp)
+ if err := c.registerResponse(resp); err != nil {
+ return err
+ }
go primitive.WithRecover(func() {
defer resp.executeInvokeCallback()
@@ -122,12 +133,12 @@ func (c *remotingClient) InvokeAsync(ctx context.Context,
addr string, request *
conn, err := c.connect(ctx, addr)
if err != nil {
- resp.Err = err
+ resp.complete(nil, err)
return
}
err = c.sendRequest(ctx, conn, request)
if err != nil {
- resp.Err = err
+ resp.complete(nil, err)
return
}
c.receiveAsync(resp)
@@ -151,18 +162,50 @@ func (c *remotingClient) InvokeOneWay(ctx
context.Context, addr string, request
return c.sendRequest(ctx, conn, request)
}
+func (c *remotingClient) registerResponse(resp *ResponseFuture) error {
+ c.lifecycleMu.Lock()
+ defer c.lifecycleMu.Unlock()
+ if c.closed {
+ return errClientClosed
+ }
+ c.responseTable.Store(resp.Opaque, resp)
+ return nil
+}
+
func (c *remotingClient) connect(ctx context.Context, addr string)
(*tcpConnWrapper, error) {
- // it needs additional locker.
c.connectionLocker.Lock()
defer c.connectionLocker.Unlock()
+ c.lifecycleMu.Lock()
+ if c.closed {
+ c.lifecycleMu.Unlock()
+ return nil, errClientClosed
+ }
conn, ok := c.connectionTable.Load(addr)
+ c.lifecycleMu.Unlock()
if ok {
return conn.(*tcpConnWrapper), nil
}
- tcpConn, err := initConn(ctx, addr, c.config)
+
+ // Cancel a pending TCP dial or TLS handshake when the transport stops.
+ dialCtx, cancel := context.WithCancel(ctx)
+ defer cancel()
+ go func() {
+ select {
+ case <-c.done:
+ cancel()
+ case <-dialCtx.Done():
+ }
+ }()
+ tcpConn, err := initConn(dialCtx, addr, c.config)
if err != nil {
return nil, err
}
+ c.lifecycleMu.Lock()
+ defer c.lifecycleMu.Unlock()
+ if c.closed {
+ tcpConn.destroy()
+ return nil, errClientClosed
+ }
c.connectionTable.Store(addr, tcpConn)
go primitive.WithRecover(func() {
c.receiveResponse(tcpConn)
@@ -243,13 +286,7 @@ func (c *remotingClient) processCMD(cmd *RemotingCommand,
r *tcpConnWrapper) {
if exist {
c.responseTable.Delete(cmd.Opaque)
responseFuture := resp.(*ResponseFuture)
- go primitive.WithRecover(func() {
- responseFuture.ResponseCommand = cmd
- if responseFuture.Done != nil {
- close(responseFuture.Done)
- }
- responseFuture.executeInvokeCallback()
- })
+ responseFuture.complete(cmd, nil)
}
} else {
f := c.processors[cmd.Code]
@@ -370,11 +407,20 @@ func (c *remotingClient) closeConnection(toCloseConn
*tcpConnWrapper) {
}
func (c *remotingClient) ShutDown() {
+ c.lifecycleMu.Lock()
+ defer c.lifecycleMu.Unlock()
+ if c.closed {
+ return
+ }
+ c.closed = true
+ close(c.done)
c.responseTable.Range(func(key, value interface{}) bool {
+ value.(*ResponseFuture).complete(nil, errClientClosed)
c.responseTable.Delete(key)
return true
})
c.connectionTable.Range(func(key, value interface{}) bool {
+ c.connectionTable.Delete(key)
conn := value.(*tcpConnWrapper)
err := conn.destroy()
if err != nil {
diff --git a/internal/remote/remote_client_test.go
b/internal/remote/remote_client_test.go
index b7171c4..5d51284 100644
--- a/internal/remote/remote_client_test.go
+++ b/internal/remote/remote_client_test.go
@@ -283,13 +283,14 @@ func TestInvokeAsyncTimeout(t *testing.T) {
var wg sync.WaitGroup
wg.Add(1)
client := NewRemotingClient(nil)
+ defer client.ShutDown()
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
var clientSend sync.WaitGroup // blocking client send message until the
server listen success.
clientSend.Add(1)
go func() {
clientSend.Wait()
- ctx, cancel := context.WithTimeout(context.Background(),
time.Duration(10*time.Second))
- defer cancel()
err := client.InvokeAsync(ctx, addr, clientSendRemtingCommand,
func(r *ResponseFuture) {
assert.NotNil(t, r.Err)
diff --git a/internal/remote/resource_lifecycle_test.go
b/internal/remote/resource_lifecycle_test.go
new file mode 100644
index 0000000..adbe432
--- /dev/null
+++ b/internal/remote/resource_lifecycle_test.go
@@ -0,0 +1,212 @@
+/*
+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 remote
+
+import (
+ "bufio"
+ "context"
+ "io"
+ "io/ioutil"
+ "net"
+ "net/http"
+ "net/http/httptest"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+)
+
+func awaitShutdownSignal(t *testing.T, signal <-chan struct{}) {
+ t.Helper()
+ select {
+ case <-signal:
+ case <-time.After(2 * time.Second):
+ t.Fatal("shutdown did not release the waiting operation")
+ }
+}
+
+func TestShutdownCompletesPendingRequests(t *testing.T) {
+ // Reuse the process while repeatedly constructing and disposing
transports.
+ for cycle := 0; cycle < 10; cycle++ {
+ listener, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ defer listener.Close()
+ client := NewRemotingClient(nil)
+ defer client.ShutDown()
+ requests := make(chan struct{}, 2)
+ disconnected := make(chan struct{})
+ go func() {
+ defer close(disconnected)
+ conn, err := listener.Accept()
+ if err != nil {
+ return
+ }
+ defer conn.Close()
+ scanner := client.createScanner(conn)
+ for scanner.Scan() {
+ requests <- struct{}{}
+ }
+ }()
+ syncDone := make(chan struct{})
+ go func() {
+ defer close(syncDone)
+ _, err := client.InvokeSync(context.Background(),
listener.Addr().String(), NewRemotingCommand(10, nil, nil))
+ if err == nil {
+ t.Error("pending sync request succeeded after
shutdown")
+ }
+ }()
+ asyncDone := make(chan struct{})
+ var callbacks int32
+ err = client.InvokeAsync(context.Background(),
listener.Addr().String(), NewRemotingCommand(10, nil, nil), func(f
*ResponseFuture) {
+ atomic.AddInt32(&callbacks, 1)
+ if f.Err == nil {
+ t.Error("pending async request succeeded after
shutdown")
+ }
+ // A callback may reenter Shutdown; it must not run
under the lifecycle lock.
+ client.ShutDown()
+ close(asyncDone)
+ })
+ require.NoError(t, err)
+ awaitShutdownSignal(t, requests)
+ awaitShutdownSignal(t, requests)
+ client.ShutDown()
+ awaitShutdownSignal(t, syncDone)
+ awaitShutdownSignal(t, asyncDone)
+ awaitShutdownSignal(t, disconnected)
+ require.EqualValues(t, 1, atomic.LoadInt32(&callbacks))
+ client.responseTable.Range(func(_, _ interface{}) bool {
t.Error("pending response retained"); return true })
+ client.connectionTable.Range(func(_, _ interface{}) bool {
t.Error("connection retained"); return true })
+ _, err = client.connect(context.Background(),
listener.Addr().String())
+ require.Equal(t, errClientClosed, err)
+ err = client.InvokeAsync(context.Background(),
listener.Addr().String(), NewRemotingCommand(10, nil, nil), nil)
+ require.Equal(t, errClientClosed, err)
+ }
+}
+
+func TestShutdownCancelsTLSHandshake(t *testing.T) {
+ listener, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ defer listener.Close()
+ config := DefaultRemotingClientConfig
+ config.UseTls = true
+ config.ConnectionTimeout = time.Minute
+ client := NewRemotingClient(&config)
+ defer client.ShutDown()
+ accepted := make(chan net.Conn, 1)
+ go func() {
+ conn, err := listener.Accept()
+ if err == nil {
+ accepted <- conn
+ }
+ }()
+ invoked := make(chan struct{})
+ go func() {
+ defer close(invoked)
+ _, err := client.connect(context.Background(),
listener.Addr().String())
+ if err == nil {
+ t.Error("stalled TLS handshake unexpectedly succeeded")
+ }
+ }()
+ var conn net.Conn
+ select {
+ case conn = <-accepted:
+ case <-time.After(2 * time.Second):
+ t.Fatal("TLS connection was not established")
+ }
+ defer conn.Close()
+ require.NoError(t, conn.SetReadDeadline(time.Now().Add(2*time.Second)))
+ // Wait until TLS handshake has started. The server deliberately never
replies.
+ _, err = io.ReadFull(conn, make([]byte, 1))
+ require.NoError(t, err)
+ stopped := make(chan struct{})
+ go func() { client.ShutDown(); close(stopped) }()
+ awaitShutdownSignal(t, stopped)
+ awaitShutdownSignal(t, invoked)
+}
+
+func TestResponseCompletionRaces(t *testing.T) {
+ for i := 0; i < 100; i++ {
+ ctx, cancel := context.WithCancel(context.Background())
+ f := NewResponseFuture(ctx, int32(i), nil)
+ cmd := NewRemotingCommand(10, nil, nil)
+ var wg sync.WaitGroup
+ wg.Add(3)
+ go func() { defer wg.Done(); f.complete(cmd, nil) }()
+ go func() { defer wg.Done(); f.complete(nil, errClientClosed)
}()
+ go func() { defer wg.Done(); cancel() }()
+ response, err := f.waitResponse()
+ wg.Wait()
+ require.True(t, response == cmd || err != nil, "completion lost
both response and error")
+ require.Same(t, response, f.ResponseCommand)
+ require.Equal(t, err, f.Err)
+ }
+}
+
+func TestConcurrentRequestRegistrationAndShutdown(t *testing.T) {
+ for i := 0; i < 20; i++ {
+ client := NewRemotingClient(nil)
+ start := make(chan struct{})
+ var wg sync.WaitGroup
+ for j := 0; j < 20; j++ {
+ wg.Add(1)
+ go func(id int32) {
+ defer wg.Done()
+ <-start
+ f := NewResponseFuture(context.Background(),
id, nil)
+ if client.registerResponse(f) == nil {
+ _, err := f.waitResponse()
+ if err != errClientClosed {
+ t.Error("registered request was
not canceled")
+ }
+ }
+ }(int32(j))
+ }
+ close(start)
+ client.ShutDown()
+ exited := make(chan struct{})
+ go func() { wg.Wait(); close(exited) }()
+ awaitShutdownSignal(t, exited)
+ }
+}
+
+// The dial context covers setup only; canceling it must not close a usable
TLS connection.
+func TestTLSConnectionSurvivesDialContextCancellation(t *testing.T) {
+ server := httptest.NewTLSServer(http.HandlerFunc(func(w
http.ResponseWriter, r *http.Request) {
+ io.WriteString(w, "connected")
+ }))
+ defer server.Close()
+ config := DefaultRemotingClientConfig
+ config.UseTls = true
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ conn, err := initConn(ctx, server.Listener.Addr().String(), &config)
+ require.NoError(t, err)
+ defer conn.destroy()
+ cancel()
+ require.NoError(t, conn.SetDeadline(time.Now().Add(2*time.Second)))
+ _, err = io.WriteString(conn, "GET / HTTP/1.1\r\nHost:
localhost\r\nConnection: close\r\n\r\n")
+ require.NoError(t, err)
+ response, err := http.ReadResponse(bufio.NewReader(conn), nil)
+ require.NoError(t, err)
+ defer response.Body.Close()
+ body, err := ioutil.ReadAll(response.Body)
+ require.NoError(t, err)
+ require.Equal(t, "connected", string(body))
+}
diff --git a/internal/remote/tcp_conn.go b/internal/remote/tcp_conn.go
index 3b7d164..2954dc1 100644
--- a/internal/remote/tcp_conn.go
+++ b/internal/remote/tcp_conn.go
@@ -38,19 +38,47 @@ func initConn(ctx context.Context, addr string, config
*RemotingClientConfig) (*
d.KeepAlive = config.KeepAliveDuration
d.Deadline = time.Now().Add(config.ConnectionTimeout)
- var conn net.Conn
- var err error
- if config.UseTls {
- conn, err = tls.DialWithDialer(&d, "tcp", addr, &tls.Config{
- InsecureSkipVerify: true,
- })
- } else {
- conn, err = d.DialContext(ctx, "tcp", addr)
- }
-
+ conn, err := d.DialContext(ctx, "tcp", addr)
if err != nil {
return nil, err
}
+ if config.UseTls {
+ host, _, _ := net.SplitHostPort(addr)
+ tlsConn := tls.Client(conn, &tls.Config{ServerName: host,
InsecureSkipVerify: true})
+ deadline := d.Deadline
+ if ctxDeadline, ok := ctx.Deadline(); ok &&
ctxDeadline.Before(deadline) {
+ deadline = ctxDeadline
+ }
+ if err := conn.SetDeadline(deadline); err != nil {
+ conn.Close()
+ return nil, err
+ }
+ handshakeDone := make(chan struct{})
+ watcherDone := make(chan struct{})
+ go func() {
+ defer close(watcherDone)
+ select {
+ case <-ctx.Done():
+ conn.Close()
+ case <-handshakeDone:
+ }
+ }()
+ err = tlsConn.Handshake()
+ close(handshakeDone)
+ <-watcherDone
+ if err == nil {
+ err = ctx.Err()
+ }
+ if err == nil {
+ err = conn.SetDeadline(time.Time{})
+ }
+ if err != nil {
+ conn.Close()
+ return nil, err
+ }
+ conn = tlsConn
+ }
+
return &tcpConnWrapper{
Conn: conn,
}, nil
diff --git a/internal/resource_lifecycle_test.go
b/internal/resource_lifecycle_test.go
new file mode 100644
index 0000000..42b6fae
--- /dev/null
+++ b/internal/resource_lifecycle_test.go
@@ -0,0 +1,141 @@
+/*
+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 internal
+
+import (
+ "bytes"
+ "context"
+ "errors"
+ "github.com/golang/mock/gomock"
+ "io/ioutil"
+ "net"
+ "runtime/pprof"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/apache/rocketmq-client-go/v2/internal/remote"
+ "github.com/apache/rocketmq-client-go/v2/primitive"
+ "github.com/stretchr/testify/require"
+)
+
+func TestLastClientReleaseClosesNameServerConnection(t *testing.T) {
+ listener, err := net.Listen("tcp", "127.0.0.1:0")
+ require.NoError(t, err)
+ defer listener.Close()
+ accepted := make(chan net.Conn, 1)
+ disconnected := make(chan struct{})
+ go func() {
+ defer close(disconnected)
+ conn, err := listener.Accept()
+ if err != nil {
+ return
+ }
+ defer conn.Close()
+ accepted <- conn
+ ioutil.ReadAll(conn)
+ }()
+ srvs, err :=
NewNamesrv(primitive.NewPassthroughResolver([]string{listener.Addr().String()}),
nil)
+ require.NoError(t, err)
+ options := DefaultClientOptions()
+ options.InstanceName = t.Name()
+ options.Namesrv = srvs
+ first := GetOrNewRocketMQClient(options, nil)
+ second := GetOrNewRocketMQClient(options, nil)
+ require.NotNil(t, first)
+ require.Same(t, first, second)
+ defer second.Shutdown()
+
+ ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+ defer cancel()
+ err = srvs.nameSrvClient.InvokeOneWay(ctx, listener.Addr().String(),
remote.NewRemotingCommand(ReqGetRouteInfoByTopic, nil, nil))
+ require.NoError(t, err)
+ select {
+ case conn := <-accepted:
+ defer conn.Close()
+ case <-ctx.Done():
+ t.Fatal("no NameServer connection was opened")
+ }
+ first.Shutdown()
+ // A remaining owner must still be able to use the connection.
+ err = srvs.nameSrvClient.InvokeOneWay(ctx, listener.Addr().String(),
remote.NewRemotingCommand(ReqGetRouteInfoByTopic, nil, nil))
+ require.NoError(t, err)
+ second.Shutdown()
+ select {
+ case <-disconnected:
+ case <-time.After(time.Second):
+ t.Fatal("NameServer connection remained open after the last
release")
+ }
+ // Late refresh work must not recreate the disposed transport's
connection.
+ err = srvs.nameSrvClient.InvokeOneWay(ctx, listener.Addr().String(),
remote.NewRemotingCommand(ReqGetRouteInfoByTopic, nil, nil))
+ require.Error(t, err)
+}
+
+func TestConcurrentShutdownAndInvoke(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ defer ctrl.Finish()
+ r := remote.NewMockRemotingClient(ctrl)
+ r.EXPECT().InvokeSync(gomock.Any(), gomock.Any(),
gomock.Any()).Return(nil, errors.New("test transport")).AnyTimes()
+ r.EXPECT().ShutDown()
+ opts := DefaultClientOptions()
+ opts.InstanceName = t.Name()
+ srvs, err :=
NewNamesrv(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"}), nil)
+ require.NoError(t, err)
+ opts.Namesrv = srvs
+ c := GetOrNewRocketMQClient(opts, nil).(*rmqClient)
+ c.remoteClient = r
+ started, stop, stopped := make(chan struct{}), make(chan struct{}),
make(chan struct{})
+ go func() {
+ defer close(stopped)
+ close(started)
+ for {
+ select {
+ case <-stop:
+ return
+ default:
+ c.InvokeSync(context.Background(), "unused",
remote.NewRemotingCommand(10, nil, nil), time.Second)
+ }
+ }
+ }()
+ <-started
+ time.Sleep(time.Millisecond)
+ c.Shutdown()
+ close(stop)
+ <-stopped
+}
+
+func clientBackgroundWorkers() int {
+ var b bytes.Buffer
+ pprof.Lookup("goroutine").WriteTo(&b, 2)
+ return strings.Count(b.String(), "(*rmqClient).Start.func")
+}
+
+func TestClientRepeatedLifecycleStopsBackgroundWorkers(t *testing.T) {
+ baseline := clientBackgroundWorkers()
+ for i := 0; i < 10; i++ {
+ srvs, err :=
NewNamesrv(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"}), nil)
+ require.NoError(t, err)
+ opts := DefaultClientOptions()
+ opts.InstanceName = t.Name()
+ opts.Namesrv = srvs
+ c := GetOrNewRocketMQClient(opts, nil)
+ c.Start()
+ c.Shutdown()
+ require.Eventually(t, func() bool { return
clientBackgroundWorkers() <= baseline }, time.Second, time.Millisecond,
"shutdown left discovery/heartbeat workers waiting")
+ }
+}
diff --git a/internal/utils/wait.go b/internal/utils/wait.go
new file mode 100644
index 0000000..233f3aa
--- /dev/null
+++ b/internal/utils/wait.go
@@ -0,0 +1,43 @@
+/*
+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 utils
+
+import "time"
+
+// WaitFor waits for a duration, returning false when shutdown interrupts it.
+func WaitFor(done <-chan struct{}, delay time.Duration) bool {
+ select {
+ case <-done:
+ return false
+ default:
+ }
+ timer := time.NewTimer(delay)
+ defer timer.Stop()
+ select {
+ case <-done:
+ return false
+ case <-timer.C:
+ // Prefer shutdown if both events became ready together.
+ select {
+ case <-done:
+ return false
+ default:
+ return true
+ }
+ }
+}
diff --git a/producer/producer.go b/producer/producer.go
index 59dc182..e490f74 100644
--- a/producer/producer.go
+++ b/producer/producer.go
@@ -44,11 +44,13 @@ type defaultProducer struct {
options producerOptions
publishInfo sync.Map
callbackCh chan interface{}
+ done chan struct{}
interceptor primitive.Interceptor
lifecycleMu sync.Mutex
startOnce sync.Once
+ startErr error
ShutdownOnce sync.Once
}
@@ -75,9 +77,10 @@ func NewDefaultProducer(opts ...Option) (*defaultProducer,
error) {
producer := &defaultProducer{
group: defaultOpts.GroupName,
callbackCh: make(chan interface{}),
+ done: make(chan struct{}),
options: defaultOpts,
}
- producer.client =
internal.GetOrNewRocketMQClient(defaultOpts.ClientOptions, producer.callbackCh)
+ producer.client =
internal.GetOrNewRocketMQClient(defaultOpts.ClientOptions, producer.callbackCh,
producer.done)
if producer.client == nil {
return nil, fmt.Errorf("GetOrNewRocketMQClient faild")
}
@@ -95,31 +98,33 @@ func (p *defaultProducer) Start() error {
if atomic.LoadInt32(&p.state) == int32(internal.StateShutdown) {
return internal.ErrServiceState
}
- var err error
defer func() {
- if err != nil &&
!internal.IsNilTraceDispatcher(p.options.TraceDispatcher) {
+ if p.startErr != nil &&
!internal.IsNilTraceDispatcher(p.options.TraceDispatcher) {
p.options.TraceDispatcher.Close()
}
}()
p.startOnce.Do(func() {
- err = p.client.RegisterProducer(p.group, p)
- if err != nil {
+ p.startErr = p.client.RegisterProducer(p.group, p)
+ if p.startErr != nil {
rlog.Error("the producer group has been created,
specify another one", map[string]interface{}{
rlog.LogKeyProducerGroup: p.group,
})
- err = errors2.ErrProducerCreated
+ p.startErr = errors2.ErrProducerCreated
return
}
p.client.Start()
atomic.StoreInt32(&p.state, int32(internal.StateRunning))
})
- return err
+ return p.startErr
}
func (p *defaultProducer) Shutdown() error {
p.lifecycleMu.Lock()
defer p.lifecycleMu.Unlock()
p.ShutdownOnce.Do(func() {
+ if p.done != nil {
+ close(p.done)
+ }
if !internal.IsNilTraceDispatcher(p.options.TraceDispatcher) {
p.options.TraceDispatcher.Close()
}
@@ -657,8 +662,9 @@ func (p *defaultProducer) IsUnitMode() bool {
}
type transactionProducer struct {
- producer *defaultProducer
- listener primitive.TransactionListener
+ producer *defaultProducer
+ listener primitive.TransactionListener
+ checkOnce sync.Once
}
// TODO: checkLocalTransaction
@@ -674,10 +680,15 @@ func NewTransactionProducer(listener
primitive.TransactionListener, opts ...Opti
}
func (tp *transactionProducer) Start() error {
- go primitive.WithRecover(func() {
- tp.checkTransactionState()
+ if err := tp.producer.Start(); err != nil {
+ return err
+ }
+ tp.checkOnce.Do(func() {
+ go primitive.WithRecover(func() {
+ tp.checkTransactionState()
+ })
})
- return tp.producer.Start()
+ return nil
}
func (tp *transactionProducer) Shutdown() error {
return tp.producer.Shutdown()
@@ -685,7 +696,20 @@ func (tp *transactionProducer) Shutdown() error {
// TODO: check addr
func (tp *transactionProducer) checkTransactionState() {
- for ch := range tp.producer.callbackCh {
+ for {
+ var ch interface{}
+ select {
+ case <-tp.producer.done:
+ return
+ case ch = <-tp.producer.callbackCh:
+ }
+ // A callback can arrive at the same time as shutdown. Do not
begin new
+ // application work once the producer has stopped.
+ select {
+ case <-tp.producer.done:
+ return
+ default:
+ }
switch callback := ch.(type) {
case *internal.CheckTransactionStateCallback:
localTransactionState :=
tp.listener.CheckLocalTransaction(callback.Msg)
diff --git a/producer/resource_lifecycle_test.go
b/producer/resource_lifecycle_test.go
new file mode 100644
index 0000000..222e7d4
--- /dev/null
+++ b/producer/resource_lifecycle_test.go
@@ -0,0 +1,86 @@
+/*
+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 producer
+
+import (
+ "bytes"
+ "runtime/pprof"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/apache/rocketmq-client-go/v2/primitive"
+ "github.com/stretchr/testify/require"
+)
+
+func TestTransactionWorkerStopsOnShutdown(t *testing.T) {
+ // Keep another owner alive: stopping this worker must not require
disposing
+ // the shared RMQ client used by the remaining producer.
+ other, err := NewDefaultProducer(WithInstanceName(t.Name()),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer other.Shutdown()
+
+ tp, err := NewTransactionProducer(nil,
+ WithInstanceName(t.Name()),
+
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"})))
+ require.NoError(t, err)
+ defer tp.Shutdown()
+
+ // Observe the actual worker's exit without relying on process-wide
+ // goroutine counts or sleeping until an unrelated SDK worker exits.
+ exited := make(chan struct{})
+ go func() {
+ defer close(exited)
+ tp.checkTransactionState()
+ }()
+ require.NoError(t, tp.Shutdown())
+ select {
+ case <-exited:
+ case <-time.After(time.Second):
+ t.Fatal("transaction worker remained blocked after shutdown")
+ }
+ require.Error(t, tp.Start())
+}
+
+func transactionWorkerCount() int {
+ var b bytes.Buffer
+ pprof.Lookup("goroutine").WriteTo(&b, 2)
+ return strings.Count(b.String(),
"(*transactionProducer).checkTransactionState(")
+}
+
+func TestTransactionRepeatedLifecycle(t *testing.T) {
+ baseline := transactionWorkerCount()
+ for i := 0; i < 10; i++ {
+ opts := []Option{WithInstanceName(t.Name()),
WithGroupName(t.Name()),
WithNsResolver(primitive.NewPassthroughResolver([]string{"127.0.0.1:9876"}))}
+ p, err := NewTransactionProducer(nil, opts...)
+ require.NoError(t, err)
+ defer p.Shutdown()
+ for j := 0; j < 3; j++ {
+ require.NoError(t, p.Start())
+ }
+ require.Eventually(t, func() bool { return
transactionWorkerCount() == baseline+1 }, time.Second, time.Millisecond)
+ failed, err := NewTransactionProducer(nil, opts...)
+ require.NoError(t, err)
+ defer failed.Shutdown()
+ require.Error(t, failed.Start())
+ require.Error(t, failed.Start(), "repeated failed Start must
not launch a transaction worker")
+ require.NoError(t, failed.Shutdown())
+ require.NoError(t, p.Shutdown())
+ require.Eventually(t, func() bool { return
transactionWorkerCount() <= baseline }, time.Second, time.Millisecond)
+ }
+}
