From 8510a88e64c4cf62198facd8fbce6495dad61223 Mon Sep 17 00:00:00 2001
From: Thoren-byte <2459501714@qq.com>
Date: Wed, 4 Dec 2024 23:22:19 +0800
Subject: [PATCH] fix: zset remove range by rank lock by read mutex #234

---
 collection/zset/zset.go      |  4 ++--
 collection/zset/zset_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/collection/zset/zset.go b/collection/zset/zset.go
index 46f46336..72801bab 100644
--- a/collection/zset/zset.go
+++ b/collection/zset/zset.go
@@ -382,8 +382,8 @@ func (z *Float64Set) RevRangeByScoreWithOpt(max, min float64, opt RangeOpt) []Fl
 //
 // RemoveRangeByRank is the replacement of ZREMRANGEBYRANK command of redis.
 func (z *Float64Set) RemoveRangeByRank(start, stop int) []Float64Node {
-	z.mu.RLock()
-	defer z.mu.RUnlock()
+	z.mu.Lock()
+	defer z.mu.Unlock()
 
 	// Convert negative rank to positive.
 	if start < 0 {
diff --git a/collection/zset/zset_test.go b/collection/zset/zset_test.go
index 87056fcd..638d1d2b 100644
--- a/collection/zset/zset_test.go
+++ b/collection/zset/zset_test.go
@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"math/rand"
 	"sort"
+	"sync"
 	"testing"
 	"time"
 
@@ -493,6 +494,33 @@ func TestFloat64SetRemoveRangeByRank(t *testing.T) {
 	assert.Equal(t, N, z.Len()+len(actualNs))
 }
 
+func TestFloat64SetRemoveRangeByRankConcurrently(t *testing.T) {
+	const N = 10000
+	z := NewFloat64()
+	for i := 0; i < N; i++ {
+		z.Add(fastrand.Float64(), fmt.Sprint(i))
+	}
+
+	start, stop := func(a, b int) (int, int) {
+		if a < b {
+			return a, b
+		} else {
+			return b, a
+		}
+	}(fastrand.Intn(N), fastrand.Intn(N))
+
+	const G = 10
+	wg := sync.WaitGroup{}
+	for i := 0; i < G; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			z.RemoveRangeByRank(start, stop)
+		}()
+	}
+	wg.Wait()
+}
+
 func TestFloat64SetRemoveRangeByScore(t *testing.T) {
 	const N = 1000
 	z := NewFloat64()
@@ -618,3 +646,4 @@ func TestInterFloat64_Simple(t *testing.T) {
 	z := InterFloat64(z1, z2, z3)
 	assert.Zero(t, z.Len())
 }
+