浏览代码

fix: prevent ring index overflow (#3738)

POABOB 1 年之前
父节点
当前提交
c220b5d886
共有 1 个文件被更改,包括 17 次插入6 次删除
  1. 17 6
      core/collection/ring.go

+ 17 - 6
core/collection/ring.go

@@ -1,6 +1,8 @@
 package collection
 
-import "sync"
+import (
+	"sync"
+)
 
 // A Ring can be used as fixed size ring.
 type Ring struct {
@@ -25,8 +27,15 @@ func (r *Ring) Add(v any) {
 	r.lock.Lock()
 	defer r.lock.Unlock()
 
-	r.elements[r.index%len(r.elements)] = v
+	ringLength := len(r.elements)
+
+	r.elements[r.index%ringLength] = v
 	r.index++
+
+	// prevent ring index overflow
+	if r.index/ringLength >= 2 {
+		r.index = r.index - ringLength
+	}
 }
 
 // Take takes all items from r.
@@ -36,16 +45,18 @@ func (r *Ring) Take() []any {
 
 	var size int
 	var start int
-	if r.index > len(r.elements) {
-		size = len(r.elements)
-		start = r.index % len(r.elements)
+	ringLength := len(r.elements)
+
+	if r.index > ringLength {
+		size = ringLength
+		start = r.index % ringLength
 	} else {
 		size = r.index
 	}
 
 	elements := make([]any, size)
 	for i := 0; i < size; i++ {
-		elements[i] = r.elements[(start+i)%len(r.elements)]
+		elements[i] = r.elements[(start+i)%ringLength]
 	}
 
 	return elements