Repository: apache/rocketmq-client-go
Pull request #172 (2019-08-23 07:48:21)

[ISSUE #130] Add Shutdown of PushConsumer
## What is the purpose of the change

Add Shutdown of PushConsumer

## Brief changelog

- Add Shutdown of PushConsumer
- Add mqChanged impl of PushConsumer

## Verifying this change

XXXX

Follow this checklist to help us incorporate your contribution quickly and easily. Notice, `it would be helpful if you could finish the following 5 checklist(the last one is not necessary)before request the community to review your PR`.

- [x] Make sure there is a [Github issue](https://github.com/apache/rocketmq/issues) filed for the change (usually before you start working on it). Trivial changes like typos do not require a Github issue. Your pull request should address just this issue, without pulling in other changes - one PR resolves one issue. 
- [x] Format the pull request title like `[ISSUE #123] Fix UnknownException when host config not exist`. Each commit in the pull request should have a meaningful subject line and body.
- [x] Write a pull request description that is detailed enough to understand what the pull request does, how, and why.
- [x] Write necessary unit-test(over 80% coverage) to verify your logic correction, more mock a little better when a cross-module dependency exists.
- [ ] If this contribution is large, please file an [Apache Individual Contributor License Agreement](http://www.apache.org/licenses/#clas).


codecov-io: # [Codecov](https://codecov.io/gh/apache/rocketmq-client-go/pull/172?src=pr&el=h1) Report
> Merging [#172](https://codecov.io/gh/apache/rocketmq-client-go/pull/172?src=pr&el=desc) into [native](https://codecov.io/gh/apache/rocketmq-client-go/commit/e572801dd746d6cae4e577d0edd5b45d4b63cccc?src=pr&el=desc) will **decrease** coverage by `0.18%`.
> The diff coverage is `2.7%`.

[![Impacted file tree graph](https://codecov.io/gh/apache/rocketmq-client-go/pull/172/graphs/tree.svg?width=650&token=CiV9ofPkG3&height=150&src=pr)](https://codecov.io/gh/apache/rocketmq-client-go/pull/172?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           native     #172      +/-   ##
==========================================
- Coverage   24.34%   24.16%   -0.19%     
==========================================
  Files          44       44              
  Lines        4608     4644      +36     
==========================================
  Hits         1122     1122              
- Misses       3338     3374      +36     
  Partials      148      148
```


| [Impacted Files](https://codecov.io/gh/apache/rocketmq-client-go/pull/172?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [consumer/consumer.go](https://codecov.io/gh/apache/rocketmq-client-go/pull/172/diff?src=pr&el=tree#diff-Y29uc3VtZXIvY29uc3VtZXIuZ28=) | `3.79% <0%> (-0.07%)` | :arrow_down: |
| [consumer/push\_consumer.go](https://codecov.io/gh/apache/rocketmq-client-go/pull/172/diff?src=pr&el=tree#diff-Y29uc3VtZXIvcHVzaF9jb25zdW1lci5nbw==) | `14.88% <3.57%> (-0.78%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/apache/rocketmq-client-go/pull/172?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/apache/rocketmq-client-go/pull/172?src=pr&el=footer). Last update [e572801...94d96f2](https://codecov.io/gh/apache/rocketmq-client-go/pull/172?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/api.go b/api.go
index ca54668..931dbcf 100644
--- a/api.go
+++ b/api.go
@@ -19,6 +19,7 @@ package rocketmq
 
 import (
 	"context"
+	"github.com/pkg/errors"
 
 	"github.com/apache/rocketmq-client-go/consumer"
 	"github.com/apache/rocketmq-client-go/primitive"
@@ -127,6 +128,10 @@ type PullConsumer interface {
 	//Resume(mqs ...primitive.MessageQueue) error
 }
 
+// The PullConsumer has not implemented completely, if you want have an experience of PullConsumer, you could use
+// consumer.NewPullConsumer(...), but it may changed in the future.
+//
+// The PullConsumer will be supported in next release
 func NewPullConsumer(opts ...consumer.Option) (PullConsumer, error) {
-	return nil, nil
+	return nil, errors.New("pull consumer has not supported")
 }
diff --git a/config.go b/config.go
index 5519947..8f1592c 100644
--- a/config.go
+++ b/config.go
@@ -1,1 +1,17 @@
+/*
+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
diff --git a/consumer/consumer.go b/consumer/consumer.go
index 9bed213..0564eaf 100644
--- a/consumer/consumer.go
+++ b/consumer/consumer.go
@@ -292,9 +292,17 @@ func (dc *defaultConsumer) start() error {
 }
 
 func (dc *defaultConsumer) shutdown() error {
-	dc.state = internal.StateRunning
+	dc.state = internal.StateShutdown
+	mqs := make([]*primitive.MessageQueue, 0)
+	dc.processQueueTable.Range(func(key, value interface{}) bool {
+		k := key.(primitive.MessageQueue)
+		pq := value.(*processQueue)
+		pq.dropped = true
+		mqs = append(mqs, &k)
+		return true
+	})
+	dc.storage.persist(mqs)
 	dc.client.Shutdown()
-
 	return nil
 }
 
@@ -397,7 +405,6 @@ func (dc *defaultConsumer) doBalance() {
 					"topic=%s, clientId=%s, mqAllSize=%d, cidAllSize=%d, rebalanceResultSize=%d, "+
 					"rebalanceResultSet=%v", dc.consumerGroup, topic, dc.client.ClientID(), len(mqAll),
 					len(cidAll), len(allocateResult), allocateResult)
-
 			}
 		}
 		return true
@@ -546,7 +553,7 @@ func (dc *defaultConsumer) unlockAll(oneway bool) {
 func (dc *defaultConsumer) doLock(addr string, body *lockBatchRequestBody) []primitive.MessageQueue {
 	data, _ := json.Marshal(body)
 	request := remote.NewRemotingCommand(internal.ReqLockBatchMQ, nil, data)
-	response, err := dc.client.InvokeSync(addr, request, 1*time.Second)
+	response, err := dc.client.InvokeSync(context.Background(), addr, request, 1*time.Second)
 	if err != nil {
 		rlog.Errorf("lock mq to broker: %s error %s", addr, err.Error())
 		return nil
@@ -566,12 +573,12 @@ func (dc *defaultConsumer) doUnlock(addr string, body *lockBatchRequestBody, one
 	data, _ := json.Marshal(body)
 	request := remote.NewRemotingCommand(internal.ReqUnlockBatchMQ, nil, data)
 	if oneway {
-		err := dc.client.InvokeOneWay(addr, request, 3*time.Second)
+		err := dc.client.InvokeOneWay(context.Background(), addr, request, 3*time.Second)
 		if err != nil {
 			rlog.Errorf("lock mq to broker with oneway: %s error %s", addr, err.Error())
 		}
 	} else {
-		response, err := dc.client.InvokeSync(addr, request, 1*time.Second)
+		response, err := dc.client.InvokeSync(context.Background(), addr, request, 1*time.Second)
 		if err != nil {
 			rlog.Errorf("lock mq to broker: %s error %s", addr, err.Error())
 		}
@@ -832,7 +839,7 @@ func (dc *defaultConsumer) findConsumerList(topic string) []string {
 			ConsumerGroup: dc.consumerGroup,
 		}
 		cmd := remote.NewRemotingCommand(internal.ReqGetConsumerListByGroup, req, nil)
-		res, err := dc.client.InvokeSync(brokerAddr, cmd, 3*time.Second) // TODO 超时机制有问题
+		res, err := dc.client.InvokeSync(context.Background(), brokerAddr, cmd, 3*time.Second) // TODO 超时机制有问题
 		if err != nil {
 			rlog.Errorf("get consumer list of [%s] from %s error: %s", dc.consumerGroup, brokerAddr, err.Error())
 			return nil
@@ -869,7 +876,7 @@ func (dc *defaultConsumer) queryMaxOffset(mq *primitive.MessageQueue) (int64, er
 	}
 
 	cmd := remote.NewRemotingCommand(internal.ReqGetMaxOffset, request, nil)
-	response, err := dc.client.InvokeSync(brokerAddr, cmd, 3*time.Second)
+	response, err := dc.client.InvokeSync(context.Background(), brokerAddr, cmd, 3*time.Second)
 	if err != nil {
 		return -1, err
 	}
@@ -899,7 +906,7 @@ func (dc *defaultConsumer) searchOffsetByTimestamp(mq *primitive.MessageQueue, t
 	}
 
 	cmd := remote.NewRemotingCommand(internal.ReqSearchOffsetByTimestamp, request, nil)
-	response, err := dc.client.InvokeSync(brokerAddr, cmd, 3*time.Second)
+	response, err := dc.client.InvokeSync(context.Background(), brokerAddr, cmd, 3*time.Second)
 	if err != nil {
 		return -1, err
 	}
diff --git a/consumer/mock_offset_store.go b/consumer/mock_offset_store.go
new file mode 100644
index 0000000..093c50d
--- /dev/null
+++ b/consumer/mock_offset_store.go
@@ -0,0 +1,76 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: offset_store.go
+
+// Package consumer is a generated GoMock package.
+package consumer
+
+import (
+	primitive "github.com/apache/rocketmq-client-go/primitive"
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+)
+
+// MockOffsetStore is a mock of OffsetStore interface
+type MockOffsetStore struct {
+	ctrl     *gomock.Controller
+	recorder *MockOffsetStoreMockRecorder
+}
+
+// MockOffsetStoreMockRecorder is the mock recorder for MockOffsetStore
+type MockOffsetStoreMockRecorder struct {
+	mock *MockOffsetStore
+}
+
+// NewMockOffsetStore creates a new mock instance
+func NewMockOffsetStore(ctrl *gomock.Controller) *MockOffsetStore {
+	mock := &MockOffsetStore{ctrl: ctrl}
+	mock.recorder = &MockOffsetStoreMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockOffsetStore) EXPECT() *MockOffsetStoreMockRecorder {
+	return m.recorder
+}
+
+// persist mocks base method
+func (m *MockOffsetStore) persist(mqs []*primitive.MessageQueue) {
+	m.ctrl.Call(m, "persist", mqs)
+}
+
+// persist indicates an expected call of persist
+func (mr *MockOffsetStoreMockRecorder) persist(mqs interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "persist", reflect.TypeOf((*MockOffsetStore)(nil).persist), mqs)
+}
+
+// remove mocks base method
+func (m *MockOffsetStore) remove(mq *primitive.MessageQueue) {
+	m.ctrl.Call(m, "remove", mq)
+}
+
+// remove indicates an expected call of remove
+func (mr *MockOffsetStoreMockRecorder) remove(mq interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "remove", reflect.TypeOf((*MockOffsetStore)(nil).remove), mq)
+}
+
+// read mocks base method
+func (m *MockOffsetStore) read(mq *primitive.MessageQueue, t readType) int64 {
+	ret := m.ctrl.Call(m, "read", mq, t)
+	ret0, _ := ret[0].(int64)
+	return ret0
+}
+
+// read indicates an expected call of read
+func (mr *MockOffsetStoreMockRecorder) read(mq, t interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "read", reflect.TypeOf((*MockOffsetStore)(nil).read), mq, t)
+}
+
+// update mocks base method
+func (m *MockOffsetStore) update(mq *primitive.MessageQueue, offset int64, increaseOnly bool) {
+	m.ctrl.Call(m, "update", mq, offset, increaseOnly)
+}
+
+// update indicates an expected call of update
+func (mr *MockOffsetStoreMockRecorder) update(mq, offset, increaseOnly interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "update", reflect.TypeOf((*MockOffsetStore)(nil).update), mq, offset, increaseOnly)
+}
diff --git a/consumer/offset_store.go b/consumer/offset_store.go
index f799c1a..71d2c87 100644
--- a/consumer/offset_store.go
+++ b/consumer/offset_store.go
@@ -18,6 +18,7 @@ limitations under the License.
 package consumer
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -51,6 +52,7 @@ func init() {
 	}
 }
 
+//go:generate mockgen -source offset_store.go -destination mock_offset_store.go -self_package github.com/apache/rocketmq-client-go/consumer  --package consumer OffsetStore
 type OffsetStore interface {
 	persist(mqs []*primitive.MessageQueue)
 	remove(mq *primitive.MessageQueue)
@@ -301,7 +303,7 @@ func (r *remoteBrokerOffsetStore) fetchConsumeOffsetFromBroker(group string, mq
 		QueueId:       mq.QueueId,
 	}
 	cmd := remote.NewRemotingCommand(internal.ReqQueryConsumerOffset, queryOffsetRequest, nil)
-	res, err := r.client.InvokeSync(broker, cmd, 3*time.Second)
+	res, err := r.client.InvokeSync(context.Background(), broker, cmd, 3*time.Second)
 	if err != nil {
 		return -1, err
 	}
@@ -335,7 +337,7 @@ func (r *remoteBrokerOffsetStore) updateConsumeOffsetToBroker(group, topic strin
 		CommitOffset:  queue.Offset,
 	}
 	cmd := remote.NewRemotingCommand(internal.ReqUpdateConsumerOffset, updateOffsetRequest, nil)
-	return r.client.InvokeOneWay(broker, cmd, 5*time.Second)
+	return r.client.InvokeOneWay(context.Background(), broker, cmd, 5*time.Second)
 }
 
 func readFromMemory(table map[string]map[int]*queueOffset, mq *primitive.MessageQueue) int64 {
diff --git a/consumer/push_consumer.go b/consumer/push_consumer.go
index d5a5ffc..dfee161 100644
--- a/consumer/push_consumer.go
+++ b/consumer/push_consumer.go
@@ -164,7 +164,7 @@ func (pc *pushConsumer) Start() error {
 }
 
 func (pc *pushConsumer) Shutdown() error {
-	return nil
+	return pc.defaultConsumer.shutdown()
 }
 
 func (pc *pushConsumer) Subscribe(topic string, selector MessageSelector,
@@ -220,7 +220,41 @@ func (pc *pushConsumer) IsUnitMode() bool {
 }
 
 func (pc *pushConsumer) messageQueueChanged(topic string, mqAll, mqDivided []*primitive.MessageQueue) {
-	// TODO
+	v, exit := pc.subscriptionDataTable.Load(topic)
+	if !exit {
+		return
+	}
+	data := v.(*internal.SubscriptionData)
+	newVersion := time.Now().UnixNano()
+	rlog.Infof("the MessageQueue changed, also update version: %d to %d", data.SubVersion, newVersion)
+	data.SubVersion = newVersion
+
+	// TODO: optimize
+	count := 0
+	pc.processQueueTable.Range(func(key, value interface{}) bool {
+		count++
+		return true
+	})
+	if count > 0 {
+		if pc.option.PullThresholdForTopic != -1 {
+			newVal := pc.option.PullThresholdForTopic / count
+			if newVal == 0 {
+				newVal = 1
+			}
+			rlog.Info("The PullThresholdForTopic is changed from %d to %d", pc.option.PullThresholdForTopic, newVal)
+			pc.option.PullThresholdForTopic = newVal
+		}
+
+		if pc.option.PullThresholdSizeForTopic != -1 {
+			newVal := pc.option.PullThresholdSizeForTopic / count
+			if newVal == 0 {
+				newVal = 1
+			}
+			rlog.Info("The PullThresholdSizeForTopic is changed from %d to %d", pc.option.PullThresholdSizeForTopic, newVal)
+			pc.option.PullThresholdSizeForTopic = newVal
+		}
+	}
+	pc.client.SendHeartbeatToAllBrokerWithLock()
 }
 
 func (pc *pushConsumer) validate() {
@@ -515,7 +549,7 @@ func (pc *pushConsumer) sendMessageBack(brokerName string, msg *primitive.Messag
 	} else {
 		brokerAddr = msg.StoreHost
 	}
-	_, err := pc.client.InvokeSync(brokerAddr, pc.buildSendBackRequest(msg, delayLevel), 3*time.Second)
+	_, err := pc.client.InvokeSync(context.Background(), brokerAddr, pc.buildSendBackRequest(msg, delayLevel), 3*time.Second)
 	if err != nil {
 		return false
 	}
diff --git a/examples/consumer/pull/main.go b/examples/consumer/pull/main.go
index c2915f9..d09da48 100644
--- a/examples/consumer/pull/main.go
+++ b/examples/consumer/pull/main.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/apache/rocketmq-client-go"
 	"github.com/apache/rocketmq-client-go/consumer"
 	"github.com/apache/rocketmq-client-go/internal/utils"
 	"github.com/apache/rocketmq-client-go/primitive"
@@ -29,14 +30,20 @@ import (
 )
 
 func main() {
-	c, err := consumer.NewPullConsumer(consumer.WithGroupName("testGroup"), consumer.WithNameServer([]string{"127.0.0.1:9876"}))
+	c, err := rocketmq.NewPullConsumer(
+		consumer.WithGroupName("testGroup"),
+		consumer.WithNameServer([]string{"127.0.0.1:9876"}),
+	)
+	if err != nil {
+		rlog.Fatal("fail to new pullConsumer: ", err)
+	}
+	err = c.Start()
 	if err != nil {
 		rlog.Fatal("fail to new pullConsumer: ", err)
 	}
-	c.Start()
 
 	ctx := context.Background()
-	queue := &primitive.MessageQueue{
+	queue := primitive.MessageQueue{
 		Topic:      "TopicTest",
 		BrokerName: "", // replace with your broker name. otherwise, pull will failed.
 		QueueId:    0,
diff --git a/internal/client.go b/internal/client.go
index 1307ff2..cb4cfb6 100644
--- a/internal/client.go
+++ b/internal/client.go
@@ -129,11 +129,11 @@ type RMQClient interface {
 	ClientID() string
 
 	RegisterProducer(group string, producer InnerProducer)
-	InvokeSync(addr string, request *remote.RemotingCommand,
+	InvokeSync(ctx context.Context, addr string, request *remote.RemotingCommand,
 		timeoutMillis time.Duration) (*remote.RemotingCommand, error)
-	InvokeAsync(addr string, request *remote.RemotingCommand,
+	InvokeAsync(ctx context.Context, addr string, request *remote.RemotingCommand,
 		timeoutMillis time.Duration, f func(*remote.RemotingCommand, error)) error
-	InvokeOneWay(addr string, request *remote.RemotingCommand,
+	InvokeOneWay(ctx context.Context, addr string, request *remote.RemotingCommand,
 		timeoutMillis time.Duration) error
 	CheckClientInBroker()
 	SendHeartbeatToAllBrokerWithLock()
@@ -160,7 +160,7 @@ type rmqClient struct {
 	consumerMap sync.Map
 	once        sync.Once
 
-	remoteClient *remote.RemotingClient
+	remoteClient remote.RemotingClient
 	hbMutex      sync.Mutex
 	close        bool
 	namesrvs     *namesrvs
@@ -291,31 +291,31 @@ func (c *rmqClient) ClientID() string {
 	return id
 }
 
-func (c *rmqClient) InvokeSync(addr string, request *remote.RemotingCommand,
+func (c *rmqClient) InvokeSync(ctx context.Context, addr string, request *remote.RemotingCommand,
 	timeoutMillis time.Duration) (*remote.RemotingCommand, error) {
 	if c.close {
 		return nil, ErrServiceState
 	}
-	return c.remoteClient.InvokeSync(addr, request, timeoutMillis)
+	return c.remoteClient.InvokeSync(ctx, addr, request, timeoutMillis)
 }
 
-func (c *rmqClient) InvokeAsync(addr string, request *remote.RemotingCommand,
+func (c *rmqClient) InvokeAsync(ctx context.Context, addr string, request *remote.RemotingCommand,
 	timeoutMillis time.Duration, f func(*remote.RemotingCommand, error)) error {
 	if c.close {
 		return ErrServiceState
 	}
-	return c.remoteClient.InvokeAsync(addr, request, timeoutMillis, func(future *remote.ResponseFuture) {
+	return c.remoteClient.InvokeAsync(ctx, addr, request, timeoutMillis, func(future *remote.ResponseFuture) {
 		f(future.ResponseCommand, future.Err)
 	})
 
 }
 
-func (c *rmqClient) InvokeOneWay(addr string, request *remote.RemotingCommand,
+func (c *rmqClient) InvokeOneWay(ctx context.Context, addr string, request *remote.RemotingCommand,
 	timeoutMillis time.Duration) error {
 	if c.close {
 		return ErrServiceState
 	}
-	return c.remoteClient.InvokeOneWay(addr, request, timeoutMillis)
+	return c.remoteClient.InvokeOneWay(ctx, addr, request, timeoutMillis)
 }
 
 func (c *rmqClient) CheckClientInBroker() {
@@ -357,7 +357,7 @@ func (c *rmqClient) SendHeartbeatToAllBrokerWithLock() {
 		data := value.(*BrokerData)
 		for id, addr := range data.BrokerAddresses {
 			cmd := remote.NewRemotingCommand(ReqHeartBeat, nil, hbData.encode())
-			response, err := c.remoteClient.InvokeSync(addr, cmd, 3*time.Second)
+			response, err := c.remoteClient.InvokeSync(context.Background(), addr, cmd, 3*time.Second)
 			if err != nil {
 				rlog.Warnf("send heart beat to broker error: %s", err.Error())
 				return true
@@ -417,7 +417,7 @@ func (c *rmqClient) SendMessageAsync(ctx context.Context, brokerAddrs, brokerNam
 func (c *rmqClient) SendMessageOneWay(ctx context.Context, brokerAddrs string, request *SendMessageRequest,
 	msgs []*primitive.Message) (*primitive.SendResult, error) {
 	cmd := remote.NewRemotingCommand(ReqSendBatchMessage, request, encodeMessages(msgs))
-	err := c.remoteClient.InvokeOneWay(brokerAddrs, cmd, 3*time.Second)
+	err := c.remoteClient.InvokeOneWay(ctx, brokerAddrs, cmd, 3*time.Second)
 	if err != nil {
 		rlog.Warnf("send messages with oneway error: %v", err)
 	}
@@ -473,7 +473,7 @@ func (c *rmqClient) ProcessSendResponse(brokerName string, cmd *remote.RemotingC
 // PullMessage with sync
 func (c *rmqClient) PullMessage(ctx context.Context, brokerAddrs string, request *PullMessageRequest) (*primitive.PullResult, error) {
 	cmd := remote.NewRemotingCommand(ReqPullMessage, request, nil)
-	res, err := c.remoteClient.InvokeSync(brokerAddrs, cmd, 10*time.Second)
+	res, err := c.remoteClient.InvokeSync(ctx, brokerAddrs, cmd, 10*time.Second)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/mock_client.go b/internal/mock_client.go
index d19f0e9..a216e93 100644
--- a/internal/mock_client.go
+++ b/internal/mock_client.go
@@ -1,3 +1,19 @@
+/*
+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.
+*/
 // Code generated by MockGen. DO NOT EDIT.
 // Source: client.go
 
@@ -5,13 +21,14 @@
 package internal
 
 import (
-	context "context"
-	reflect "reflect"
-	time "time"
+	"context"
+	"reflect"
+	"time"
 
-	remote "github.com/apache/rocketmq-client-go/internal/remote"
-	primitive "github.com/apache/rocketmq-client-go/primitive"
-	gomock "github.com/golang/mock/gomock"
+	"github.com/golang/mock/gomock"
+
+	"github.com/apache/rocketmq-client-go/internal/remote"
+	"github.com/apache/rocketmq-client-go/primitive"
 )
 
 // MockInnerProducer is a mock of InnerProducer interface
@@ -39,7 +56,6 @@ func (m *MockInnerProducer) EXPECT() *MockInnerProducerMockRecorder {
 
 // PublishTopicList mocks base method
 func (m *MockInnerProducer) PublishTopicList() []string {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "PublishTopicList")
 	ret0, _ := ret[0].([]string)
 	return ret0
@@ -47,25 +63,21 @@ func (m *MockInnerProducer) PublishTopicList() []string {
 
 // PublishTopicList indicates an expected call of PublishTopicList
 func (mr *MockInnerProducerMockRecorder) PublishTopicList() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PublishTopicList", reflect.TypeOf((*MockInnerProducer)(nil).PublishTopicList))
 }
 
 // UpdateTopicPublishInfo mocks base method
 func (m *MockInnerProducer) UpdateTopicPublishInfo(topic string, info *TopicPublishInfo) {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "UpdateTopicPublishInfo", topic, info)
 }
 
 // UpdateTopicPublishInfo indicates an expected call of UpdateTopicPublishInfo
 func (mr *MockInnerProducerMockRecorder) UpdateTopicPublishInfo(topic, info interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTopicPublishInfo", reflect.TypeOf((*MockInnerProducer)(nil).UpdateTopicPublishInfo), topic, info)
 }
 
 // IsPublishTopicNeedUpdate mocks base method
 func (m *MockInnerProducer) IsPublishTopicNeedUpdate(topic string) bool {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "IsPublishTopicNeedUpdate", topic)
 	ret0, _ := ret[0].(bool)
 	return ret0
@@ -73,13 +85,11 @@ func (m *MockInnerProducer) IsPublishTopicNeedUpdate(topic string) bool {
 
 // IsPublishTopicNeedUpdate indicates an expected call of IsPublishTopicNeedUpdate
 func (mr *MockInnerProducerMockRecorder) IsPublishTopicNeedUpdate(topic interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsPublishTopicNeedUpdate", reflect.TypeOf((*MockInnerProducer)(nil).IsPublishTopicNeedUpdate), topic)
 }
 
 // IsUnitMode mocks base method
 func (m *MockInnerProducer) IsUnitMode() bool {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "IsUnitMode")
 	ret0, _ := ret[0].(bool)
 	return ret0
@@ -87,7 +97,6 @@ func (m *MockInnerProducer) IsUnitMode() bool {
 
 // IsUnitMode indicates an expected call of IsUnitMode
 func (mr *MockInnerProducerMockRecorder) IsUnitMode() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsUnitMode", reflect.TypeOf((*MockInnerProducer)(nil).IsUnitMode))
 }
 
@@ -116,7 +125,6 @@ func (m *MockInnerConsumer) EXPECT() *MockInnerConsumerMockRecorder {
 
 // PersistConsumerOffset mocks base method
 func (m *MockInnerConsumer) PersistConsumerOffset() error {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "PersistConsumerOffset")
 	ret0, _ := ret[0].(error)
 	return ret0
@@ -124,25 +132,21 @@ func (m *MockInnerConsumer) PersistConsumerOffset() error {
 
 // PersistConsumerOffset indicates an expected call of PersistConsumerOffset
 func (mr *MockInnerConsumerMockRecorder) PersistConsumerOffset() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PersistConsumerOffset", reflect.TypeOf((*MockInnerConsumer)(nil).PersistConsumerOffset))
 }
 
 // UpdateTopicSubscribeInfo mocks base method
 func (m *MockInnerConsumer) UpdateTopicSubscribeInfo(topic string, mqs []*primitive.MessageQueue) {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "UpdateTopicSubscribeInfo", topic, mqs)
 }
 
 // UpdateTopicSubscribeInfo indicates an expected call of UpdateTopicSubscribeInfo
 func (mr *MockInnerConsumerMockRecorder) UpdateTopicSubscribeInfo(topic, mqs interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTopicSubscribeInfo", reflect.TypeOf((*MockInnerConsumer)(nil).UpdateTopicSubscribeInfo), topic, mqs)
 }
 
 // IsSubscribeTopicNeedUpdate mocks base method
 func (m *MockInnerConsumer) IsSubscribeTopicNeedUpdate(topic string) bool {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "IsSubscribeTopicNeedUpdate", topic)
 	ret0, _ := ret[0].(bool)
 	return ret0
@@ -150,13 +154,11 @@ func (m *MockInnerConsumer) IsSubscribeTopicNeedUpdate(topic string) bool {
 
 // IsSubscribeTopicNeedUpdate indicates an expected call of IsSubscribeTopicNeedUpdate
 func (mr *MockInnerConsumerMockRecorder) IsSubscribeTopicNeedUpdate(topic interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsSubscribeTopicNeedUpdate", reflect.TypeOf((*MockInnerConsumer)(nil).IsSubscribeTopicNeedUpdate), topic)
 }
 
 // SubscriptionDataList mocks base method
 func (m *MockInnerConsumer) SubscriptionDataList() []*SubscriptionData {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "SubscriptionDataList")
 	ret0, _ := ret[0].([]*SubscriptionData)
 	return ret0
@@ -164,25 +166,21 @@ func (m *MockInnerConsumer) SubscriptionDataList() []*SubscriptionData {
 
 // SubscriptionDataList indicates an expected call of SubscriptionDataList
 func (mr *MockInnerConsumerMockRecorder) SubscriptionDataList() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SubscriptionDataList", reflect.TypeOf((*MockInnerConsumer)(nil).SubscriptionDataList))
 }
 
 // Rebalance mocks base method
 func (m *MockInnerConsumer) Rebalance() {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "Rebalance")
 }
 
 // Rebalance indicates an expected call of Rebalance
 func (mr *MockInnerConsumerMockRecorder) Rebalance() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Rebalance", reflect.TypeOf((*MockInnerConsumer)(nil).Rebalance))
 }
 
 // IsUnitMode mocks base method
 func (m *MockInnerConsumer) IsUnitMode() bool {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "IsUnitMode")
 	ret0, _ := ret[0].(bool)
 	return ret0
@@ -190,7 +188,6 @@ func (m *MockInnerConsumer) IsUnitMode() bool {
 
 // IsUnitMode indicates an expected call of IsUnitMode
 func (mr *MockInnerConsumerMockRecorder) IsUnitMode() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsUnitMode", reflect.TypeOf((*MockInnerConsumer)(nil).IsUnitMode))
 }
 
@@ -219,31 +216,26 @@ func (m *MockRMQClient) EXPECT() *MockRMQClientMockRecorder {
 
 // Start mocks base method
 func (m *MockRMQClient) Start() {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "Start")
 }
 
 // Start indicates an expected call of Start
 func (mr *MockRMQClientMockRecorder) Start() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Start", reflect.TypeOf((*MockRMQClient)(nil).Start))
 }
 
 // Shutdown mocks base method
 func (m *MockRMQClient) Shutdown() {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "Shutdown")
 }
 
 // Shutdown indicates an expected call of Shutdown
 func (mr *MockRMQClientMockRecorder) Shutdown() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Shutdown", reflect.TypeOf((*MockRMQClient)(nil).Shutdown))
 }
 
 // ClientID mocks base method
 func (m *MockRMQClient) ClientID() string {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "ClientID")
 	ret0, _ := ret[0].(string)
 	return ret0
@@ -251,104 +243,88 @@ func (m *MockRMQClient) ClientID() string {
 
 // ClientID indicates an expected call of ClientID
 func (mr *MockRMQClientMockRecorder) ClientID() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ClientID", reflect.TypeOf((*MockRMQClient)(nil).ClientID))
 }
 
 // RegisterProducer mocks base method
 func (m *MockRMQClient) RegisterProducer(group string, producer InnerProducer) {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "RegisterProducer", group, producer)
 }
 
 // RegisterProducer indicates an expected call of RegisterProducer
 func (mr *MockRMQClientMockRecorder) RegisterProducer(group, producer interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterProducer", reflect.TypeOf((*MockRMQClient)(nil).RegisterProducer), group, producer)
 }
 
 // InvokeSync mocks base method
-func (m *MockRMQClient) InvokeSync(addr string, request *remote.RemotingCommand, timeoutMillis time.Duration) (*remote.RemotingCommand, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "InvokeSync", addr, request, timeoutMillis)
+func (m *MockRMQClient) InvokeSync(ctx context.Context, addr string, request *remote.RemotingCommand, timeoutMillis time.Duration) (*remote.RemotingCommand, error) {
+	ret := m.ctrl.Call(m, "InvokeSync", ctx, addr, request, timeoutMillis)
 	ret0, _ := ret[0].(*remote.RemotingCommand)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // InvokeSync indicates an expected call of InvokeSync
-func (mr *MockRMQClientMockRecorder) InvokeSync(addr, request, timeoutMillis interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeSync", reflect.TypeOf((*MockRMQClient)(nil).InvokeSync), addr, request, timeoutMillis)
+func (mr *MockRMQClientMockRecorder) InvokeSync(ctx, addr, request, timeoutMillis interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeSync", reflect.TypeOf((*MockRMQClient)(nil).InvokeSync), ctx, addr, request, timeoutMillis)
 }
 
 // InvokeAsync mocks base method
-func (m *MockRMQClient) InvokeAsync(addr string, request *remote.RemotingCommand, timeoutMillis time.Duration, f func(*remote.RemotingCommand, error)) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "InvokeAsync", addr, request, timeoutMillis, f)
+func (m *MockRMQClient) InvokeAsync(ctx context.Context, addr string, request *remote.RemotingCommand, timeoutMillis time.Duration, f func(*remote.RemotingCommand, error)) error {
+	ret := m.ctrl.Call(m, "InvokeAsync", ctx, addr, request, timeoutMillis, f)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // InvokeAsync indicates an expected call of InvokeAsync
-func (mr *MockRMQClientMockRecorder) InvokeAsync(addr, request, timeoutMillis, f interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeAsync", reflect.TypeOf((*MockRMQClient)(nil).InvokeAsync), addr, request, timeoutMillis, f)
+func (mr *MockRMQClientMockRecorder) InvokeAsync(ctx, addr, request, timeoutMillis, f interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeAsync", reflect.TypeOf((*MockRMQClient)(nil).InvokeAsync), ctx, addr, request, timeoutMillis, f)
 }
 
 // InvokeOneWay mocks base method
-func (m *MockRMQClient) InvokeOneWay(addr string, request *remote.RemotingCommand, timeoutMillis time.Duration) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "InvokeOneWay", addr, request, timeoutMillis)
+func (m *MockRMQClient) InvokeOneWay(ctx context.Context, addr string, request *remote.RemotingCommand, timeoutMillis time.Duration) error {
+	ret := m.ctrl.Call(m, "InvokeOneWay", ctx, addr, request, timeoutMillis)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // InvokeOneWay indicates an expected call of InvokeOneWay
-func (mr *MockRMQClientMockRecorder) InvokeOneWay(addr, request, timeoutMillis interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeOneWay", reflect.TypeOf((*MockRMQClient)(nil).InvokeOneWay), addr, request, timeoutMillis)
+func (mr *MockRMQClientMockRecorder) InvokeOneWay(ctx, addr, request, timeoutMillis interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeOneWay", reflect.TypeOf((*MockRMQClient)(nil).InvokeOneWay), ctx, addr, request, timeoutMillis)
 }
 
 // CheckClientInBroker mocks base method
 func (m *MockRMQClient) CheckClientInBroker() {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "CheckClientInBroker")
 }
 
 // CheckClientInBroker indicates an expected call of CheckClientInBroker
 func (mr *MockRMQClientMockRecorder) CheckClientInBroker() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckClientInBroker", reflect.TypeOf((*MockRMQClient)(nil).CheckClientInBroker))
 }
 
 // SendHeartbeatToAllBrokerWithLock mocks base method
 func (m *MockRMQClient) SendHeartbeatToAllBrokerWithLock() {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "SendHeartbeatToAllBrokerWithLock")
 }
 
 // SendHeartbeatToAllBrokerWithLock indicates an expected call of SendHeartbeatToAllBrokerWithLock
 func (mr *MockRMQClientMockRecorder) SendHeartbeatToAllBrokerWithLock() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendHeartbeatToAllBrokerWithLock", reflect.TypeOf((*MockRMQClient)(nil).SendHeartbeatToAllBrokerWithLock))
 }
 
 // UpdateTopicRouteInfo mocks base method
 func (m *MockRMQClient) UpdateTopicRouteInfo() {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "UpdateTopicRouteInfo")
 }
 
 // UpdateTopicRouteInfo indicates an expected call of UpdateTopicRouteInfo
 func (mr *MockRMQClientMockRecorder) UpdateTopicRouteInfo() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTopicRouteInfo", reflect.TypeOf((*MockRMQClient)(nil).UpdateTopicRouteInfo))
 }
 
 // ProcessSendResponse mocks base method
 func (m *MockRMQClient) ProcessSendResponse(brokerName string, cmd *remote.RemotingCommand, resp *primitive.SendResult, msgs ...*primitive.Message) error {
-	m.ctrl.T.Helper()
 	varargs := []interface{}{brokerName, cmd, resp}
 	for _, a := range msgs {
 		varargs = append(varargs, a)
@@ -360,14 +336,12 @@ func (m *MockRMQClient) ProcessSendResponse(brokerName string, cmd *remote.Remot
 
 // ProcessSendResponse indicates an expected call of ProcessSendResponse
 func (mr *MockRMQClientMockRecorder) ProcessSendResponse(brokerName, cmd, resp interface{}, msgs ...interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	varargs := append([]interface{}{brokerName, cmd, resp}, msgs...)
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ProcessSendResponse", reflect.TypeOf((*MockRMQClient)(nil).ProcessSendResponse), varargs...)
 }
 
 // RegisterConsumer mocks base method
 func (m *MockRMQClient) RegisterConsumer(group string, consumer InnerConsumer) error {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "RegisterConsumer", group, consumer)
 	ret0, _ := ret[0].(error)
 	return ret0
@@ -375,25 +349,21 @@ func (m *MockRMQClient) RegisterConsumer(group string, consumer InnerConsumer) e
 
 // RegisterConsumer indicates an expected call of RegisterConsumer
 func (mr *MockRMQClientMockRecorder) RegisterConsumer(group, consumer interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterConsumer", reflect.TypeOf((*MockRMQClient)(nil).RegisterConsumer), group, consumer)
 }
 
 // UnregisterConsumer mocks base method
 func (m *MockRMQClient) UnregisterConsumer(group string) {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "UnregisterConsumer", group)
 }
 
 // UnregisterConsumer indicates an expected call of UnregisterConsumer
 func (mr *MockRMQClientMockRecorder) UnregisterConsumer(group interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UnregisterConsumer", reflect.TypeOf((*MockRMQClient)(nil).UnregisterConsumer), group)
 }
 
 // PullMessage mocks base method
 func (m *MockRMQClient) PullMessage(ctx context.Context, brokerAddrs string, request *PullMessageRequest) (*primitive.PullResult, error) {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "PullMessage", ctx, brokerAddrs, request)
 	ret0, _ := ret[0].(*primitive.PullResult)
 	ret1, _ := ret[1].(error)
@@ -402,13 +372,11 @@ func (m *MockRMQClient) PullMessage(ctx context.Context, brokerAddrs string, req
 
 // PullMessage indicates an expected call of PullMessage
 func (mr *MockRMQClientMockRecorder) PullMessage(ctx, brokerAddrs, request interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PullMessage", reflect.TypeOf((*MockRMQClient)(nil).PullMessage), ctx, brokerAddrs, request)
 }
 
 // PullMessageAsync mocks base method
 func (m *MockRMQClient) PullMessageAsync(ctx context.Context, brokerAddrs string, request *PullMessageRequest, f func(*primitive.PullResult)) error {
-	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "PullMessageAsync", ctx, brokerAddrs, request, f)
 	ret0, _ := ret[0].(error)
 	return ret0
@@ -416,30 +384,25 @@ func (m *MockRMQClient) PullMessageAsync(ctx context.Context, brokerAddrs string
 
 // PullMessageAsync indicates an expected call of PullMessageAsync
 func (mr *MockRMQClientMockRecorder) PullMessageAsync(ctx, brokerAddrs, request, f interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PullMessageAsync", reflect.TypeOf((*MockRMQClient)(nil).PullMessageAsync), ctx, brokerAddrs, request, f)
 }
 
 // RebalanceImmediately mocks base method
 func (m *MockRMQClient) RebalanceImmediately() {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "RebalanceImmediately")
 }
 
 // RebalanceImmediately indicates an expected call of RebalanceImmediately
 func (mr *MockRMQClientMockRecorder) RebalanceImmediately() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RebalanceImmediately", reflect.TypeOf((*MockRMQClient)(nil).RebalanceImmediately))
 }
 
 // UpdatePublishInfo mocks base method
 func (m *MockRMQClient) UpdatePublishInfo(topic string, data *TopicRouteData) {
-	m.ctrl.T.Helper()
 	m.ctrl.Call(m, "UpdatePublishInfo", topic, data)
 }
 
 // UpdatePublishInfo indicates an expected call of UpdatePublishInfo
 func (mr *MockRMQClientMockRecorder) UpdatePublishInfo(topic, data interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdatePublishInfo", reflect.TypeOf((*MockRMQClient)(nil).UpdatePublishInfo), topic, data)
 }
diff --git a/internal/mock_namesrv.go b/internal/mock_namesrv.go
index 6333e48..630d4cc 100644
--- a/internal/mock_namesrv.go
+++ b/internal/mock_namesrv.go
@@ -1,3 +1,20 @@
+/*
+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.
+*/
+
 // Code generated by MockGen. DO NOT EDIT.
 // Source: namesrv.go
 
diff --git a/internal/namesrv.go b/internal/namesrv.go
index 3fc84dd..3c20968 100644
--- a/internal/namesrv.go
+++ b/internal/namesrv.go
@@ -59,10 +59,10 @@ type namesrvs struct {
 	// namesrv addr list
 	srvs []string
 
-	// lock for getNamesrv in case of update index race condition
+	// lock for getNameServerAddress in case of update index race condition
 	lock sync.Locker
 
-	// index indicate the next position for getNamesrv
+	// index indicate the next position for getNameServerAddress
 	index int
 
 	// brokerName -> *BrokerData
@@ -76,7 +76,7 @@ type namesrvs struct {
 
 	lockNamesrv sync.Mutex
 
-	nameSrvClient *remote.RemotingClient
+	nameSrvClient remote.RemotingClient
 }
 
 var _ Namesrvs = &namesrvs{}
@@ -94,8 +94,8 @@ func NewNamesrv(addr primitive.NamesrvAddr) (*namesrvs, error) {
 	}, nil
 }
 
-// getNamesrv return namesrv using round-robin strategy.
-func (s *namesrvs) getNamesrv() string {
+// getNameServerAddress return namesrv using round-robin strategy.
+func (s *namesrvs) getNameServerAddress() string {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
@@ -119,25 +119,3 @@ func (s *namesrvs) String() string {
 func (s *namesrvs) SetCredentials(credentials primitive.Credentials) {
 	s.nameSrvClient.RegisterInterceptor(remote.ACLInterceptor(credentials))
 }
-
-var (
-	httpPrefixRegex, _ = regexp.Compile("^(http|https)://")
-)
-
-func verifyIP(ip string) error {
-	if httpPrefixRegex.MatchString(ip) {
-		return nil
-	}
-	if strings.Contains(ip, ";") {
-		return ErrMultiIP
-	}
-	ips := ipRegex.FindAllString(ip, -1)
-	if len(ips) == 0 {
-		return ErrIllegalIP
-	}
-
-	if len(ips) > 1 {
-		return ErrMultiIP
-	}
-	return nil
-}
diff --git a/internal/remote/future.go b/internal/remote/future.go
index 8690644..4d3008f 100644
--- a/internal/remote/future.go
+++ b/internal/remote/future.go
@@ -18,6 +18,7 @@ limitations under the License.
 package remote
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -35,16 +36,18 @@ type ResponseFuture struct {
 	BeginTimestamp  time.Duration
 	Done            chan bool
 	callbackOnce    sync.Once
+	ctx             context.Context
 }
 
 // NewResponseFuture create ResponseFuture with opaque, timeout and callback
-func NewResponseFuture(opaque int32, timeout time.Duration, callback func(*ResponseFuture)) *ResponseFuture {
+func NewResponseFuture(ctx context.Context, opaque int32, timeout time.Duration, callback func(*ResponseFuture)) *ResponseFuture {
 	return &ResponseFuture{
 		Opaque:         opaque,
 		Done:           make(chan bool),
 		Timeout:        timeout,
 		callback:       callback,
 		BeginTimestamp: time.Duration(time.Now().Unix()) * time.Second,
+		ctx:            ctx,
 	}
 }
 
@@ -66,19 +69,19 @@ func (r *ResponseFuture) waitResponse() (*RemotingCommand, error) {
 		cmd *RemotingCommand
 		err error
 	)
-	timer := time.NewTimer(r.Timeout)
+	ctx, cancel := context.WithTimeout(r.ctx, r.Timeout)
+	defer cancel()
 	for {
 		select {
 		case <-r.Done:
 			cmd, err = r.ResponseCommand, r.Err
 			goto done
-		case <-timer.C:
+		case <-ctx.Done():
 			err = utils.ErrRequestTimeout
 			r.Err = err
 			goto done
 		}
 	}
 done:
-	timer.Stop()
 	return cmd, err
 }
diff --git a/internal/remote/mock_remote_client.go b/internal/remote/mock_remote_client.go
new file mode 100644
index 0000000..e264ab2
--- /dev/null
+++ b/internal/remote/mock_remote_client.go
@@ -0,0 +1,123 @@
+/*
+ * 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.
+ */
+// Code generated by MockGen. DO NOT EDIT.
+// Source: remote_client.go
+
+// Package remote is a generated GoMock package.
+package remote
+
+import (
+	context "context"
+	primitive "github.com/apache/rocketmq-client-go/primitive"
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+	time "time"
+)
+
+// MockRemotingClient is a mock of RemotingClient interface
+type MockRemotingClient struct {
+	ctrl     *gomock.Controller
+	recorder *MockRemotingClientMockRecorder
+}
+
+// MockRemotingClientMockRecorder is the mock recorder for MockRemotingClient
+type MockRemotingClientMockRecorder struct {
+	mock *MockRemotingClient
+}
+
+// NewMockRemotingClient creates a new mock instance
+func NewMockRemotingClient(ctrl *gomock.Controller) *MockRemotingClient {
+	mock := &MockRemotingClient{ctrl: ctrl}
+	mock.recorder = &MockRemotingClientMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockRemotingClient) EXPECT() *MockRemotingClientMockRecorder {
+	return m.recorder
+}
+
+// RegisterRequestFunc mocks base method
+func (m *MockRemotingClient) RegisterRequestFunc(code int16, f ClientRequestFunc) {
+	m.ctrl.Call(m, "RegisterRequestFunc", code, f)
+}
+
+// RegisterRequestFunc indicates an expected call of RegisterRequestFunc
+func (mr *MockRemotingClientMockRecorder) RegisterRequestFunc(code, f interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterRequestFunc", reflect.TypeOf((*MockRemotingClient)(nil).RegisterRequestFunc), code, f)
+}
+
+// RegisterInterceptor mocks base method
+func (m *MockRemotingClient) RegisterInterceptor(interceptors ...primitive.Interceptor) {
+	varargs := []interface{}{}
+	for _, a := range interceptors {
+		varargs = append(varargs, a)
+	}
+	m.ctrl.Call(m, "RegisterInterceptor", varargs...)
+}
+
+// RegisterInterceptor indicates an expected call of RegisterInterceptor
+func (mr *MockRemotingClientMockRecorder) RegisterInterceptor(interceptors ...interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterInterceptor", reflect.TypeOf((*MockRemotingClient)(nil).RegisterInterceptor), interceptors...)
+}
+
+// InvokeSync mocks base method
+func (m *MockRemotingClient) InvokeSync(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration) (*RemotingCommand, error) {
+	ret := m.ctrl.Call(m, "InvokeSync", ctx, addr, request, timeout)
+	ret0, _ := ret[0].(*RemotingCommand)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// InvokeSync indicates an expected call of InvokeSync
+func (mr *MockRemotingClientMockRecorder) InvokeSync(ctx, addr, request, timeout interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeSync", reflect.TypeOf((*MockRemotingClient)(nil).InvokeSync), ctx, addr, request, timeout)
+}
+
+// InvokeAsync mocks base method
+func (m *MockRemotingClient) InvokeAsync(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration, callback func(*ResponseFuture)) error {
+	ret := m.ctrl.Call(m, "InvokeAsync", ctx, addr, request, timeout, callback)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// InvokeAsync indicates an expected call of InvokeAsync
+func (mr *MockRemotingClientMockRecorder) InvokeAsync(ctx, addr, request, timeout, callback interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeAsync", reflect.TypeOf((*MockRemotingClient)(nil).InvokeAsync), ctx, addr, request, timeout, callback)
+}
+
+// InvokeOneWay mocks base method
+func (m *MockRemotingClient) InvokeOneWay(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration) error {
+	ret := m.ctrl.Call(m, "InvokeOneWay", ctx, addr, request, timeout)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// InvokeOneWay indicates an expected call of InvokeOneWay
+func (mr *MockRemotingClientMockRecorder) InvokeOneWay(ctx, addr, request, timeout interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeOneWay", reflect.TypeOf((*MockRemotingClient)(nil).InvokeOneWay), ctx, addr, request, timeout)
+}
+
+// ShutDown mocks base method
+func (m *MockRemotingClient) ShutDown() {
+	m.ctrl.Call(m, "ShutDown")
+}
+
+// ShutDown indicates an expected call of ShutDown
+func (mr *MockRemotingClientMockRecorder) ShutDown() *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ShutDown", reflect.TypeOf((*MockRemotingClient)(nil).ShutDown))
+}
diff --git a/internal/remote/remote_client.go b/internal/remote/remote_client.go
index 4159808..96f67a6 100644
--- a/internal/remote/remote_client.go
+++ b/internal/remote/remote_client.go
@@ -37,7 +37,19 @@ type TcpOption struct {
 	// TODO
 }
 
-type RemotingClient struct {
+//go:generate mockgen -source remote_client.go -destination mock_remote_client.go -self_package github.com/apache/rocketmq-client-go/internal/remote  --package remote RemotingClient
+type RemotingClient interface {
+	RegisterRequestFunc(code int16, f ClientRequestFunc)
+	RegisterInterceptor(interceptors ...primitive.Interceptor)
+	InvokeSync(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration) (*RemotingCommand, error)
+	InvokeAsync(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration, callback func(*ResponseFuture)) error
+	InvokeOneWay(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration) error
+	ShutDown()
+}
+
+var _ RemotingClient = &remotingClient{}
+
+type remotingClient struct {
 	responseTable    sync.Map
 	connectionTable  sync.Map
 	option           TcpOption
@@ -46,23 +58,23 @@ type RemotingClient struct {
 	interceptor      primitive.Interceptor
 }
 
-func NewRemotingClient() *RemotingClient {
-	return &RemotingClient{
+func NewRemotingClient() *remotingClient {
+	return &remotingClient{
 		processors: make(map[int16]ClientRequestFunc),
 	}
 }
 
-func (c *RemotingClient) RegisterRequestFunc(code int16, f ClientRequestFunc) {
+func (c *remotingClient) RegisterRequestFunc(code int16, f ClientRequestFunc) {
 	c.processors[code] = f
 }
 
 // TODO: merge sync and async model. sync should run on async model by blocking on chan
-func (c *RemotingClient) InvokeSync(addr string, request *RemotingCommand, timeout time.Duration) (*RemotingCommand, error) {
-	conn, err := c.connect(addr)
+func (c *remotingClient) InvokeSync(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration) (*RemotingCommand, error) {
+	conn, err := c.connect(ctx, addr)
 	if err != nil {
 		return nil, err
 	}
-	resp := NewResponseFuture(request.Opaque, timeout, nil)
+	resp := NewResponseFuture(ctx, request.Opaque, timeout, nil)
 	c.responseTable.Store(resp.Opaque, resp)
 	defer c.responseTable.Delete(request.Opaque)
 	err = c.sendRequest(conn, request)
@@ -74,12 +86,12 @@ func (c *RemotingClient) InvokeSync(addr string, request *RemotingCommand, timeo
 }
 
 // InvokeAsync send request without blocking, just return immediately.
-func (c *RemotingClient) InvokeAsync(addr string, request *RemotingCommand, timeout time.Duration, callback func(*ResponseFuture)) error {
-	conn, err := c.connect(addr)
+func (c *remotingClient) InvokeAsync(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration, callback func(*ResponseFuture)) error {
+	conn, err := c.connect(ctx, addr)
 	if err != nil {
 		return err
 	}
-	resp := NewResponseFuture(request.Opaque, timeout, callback)
+	resp := NewResponseFuture(ctx, request.Opaque, timeout, callback)
 	c.responseTable.Store(resp.Opaque, resp)
 	err = c.sendRequest(conn, request)
 	if err != nil {
@@ -90,22 +102,22 @@ func (c *RemotingClient) InvokeAsync(addr string, request *RemotingCommand, time
 	return nil
 }
 
-func (c *RemotingClient) receiveAsync(f *ResponseFuture) {
+func (c *remotingClient) receiveAsync(f *ResponseFuture) {
 	_, err := f.waitResponse()
 	if err != nil {
 		f.executeInvokeCallback()
 	}
 }
 
-func (c *RemotingClient) InvokeOneWay(addr string, request *RemotingCommand, timeout time.Duration) error {
-	conn, err := c.connect(addr)
+func (c *remotingClient) InvokeOneWay(ctx context.Context, addr string, request *RemotingCommand, timeout time.Duration) error {
+	conn, err := c.connect(ctx, addr)
 	if err != nil {
 		return err
 	}
 	return c.sendRequest(conn, request)
 }
 
-func (c *RemotingClient) connect(addr string) (net.Conn, error) {
+func (c *remotingClient) connect(ctx context.Context, addr string) (net.Conn, error) {
 	//it needs additional locker.
 	c.connectionLocker.Lock()
 	defer c.connectionLocker.Unlock()
@@ -113,7 +125,8 @@ func (c *RemotingClient) connect(addr string) (net.Conn, error) {
 	if ok {
 		return conn.(net.Conn), nil
 	}
-	tcpConn, err := net.Dial("tcp", addr)
+	var d net.Dialer
+	tcpConn, err := d.DialContext(ctx, "tcp", addr)
 	if err != nil {
 		return nil, err
 	}
@@ -122,7 +135,7 @@ func (c *RemotingClient) connect(addr string) (net.Conn, error) {
 	return tcpConn, nil
 }
 
-func (c *RemotingClient) receiveResponse(r net.Conn) {
+func (c *remotingClient) receiveResponse(r net.Conn) {
 	scanner := c.createScanner(r)
 	for scanner.Scan() {
 		cmd, err := decode(scanner.Bytes())
@@ -170,7 +183,7 @@ func (c *RemotingClient) receiveResponse(r net.Conn) {
 	}
 }
 
-func (c *RemotingClient) createScanner(r io.Reader) *bufio.Scanner {
+func (c *remotingClient) createScanner(r io.Reader) *bufio.Scanner {
 	scanner := bufio.NewScanner(r)
 
 	// max batch size: 32, max message size: 4Mb
@@ -200,7 +213,7 @@ func (c *RemotingClient) createScanner(r io.Reader) *bufio.Scanner {
 	return scanner
 }
 
-func (c *RemotingClient) sendRequest(conn net.Conn, request *RemotingCommand) error {
+func (c *remotingClient) sendRequest(conn net.Conn, request *RemotingCommand) error {
 	var err error
 	if c.interceptor != nil {
 		err = c.interceptor(context.Background(), request, nil, func(ctx context.Context, req, reply interface{}) error {
@@ -212,7 +225,7 @@ func (c *RemotingClient) sendRequest(conn net.Conn, request *RemotingCommand) er
 	return err
 }
 
-func (c *RemotingClient) doRequest(conn net.Conn, request *RemotingCommand) error {
+func (c *remotingClient) doRequest(conn net.Conn, request *RemotingCommand) error {
 	content, err := encode(request)
 	if err != nil {
 		return err
@@ -225,7 +238,7 @@ func (c *RemotingClient) doRequest(conn net.Conn, request *RemotingCommand) erro
 	return nil
 }
 
-func (c *RemotingClient) closeConnection(toCloseConn net.Conn) {
+func (c *remotingClient) closeConnection(toCloseConn net.Conn) {
 	c.connectionTable.Range(func(key, value interface{}) bool {
 		if value == toCloseConn {
 			c.connectionTable.Delete(key)
@@ -236,7 +249,7 @@ func (c *RemotingClient) closeConnection(toCloseConn net.Conn) {
 	})
 }
 
-func (c *RemotingClient) ShutDown() {
+func (c *remotingClient) ShutDown() {
 	c.responseTable.Range(func(key, value interface{}) bool {
 		c.responseTable.Delete(key)
 		return true
@@ -248,7 +261,7 @@ func (c *RemotingClient) ShutDown() {
 	})
 }
 
-func (c *RemotingClient) RegisterInterceptor(interceptors ...primitive.Interceptor) {
+func (c *remotingClient) RegisterInterceptor(interceptors ...primitive.Interceptor) {
 
 	c.interceptor = primitive.ChainInterceptors(interceptors...)
 
diff --git a/internal/route.go b/internal/route.go
index 8ed00b0..3611bb9 100644
--- a/internal/route.go
+++ b/internal/route.go
@@ -18,6 +18,7 @@ limitations under the License.
 package internal
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"math/rand"
@@ -304,10 +305,10 @@ func (s *namesrvs) queryTopicRouteInfoFromServer(topic string) (*TopicRouteData,
 	)
 	for i := 0; i < s.Size(); i++ {
 		rc := remote.NewRemotingCommand(ReqGetRouteInfoByTopic, request, nil)
-		response, err = s.nameSrvClient.InvokeSync(s.getNameServerAddress(), rc, requestTimeout)
+		response, err = s.nameSrvClient.InvokeSync(context.Background(), s.getNameServerAddress(), rc, requestTimeout)
 
-		if err != nil {
-			continue
+		if err == nil {
+			break
 		}
 	}
 	if err != nil {
@@ -418,10 +419,6 @@ func (s *namesrvs) routeData2PublishInfo(topic string, data *TopicRouteData) *To
 	return publishInfo
 }
 
-func (s *namesrvs) getNameServerAddress() string {
-	return s.getNamesrv()
-}
-
 // TopicRouteData TopicRouteData
 type TopicRouteData struct {
 	OrderTopicConf string
diff --git a/internal/trace.go b/internal/trace.go
index ca60eea..8bc6748 100644
--- a/internal/trace.go
+++ b/internal/trace.go
@@ -397,7 +397,7 @@ func (td *traceDispatcher) sendTraceDataByMQ(keyset Keyset, regionID string, dat
 	}
 
 	var req = td.buildSendRequest(mq, msg)
-	td.cli.InvokeAsync(addr, req, 5000*time.Millisecond, func(command *remote.RemotingCommand, e error) {
+	td.cli.InvokeAsync(context.Background(), addr, req, 5000*time.Millisecond, func(command *remote.RemotingCommand, e error) {
 		if e != nil {
 			rlog.Error("send trace data ,the traceData is %v", data)
 		}
diff --git a/internal/utils/net.go b/internal/utils/net.go
index 65f2731..c9444a9 100644
--- a/internal/utils/net.go
+++ b/internal/utils/net.go
@@ -1,3 +1,20 @@
+/*
+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 (
diff --git a/producer/producer.go b/producer/producer.go
index 0b0367f..25e3210 100644
--- a/producer/producer.go
+++ b/producer/producer.go
@@ -166,7 +166,7 @@ func (p *defaultProducer) sendSync(ctx context.Context, msg *primitive.Message,
 			producerCtx.MQ = *mq
 		}
 
-		res, _err := p.client.InvokeSync(addr, p.buildSendRequest(mq, msg), 3*time.Second)
+		res, _err := p.client.InvokeSync(ctx, addr, p.buildSendRequest(mq, msg), 3*time.Second)
 		if _err != nil {
 			err = _err
 			continue
@@ -205,7 +205,7 @@ func (p *defaultProducer) sendAsync(ctx context.Context, msg *primitive.Message,
 		return errors.Errorf("topic=%s route info not found", mq.Topic)
 	}
 
-	return p.client.InvokeAsync(addr, p.buildSendRequest(mq, msg), 3*time.Second, func(command *remote.RemotingCommand, err error) {
+	return p.client.InvokeAsync(ctx, addr, p.buildSendRequest(mq, msg), 3*time.Second, func(command *remote.RemotingCommand, err error) {
 		resp := new(primitive.SendResult)
 		if err != nil {
 			h(ctx, nil, err)
@@ -251,7 +251,7 @@ func (p *defaultProducer) sendOneWay(ctx context.Context, msg *primitive.Message
 			return fmt.Errorf("topic=%s route info not found", mq.Topic)
 		}
 
-		_err := p.client.InvokeOneWay(addr, p.buildSendRequest(mq, msg), 3*time.Second)
+		_err := p.client.InvokeOneWay(ctx, addr, p.buildSendRequest(mq, msg), 3*time.Second)
 		if _err != nil {
 			err = _err
 			continue
@@ -398,7 +398,7 @@ func (tp *transactionProducer) checkTransactionState() {
 			req := remote.NewRemotingCommand(internal.ReqENDTransaction, header, nil)
 			req.Remark = tp.errRemark(nil)
 
-			tp.producer.client.InvokeOneWay(callback.Addr.String(), req, tp.producer.options.SendMsgTimeout)
+			tp.producer.client.InvokeOneWay(context.Background(), callback.Addr.String(), req, tp.producer.options.SendMsgTimeout)
 		default:
 			rlog.Error("unknow type %v", ch)
 		}
@@ -467,7 +467,7 @@ func (tp *transactionProducer) endTransaction(result primitive.SendResult, err e
 	req := remote.NewRemotingCommand(internal.ReqENDTransaction, requestHeader, nil)
 	req.Remark = tp.errRemark(err)
 
-	return tp.producer.client.InvokeOneWay(brokerAddr, req, tp.producer.options.SendMsgTimeout)
+	return tp.producer.client.InvokeOneWay(context.Background(), brokerAddr, req, tp.producer.options.SendMsgTimeout)
 }
 
 func (tp *transactionProducer) errRemark(err error) string {


Test patch:
diff --git a/consumer/consumer_test.go b/consumer/consumer_test.go
index a837b24..665547d 100644
--- a/consumer/consumer_test.go
+++ b/consumer/consumer_test.go
@@ -18,10 +18,17 @@ limitations under the License.
 package consumer
 
 import (
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/golang/mock/gomock"
+	. "github.com/smartystreets/goconvey/convey"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/apache/rocketmq-client-go/internal"
+	"github.com/apache/rocketmq-client-go/internal/remote"
+	"github.com/apache/rocketmq-client-go/primitive"
 )
 
 func TestParseTimestamp(t *testing.T) {
@@ -30,3 +37,134 @@ func TestParseTimestamp(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, int64(1556652849), timestamp.Unix())
 }
+
+func TestDoRebalance(t *testing.T) {
+	Convey("Given a defaultConsumer", t, func() {
+		dc := &defaultConsumer{
+			model: Clustering,
+		}
+
+		topic := "test"
+		broker := "127.0.0.1:8889"
+		clientID := "clientID"
+		mqs := []*primitive.MessageQueue{
+			{
+				Topic:      topic,
+				BrokerName: "",
+				QueueId:    0,
+			},
+			{
+				Topic:      topic,
+				BrokerName: "",
+				QueueId:    1,
+			},
+		}
+		dc.topicSubscribeInfoTable.Store(topic, mqs)
+		sub := &internal.SubscriptionData{}
+		dc.subscriptionDataTable.Store(topic, sub)
+
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+		namesrvCli := internal.NewMockNamesrvs(ctrl)
+		namesrvCli.EXPECT().FindBrokerAddrByTopic(gomock.Any()).Return(broker)
+		dc.namesrv = namesrvCli
+
+		rmqCli := internal.NewMockRMQClient(ctrl)
+		rmqCli.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any()).
+			Return(&remote.RemotingCommand{
+				Body: []byte("{\"consumerIdList\": [\"a1\", \"a2\", \"a3\"] }"),
+			}, nil)
+		rmqCli.EXPECT().ClientID().Return(clientID)
+		dc.client = rmqCli
+
+		var wg sync.WaitGroup
+		wg.Add(1)
+		dc.allocate = func(cg string, clientID string, mqAll []*primitive.MessageQueue, cidAll []string) []*primitive.MessageQueue {
+			assert.Equal(t, cidAll, []string{"a1", "a2", "a3"})
+			wg.Done()
+			return nil
+		}
+
+		dc.doBalance()
+
+		wg.Wait()
+	})
+}
+
+func TestComputePullFromWhere(t *testing.T) {
+	Convey("Given a defaultConsumer", t, func() {
+		dc := &defaultConsumer{
+			model: Clustering,
+			cType: _PushConsume,
+		}
+
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+
+		offsetStore := NewMockOffsetStore(ctrl)
+		dc.storage = offsetStore
+
+		mq := &primitive.MessageQueue{
+			Topic: "test",
+		}
+
+		namesrvCli := internal.NewMockNamesrvs(ctrl)
+		dc.namesrv = namesrvCli
+
+		rmqCli := internal.NewMockRMQClient(ctrl)
+		dc.client = rmqCli
+
+		Convey("get effective offset", func() {
+			offsetStore.EXPECT().read(gomock.Any(), gomock.Any()).Return(int64(10))
+			res := dc.computePullFromWhere(mq)
+			assert.Equal(t, int64(10), res)
+		})
+
+		Convey("ConsumeFromLastOffset for normal topic", func() {
+			offsetStore.EXPECT().read(gomock.Any(), gomock.Any()).Return(int64(-1))
+			dc.option.FromWhere = ConsumeFromLastOffset
+
+			broker := "a"
+			namesrvCli.EXPECT().FindBrokerAddrByName(gomock.Any()).Return(broker)
+
+			rmqCli.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any()).
+				Return(&remote.RemotingCommand{
+					ExtFields: map[string]string{
+						"offset": "20",
+					},
+				}, nil)
+
+			res := dc.computePullFromWhere(mq)
+			assert.Equal(t, int64(20), res)
+		})
+
+		Convey("ConsumeFromFirstOffset for normal topic", func() {
+			offsetStore.EXPECT().read(gomock.Any(), gomock.Any()).Return(int64(-1))
+			dc.option.FromWhere = ConsumeFromFirstOffset
+
+			res := dc.computePullFromWhere(mq)
+			assert.Equal(t, int64(0), res)
+		})
+
+		Convey("ConsumeFromTimestamp for normal topic", func() {
+			offsetStore.EXPECT().read(gomock.Any(), gomock.Any()).Return(int64(-1))
+			dc.option.FromWhere = ConsumeFromTimestamp
+
+			dc.option.ConsumeTimestamp = "20060102150405"
+
+			broker := "a"
+			namesrvCli.EXPECT().FindBrokerAddrByName(gomock.Any()).Return(broker)
+
+			rmqCli.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any()).
+				Return(&remote.RemotingCommand{
+					ExtFields: map[string]string{
+						"offset": "30",
+					},
+				}, nil)
+
+			res := dc.computePullFromWhere(mq)
+			assert.Equal(t, int64(30), res)
+		})
+
+	})
+}
diff --git a/consumer/offset_store_test.go b/consumer/offset_store_test.go
index 93fbb59..4f6e5ef 100644
--- a/consumer/offset_store_test.go
+++ b/consumer/offset_store_test.go
@@ -1,3 +1,20 @@
+/*
+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 (
@@ -262,7 +279,7 @@ func TestRemoteBrokerOffsetStore(t *testing.T) {
 					"offset": "1",
 				},
 			}
-			rmqClient.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any()).Return(ret, nil)
+			rmqClient.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(ret, nil)
 
 			remoteStore.persist(queues)
 			offset := remoteStore.read(mq, _ReadFromStore)
diff --git a/consumer/statistics_test.go b/consumer/statistics_test.go
index 370940f..f70ed0b 100644
--- a/consumer/statistics_test.go
+++ b/consumer/statistics_test.go
@@ -1,3 +1,20 @@
+/*
+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 (
diff --git a/consumer/strategy_test.go b/consumer/strategy_test.go
index 5d38004..8b4f9e4 100644
--- a/consumer/strategy_test.go
+++ b/consumer/strategy_test.go
@@ -1,3 +1,20 @@
+/*
+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 (
diff --git a/internal/namesrv_test.go b/internal/namesrv_test.go
index 30fceb9..b6db2a0 100644
--- a/internal/namesrv_test.go
+++ b/internal/namesrv_test.go
@@ -31,15 +31,15 @@ func TestSelector(t *testing.T) {
 	namesrv, err := NewNamesrv(srvs)
 	assert.Nil(t, err)
 
-	assert.Equal(t, srvs[0], namesrv.getNamesrv())
-	assert.Equal(t, srvs[1], namesrv.getNamesrv())
-	assert.Equal(t, srvs[2], namesrv.getNamesrv())
-	assert.Equal(t, srvs[3], namesrv.getNamesrv())
-	assert.Equal(t, srvs[0], namesrv.getNamesrv())
-	assert.Equal(t, srvs[1], namesrv.getNamesrv())
-	assert.Equal(t, srvs[2], namesrv.getNamesrv())
-	assert.Equal(t, srvs[3], namesrv.getNamesrv())
-	assert.Equal(t, srvs[0], namesrv.getNamesrv())
+	assert.Equal(t, srvs[0], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[1], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[2], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[3], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[0], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[1], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[2], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[3], namesrv.getNameServerAddress())
+	assert.Equal(t, srvs[0], namesrv.getNameServerAddress())
 }
 
 func TestGetNamesrv(t *testing.T) {
@@ -55,10 +55,10 @@ func TestGetNamesrv(t *testing.T) {
 		}
 
 		index1 := ns.index
-		IP1 := ns.getNamesrv()
+		IP1 := ns.getNameServerAddress()
 
 		index2 := ns.index
-		IP2 := ns.getNamesrv()
+		IP2 := ns.getNameServerAddress()
 
 		So(index1+1, ShouldEqual, index2)
 		So(IP1, ShouldEqual, ns.srvs[index1])
diff --git a/internal/remote/remote_client_test.go b/internal/remote/remote_client_test.go
index 9cbb117..1d2c033 100644
--- a/internal/remote/remote_client_test.go
+++ b/internal/remote/remote_client_test.go
@@ -18,6 +18,7 @@ package remote
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"math/rand"
 	"net"
@@ -32,7 +33,7 @@ import (
 )
 
 func TestNewResponseFuture(t *testing.T) {
-	future := NewResponseFuture(10, time.Duration(1000), nil)
+	future := NewResponseFuture(context.Background(), 10, time.Duration(1000), nil)
 	if future.Opaque != 10 {
 		t.Errorf("wrong ResponseFuture's opaque. want=%d, got=%d", 10, future.Opaque)
 	}
@@ -62,7 +63,7 @@ func TestResponseFutureTimeout(t *testing.T) {
 			r.ResponseCommand.Remark = r.ResponseCommand.Remark + "Go Client"
 		}
 	}
-	future := NewResponseFuture(10, time.Duration(1000), callback)
+	future := NewResponseFuture(context.Background(), 10, time.Duration(1000), callback)
 	future.ResponseCommand = NewRemotingCommand(200,
 		nil, nil)
 
@@ -83,7 +84,7 @@ func TestResponseFutureTimeout(t *testing.T) {
 }
 
 func TestResponseFutureIsTimeout(t *testing.T) {
-	future := NewResponseFuture(10, 500*time.Millisecond, nil)
+	future := NewResponseFuture(context.Background(), 10, 500*time.Millisecond, nil)
 	if future.isTimeout() != false {
 		t.Errorf("wrong ResponseFuture's istimeout. want=%t, got=%t", false, future.isTimeout())
 	}
@@ -92,12 +93,12 @@ func TestResponseFutureIsTimeout(t *testing.T) {
 }
 
 func TestResponseFutureWaitResponse(t *testing.T) {
-	future := NewResponseFuture(10, 500*time.Millisecond, nil)
+	future := NewResponseFuture(context.Background(), 10, 500*time.Millisecond, nil)
 	if _, err := future.waitResponse(); err != utils.ErrRequestTimeout {
 		t.Errorf("wrong ResponseFuture waitResponse. want=%v, got=%v",
 			utils.ErrRequestTimeout, err)
 	}
-	future = NewResponseFuture(10, 500*time.Millisecond, nil)
+	future = NewResponseFuture(context.Background(), 10, 500*time.Millisecond, nil)
 	responseError := errors.New("response error")
 	go func() {
 		time.Sleep(100 * time.Millisecond)
@@ -108,7 +109,7 @@ func TestResponseFutureWaitResponse(t *testing.T) {
 		t.Errorf("wrong ResponseFuture waitResponse. want=%v. got=%v",
 			responseError, err)
 	}
-	future = NewResponseFuture(10, 500*time.Millisecond, nil)
+	future = NewResponseFuture(context.Background(), 10, 500*time.Millisecond, nil)
 	responseRemotingCommand := NewRemotingCommand(202, nil, nil)
 	go func() {
 		time.Sleep(100 * time.Millisecond)
@@ -173,7 +174,7 @@ func TestInvokeSync(t *testing.T) {
 
 	go func() {
 		clientSend.Wait()
-		receiveCommand, err := client.InvokeSync(addr,
+		receiveCommand, err := client.InvokeSync(context.Background(), addr,
 			clientSendRemtingCommand, time.Second)
 		if err != nil {
 			t.Fatalf("failed to invoke synchronous. %s", err)
@@ -237,7 +238,7 @@ func TestInvokeAsync(t *testing.T) {
 			time.Sleep(time.Duration(rand.Intn(100)) * time.Millisecond)
 			t.Logf("[Send: %d] asychronous message", index)
 			sendRemotingCommand := randomNewRemotingCommand()
-			err := client.InvokeAsync(addr, sendRemotingCommand, time.Second, func(r *ResponseFuture) {
+			err := client.InvokeAsync(context.Background(), addr, sendRemotingCommand, time.Second, func(r *ResponseFuture) {
 				t.Logf("[Receive: %d] asychronous message response", index)
 				if string(sendRemotingCommand.Body) != string(r.ResponseCommand.Body) {
 					t.Errorf("wrong response message. want=%s, got=%s", string(sendRemotingCommand.Body),
@@ -303,7 +304,7 @@ func TestInvokeAsyncTimeout(t *testing.T) {
 	clientSend.Add(1)
 	go func() {
 		clientSend.Wait()
-		err := client.InvokeAsync(addr, clientSendRemtingCommand,
+		err := client.InvokeAsync(context.Background(), addr, clientSendRemtingCommand,
 			time.Duration(1000), func(r *ResponseFuture) {
 				assert.NotNil(t, r.Err)
 				assert.Equal(t, utils.ErrRequestTimeout, r.Err)
@@ -348,7 +349,7 @@ func TestInvokeOneWay(t *testing.T) {
 	clientSend.Add(1)
 	go func() {
 		clientSend.Wait()
-		err := client.InvokeOneWay(addr, clientSendRemtingCommand, 3*time.Second)
+		err := client.InvokeOneWay(context.Background(), addr, clientSendRemtingCommand, 3*time.Second)
 		if err != nil {
 			t.Fatalf("failed to invoke synchronous. %s", err)
 		}
diff --git a/internal/route_test.go b/internal/route_test.go
new file mode 100644
index 0000000..aa944b0
--- /dev/null
+++ b/internal/route_test.go
@@ -0,0 +1,68 @@
+/*
+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 (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/pkg/errors"
+	. "github.com/smartystreets/goconvey/convey"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/apache/rocketmq-client-go/internal/remote"
+	"github.com/apache/rocketmq-client-go/primitive"
+)
+
+func TestQueryTopicRouteInfoFromServer(t *testing.T) {
+	Convey("marshal of TraceContext", t, func() {
+
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+
+		remotingCli := remote.NewMockRemotingClient(ctrl)
+
+		addr, err := primitive.NewNamesrvAddr("1.1.1.1:8880", "1.1.1.2:8880", "1.1.1.3:8880")
+		assert.Nil(t, err)
+
+		namesrv, err := NewNamesrv(addr)
+		assert.Nil(t, err)
+		namesrv.nameSrvClient = remotingCli
+
+		Convey("When marshal producer trace data", func() {
+
+			count := 0
+			remotingCli.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
+				func(ctx context.Context, addr string, request *remote.RemotingCommand, timeout time.Duration) (*remote.RemotingCommand, error) {
+					count++
+					if count < 3 {
+						return nil, errors.New("not existed")
+					}
+					return &remote.RemotingCommand{
+						Code: ResTopicNotExist,
+					}, nil
+				}).Times(3)
+
+			data, err := namesrv.queryTopicRouteInfoFromServer("notexisted")
+			assert.Nil(t, data)
+			assert.Equal(t, ErrTopicNotExist, err)
+		})
+	})
+}
diff --git a/internal/utils/net_test.go b/internal/utils/net_test.go
index 44454ad..3e65a75 100644
--- a/internal/utils/net_test.go
+++ b/internal/utils/net_test.go
@@ -1,3 +1,20 @@
+/*
+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 "testing"
diff --git a/internal/utils/string_test.go b/internal/utils/string_test.go
index d4b585b..4e9e4ea 100644
--- a/internal/utils/string_test.go
+++ b/internal/utils/string_test.go
@@ -1,1 +1,18 @@
+/*
+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
diff --git a/producer/producer_test.go b/producer/producer_test.go
index 849978a..d4c58c2 100644
--- a/producer/producer_test.go
+++ b/producer/producer_test.go
@@ -130,7 +130,7 @@ func TestSync(t *testing.T) {
 
 	mockB4Send(p)
 
-	client.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
+	client.EXPECT().InvokeSync(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
 	client.EXPECT().ProcessSendResponse(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Do(
 		func(brokerName string, cmd *remote.RemotingCommand, resp *primitive.SendResult, msgs ...*primitive.Message) {
 			resp.Status = expectedResp.Status
@@ -181,8 +181,8 @@ func TestASync(t *testing.T) {
 
 	mockB4Send(p)
 
-	client.EXPECT().InvokeAsync(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addr string, request *remote.RemotingCommand,
+	client.EXPECT().InvokeAsync(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
+		func(ctx context.Context, addr string, request *remote.RemotingCommand,
 			timeoutMillis time.Duration, f func(*remote.RemotingCommand, error)) error {
 			// mock invoke callback
 			f(nil, nil)
@@ -226,7 +226,7 @@ func TestOneway(t *testing.T) {
 
 	mockB4Send(p)
 
-	client.EXPECT().InvokeOneWay(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	client.EXPECT().InvokeOneWay(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	err = p.SendOneWay(ctx, msg)
 	assert.Nil(t, err)