mirror of
https://github.com/zeromicro/go-zero.git
synced 2025-01-23 09:00:20 +08:00
fix: prevent ring index overflow (#3738)
This commit is contained in:
parent
5e8e21b257
commit
c220b5d886
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user