modify review

db/set.go

@@ -3,6 +3,8 @@ package db
 import (
 	"bytes"
 	"encoding/binary"
+	"math/rand"
+	"time"
 )
 
 // SetNilValue is the value set to a tikv key for tikv do not support a real empty value
@@ -224,12 +226,18 @@ func (set *Set) SIsmember(member []byte) (int64, error) {
 
 // SPop removes and returns one or more random elements from the set value store at key.
 func (set *Set) SPop(count int64) (members [][]byte, err error) {
-	// TODO BUG  No rand
+	// Gets a random number
+	rand.Seed(time.Now().Unix())
+	randNum := int64(rand.Intn(1000))
+	randNum = randNum % (set.meta.Len)
+
 	if !set.Exists() {
 		return nil, nil
 	}
+
 	dkey := DataKey(set.txn.db, set.meta.ID)
 	prefix := append(dkey, ':')
+
 	iter, err := set.txn.t.Iter(prefix, nil)
 	if err != nil {
 		return nil, err
@@ -238,33 +246,44 @@ func (set *Set) SPop(count int64) (members [][]byte, err error) {
 
 	var del int64
 	var ms [][]byte
-	if count == 0 {
-		if iter.Valid() && iter.Key().HasPrefix(prefix) {
+	for iter.Valid() && iter.Key().HasPrefix(prefix) {
+		if count == 0 {
+			for randNum > 0 {
+				if err := iter.Next(); err != nil {
+					return nil, err
+				}
+				randNum--
+				continue
+			}
 			ms = append(ms, iter.Key()[len(prefix):])
 			if err := set.txn.t.Delete([]byte(iter.Key())); err != nil {
 				return nil, err
 			}
-			del++
-		}
-	} else {
-		for iter.Valid() && iter.Key().HasPrefix(prefix) && count != 0 {
+			set.meta.Len--
+			break
+		} else {
 			ms = append(ms, iter.Key()[len(prefix):])
+
 			if err := set.txn.t.Delete([]byte(iter.Key())); err != nil {
 				return nil, err
 			}
 			del++
 			count--
+			if count == 0 {
+				break
+			}
 			if err := iter.Next(); err != nil {
 				return nil, err
-
 			}
 		}
 	}
-	if count < set.meta.Len {
+
+	if count == 0 {
 		set.meta.Len -= del
 	} else {
 		set.meta.Len = 0
 	}
+
 	if err := set.updateMeta(); err != nil {
 		return nil, err
 	}

db/set_test.go

@@ -346,7 +346,7 @@ func TestSet_SCard(t *testing.T) {
 			assert.NoError(t, err)
 
 			if err = txn.Commit(context.TODO()); err != nil {
-				t.Errorf("SIsmember() txn.Commit error = %v", err)
+				t.Errorf("SCard() txn.Commit error = %v", err)
 				return
 			}
 			assert.Equal(t, got, tt.want)
@@ -458,7 +458,7 @@ func TestSet_SPop(t *testing.T) {
 			assert.NotNil(t, gotMembers)
 
 			if err = txn.Commit(context.TODO()); err != nil {
-				t.Errorf("SIsmember() txn.Commit error = %v", err)
+				t.Errorf("SPop() txn.Commit error = %v", err)
 				return
 			}
 			assert.Equal(t, int64(len(gotMembers)), tt.wantMembersCount)
@@ -512,7 +512,7 @@ func TestSet_SRem(t *testing.T) {
 			assert.NotNil(t, got)
 
 			if err = txn.Commit(context.TODO()); err != nil {
-				t.Errorf("SIsmember() txn.Commit error = %v", err)
+				t.Errorf("SRem() txn.Commit error = %v", err)
 				return
 			}
 			assert.Equal(t, got, tt.want)
@@ -608,7 +608,7 @@ func TestSet_SMove(t *testing.T) {
 			assert.NoError(t, err)
 			assert.NotNil(t, got)
 			if err = txn.Commit(context.TODO()); err != nil {
-				t.Errorf("SIsmember() txn.Commit error = %v", err)
+				t.Errorf("SMove() txn.Commit error = %v", err)
 				return
 			}
 			iss, err := set.SIsmember(tt.args.member)