Repository: apache/rocketmq-client-go
Pull request #574 (2020-12-29 04:02:05)

fix possible infinite loop after rebalance
## What is the purpose of the change

fix possible infinite loop after rebalance

## Brief changelog

it's very likely there will be cpu/memory drain caused by dead loop after a rebalance
this fix addressed the issue by check whether the process queue is already dropped before starting next consume





Patch:
diff --git a/.gitignore b/.gitignore
index a0c292d..99fd04d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,5 @@ go.mod
 go.sum
 vendor/
 coverage.txt
-examples/test
\ No newline at end of file
+examples/test
+/.vscode
\ No newline at end of file
diff --git a/LICENSE b/LICENSE
index 0130891..e6e77b0 100644
--- a/LICENSE
+++ b/LICENSE
@@ -15,7 +15,7 @@
       "Legal Entity" shall mean the union of the acting entity and all
       other entities that control, are controlled by, or are under common
       control with that entity. For the purposes of this definition,
-      "control" means (properties) the power, direct or indirect, to cause the
+      "control" means (i) the power, direct or indirect, to cause the
       direction or management of such entity, whether by contract or
       otherwise, or (ii) ownership of fifty percent (50%) or more of the
       outstanding shares, or (iii) beneficial ownership of such entity.
@@ -198,4 +198,4 @@
    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.
\ No newline at end of file
+   limitations under the License.
diff --git a/README.md b/README.md
index 0f30d5c..08494a6 100644
--- a/README.md
+++ b/README.md
@@ -52,4 +52,4 @@ For 2.0.0 version, it supports:
    
 ----------
 ## License
-  [Apache License, Version 2.0](http://www.apache.org/licenses/LICENSE-2.0.html) Copyright (C) Apache Software Foundation
+  [Apache License, Version 2.0](http://www.apache.org/licenses/LICENSE-2.0.html) Copyright (C) Apache Software Foundation
\ No newline at end of file
diff --git a/consumer/offset_store.go b/consumer/offset_store.go
index 17f5d76..4d0b271 100644
--- a/consumer/offset_store.go
+++ b/consumer/offset_store.go
@@ -26,13 +26,12 @@ import (
 	"sync"
 	"time"
 
-	jsoniter "github.com/json-iterator/go"
-
 	"github.com/apache/rocketmq-client-go/v2/internal"
 	"github.com/apache/rocketmq-client-go/v2/internal/remote"
 	"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"
+	jsoniter "github.com/json-iterator/go"
 )
 
 type readType int
@@ -101,7 +100,7 @@ func (mq *MessageQueueKey) UnmarshalText(text []byte) error {
 type localFileOffsetStore struct {
 	group       string
 	path        string
-	OffsetTable map[MessageQueueKey]int64
+	OffsetTable *sync.Map // concurrent safe , map[MessageQueueKey]int64
 	// mutex for offset file
 	mutex sync.Mutex
 }
@@ -110,7 +109,7 @@ func NewLocalFileOffsetStore(clientID, group string) OffsetStore {
 	store := &localFileOffsetStore{
 		group:       group,
 		path:        filepath.Join(_LocalOffsetStorePath, clientID, group, "offset.json"),
-		OffsetTable: make(map[MessageQueueKey]int64),
+		OffsetTable: new(sync.Map),
 	}
 	store.load()
 	return store
@@ -151,7 +150,9 @@ func (local *localFileOffsetStore) load() {
 	}
 
 	if datas != nil {
-		local.OffsetTable = datas
+		for k, v := range datas {
+			local.OffsetTable.Store(k, v)
+		}
 	}
 }
 
@@ -180,17 +181,17 @@ func (local *localFileOffsetStore) update(mq *primitive.MessageQueue, offset int
 		"new_offset":            offset,
 	})
 	key := MessageQueueKey(*mq)
-	localOffset, exist := local.OffsetTable[key]
+	localOffset, exist := local.OffsetTable.Load(key)
 	if !exist {
-		local.OffsetTable[key] = offset
+		local.OffsetTable.Store(key, offset)
 		return
 	}
 	if increaseOnly {
-		if localOffset < offset {
-			local.OffsetTable[key] = offset
+		if localOffset.(int64) < offset {
+			local.OffsetTable.Store(key, offset)
 		}
 	} else {
-		local.OffsetTable[key] = offset
+		local.OffsetTable.Store(key, offset)
 	}
 }
 
@@ -201,10 +202,17 @@ func (local *localFileOffsetStore) persist(mqs []*primitive.MessageQueue) {
 	local.mutex.Lock()
 	defer local.mutex.Unlock()
 
+	datas := make(map[MessageQueueKey]int64)
+	local.OffsetTable.Range(func(key, value interface{}) bool {
+		k := key.(MessageQueueKey)
+		v := value.(int64)
+		datas[k] = v
+		return true
+	})
+
 	wrapper := OffsetSerializeWrapper{
-		OffsetTable: local.OffsetTable,
+		OffsetTable: datas,
 	}
-
 	data, _ := jsoniter.Marshal(wrapper)
 	utils.CheckError(fmt.Sprintf("persist offset to %s", local.path), utils.WriteToFile(local.path, data))
 }
@@ -384,11 +392,11 @@ func (r *remoteBrokerOffsetStore) updateConsumeOffsetToBroker(group string, mq p
 	return r.client.InvokeOneWay(context.Background(), broker, cmd, 5*time.Second)
 }
 
-func readFromMemory(table map[MessageQueueKey]int64, mq *primitive.MessageQueue) int64 {
-	localOffset, exist := table[MessageQueueKey(*mq)]
+func readFromMemory(table *sync.Map, mq *primitive.MessageQueue) int64 {
+	localOffset, exist := table.Load(MessageQueueKey(*mq))
 	if !exist {
 		return -1
 	}
 
-	return localOffset
+	return localOffset.(int64)
 }
diff --git a/consumer/push_consumer.go b/consumer/push_consumer.go
index 931fefa..f13d5e2 100644
--- a/consumer/push_consumer.go
+++ b/consumer/push_consumer.go
@@ -465,6 +465,12 @@ func (pc *pushConsumer) pullMessage(request *PullRequest) {
 				return
 			default:
 				pc.submitToConsume(request.pq, request.mq)
+				if request.pq.IsDroppd() {
+					rlog.Info("push consumer quit pullMessage for dropped queue.", map[string]interface{}{
+						rlog.LogKeyConsumerGroup: pc.consumerGroup,
+					})
+					return
+				}
 			}
 		}
 	})
@@ -516,7 +522,7 @@ func (pc *pushConsumer) pullMessage(request *PullRequest) {
 					"PullThresholdForQueue": pc.option.PullThresholdForQueue,
 					"minOffset":             pq.Min(),
 					"maxOffset":             pq.Max(),
-					"count":                 pq.msgCache,
+					"count":                 pq.cachedMsgCount,
 					"size(MiB)":             cachedMessageSizeInMiB,
 					"flowControlTimes":      pc.queueFlowControlTimes,
 					rlog.LogKeyPullRequest:  request.String(),
@@ -533,7 +539,7 @@ func (pc *pushConsumer) pullMessage(request *PullRequest) {
 					"PullThresholdSizeForQueue": pc.option.PullThresholdSizeForQueue,
 					"minOffset":                 pq.Min(),
 					"maxOffset":                 pq.Max(),
-					"count":                     pq.msgCache,
+					"count":                     pq.cachedMsgCount,
 					"size(MiB)":                 cachedMessageSizeInMiB,
 					"flowControlTimes":          pc.queueFlowControlTimes,
 					rlog.LogKeyPullRequest:      request.String(),
diff --git a/primitive/nsresolver.go b/primitive/nsresolver.go
index d844373..4e5917c 100644
--- a/primitive/nsresolver.go
+++ b/primitive/nsresolver.go
@@ -130,11 +130,15 @@ func (h *HttpResolver) Description() string {
 
 func (h *HttpResolver) get() []string {
 	resp, err := h.cli.Get(h.domain)
-	if err != nil {
-		rlog.Error("name server http fetch failed", map[string]interface{}{
+	if err != nil || resp == nil || resp.StatusCode != 200 {
+		data := map[string]interface{}{
 			"NameServerDomain": h.domain,
 			"err":              err,
-		})
+		}
+		if resp != nil {
+			data["StatusCode"] = resp.StatusCode
+		}
+		rlog.Error("name server http fetch failed", data)
 		return nil
 	}
 


Test patch:
diff --git a/consumer/offset_store_test.go b/consumer/offset_store_test.go
index 27b98d9..cfa0eaa 100644
--- a/consumer/offset_store_test.go
+++ b/consumer/offset_store_test.go
@@ -21,12 +21,11 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/golang/mock/gomock"
-	. "github.com/smartystreets/goconvey/convey"
-
 	"github.com/apache/rocketmq-client-go/v2/internal"
 	"github.com/apache/rocketmq-client-go/v2/internal/remote"
 	"github.com/apache/rocketmq-client-go/v2/primitive"
+	"github.com/golang/mock/gomock"
+	. "github.com/smartystreets/goconvey/convey"
 )
 
 func TestNewLocalFileOffsetStore(t *testing.T) {
@@ -136,7 +135,7 @@ func TestLocalFileOffsetStore(t *testing.T) {
 			offset = localStore.read(mq, _ReadFromStore)
 			So(offset, ShouldEqual, 1)
 
-			delete(localStore.(*localFileOffsetStore).OffsetTable, MessageQueueKey(*mq))
+			localStore.(*localFileOffsetStore).OffsetTable.Delete(MessageQueueKey(*mq))
 			offset = localStore.read(mq, _ReadMemoryThenStore)
 			So(offset, ShouldEqual, 1)
 		})