feat: a new userData API Remove

server.go

@@ -698,6 +698,16 @@ func (ctx *RequestCtx) ResetUserValues() {
 	ctx.userValues.Reset()
 }
 
+// RemoveUserValue removes the given key and the value under it in ctx.
+func (ctx *RequestCtx) RemoveUserValue(key string) {
+	ctx.userValues.Remove(key)
+}
+
+// RemoveUserValueBytes removes the given key and the value under it in ctx.
+func (ctx *RequestCtx) RemoveUserValueBytes(key []byte) {
+	ctx.userValues.RemoveBytes(key)
+}
+
 type connTLSer interface {
 	Handshake() error
 	ConnectionState() tls.ConnectionState

userdata.go

@@ -2,25 +2,61 @@ package fasthttp
 
 import (
 	"io"
+	"sync/atomic"
+)
+
+const (
+	statusAlive = iota
+	statusLocked
+	statusDeleted
 )
 
 type userDataKV struct {
-	key   []byte
-	value interface{}
+	key    []byte
+	value  interface{}
+	status uint32
 }
 
 type userData []userDataKV
 
 func (d *userData) Set(key string, value interface{}) {
 	args := *d
 	n := len(args)
+	idx := -1 // the index of a deleted userDataKV in userData. for memory reuse.
+
+	defer func() {
+		if idx == -1 {
+			return
+		}
+		// unlock the locked userDataKV
+		atomic.CompareAndSwapUint32(&args[idx].status, statusLocked, statusDeleted)
+	}()
+
 	for i := 0; i < n; i++ {
 		kv := &args[i]
 		if string(kv.key) == key {
 			kv.value = value
+			kv.status = statusAlive
 			return
 		}
+		// lock the locked userDataKV and record its index.
+		if idx == -1 {
+			if ok := atomic.CompareAndSwapUint32(&kv.status, statusDeleted, statusLocked); ok {
+				idx = i
+			}
+		}
 	}
+
+	// reuse
+	if idx != -1 {
+		kv := &args[idx]
+		if kv.status == statusLocked {
+			kv.key = append(kv.key[:0], key...)
+			kv.value = value
+			kv.status = statusAlive
+		}
+	}
+
 	if value == nil {
 		return
 	}
@@ -50,7 +86,7 @@ func (d *userData) Get(key string) interface{} {
 	n := len(args)
 	for i := 0; i < n; i++ {
 		kv := &args[i]
-		if string(kv.key) == key {
+		if string(kv.key) == key && kv.status == statusAlive {
 			return kv.value
 		}
 	}
@@ -72,3 +108,22 @@ func (d *userData) Reset() {
 	}
 	*d = (*d)[:0]
 }
+
+func (d *userData) Remove(key string) {
+	args := *d
+	n := len(args)
+	for i := 0; i < n; i++ {
+		kv := &args[i]
+		if string(kv.key) == key {
+			kv.status = statusDeleted
+			if vc, ok := kv.value.(io.Closer); ok {
+				vc.Close()
+			}
+			return
+		}
+	}
+}
+
+func (d *userData) RemoveBytes(key []byte) {
+	d.Remove(b2s(key))
+}

userdata_test.go

@@ -76,3 +76,31 @@ func (cv *closerValue) Close() error {
 	(*cv.closeCalls)++
 	return nil
 }
+
+func TestUserDataDelete(t *testing.T) {
+	t.Parallel()
+
+	var u userData
+
+	for i := 0; i < 10; i++ {
+		key := fmt.Sprintf("key_%d", i)
+		u.Set(key, i)
+		testUserDataGet(t, &u, []byte(key), i)
+	}
+
+	for i := 0; i < 10; i += 2 {
+		k := fmt.Sprintf("key_%d", i)
+		u.Remove(k)
+		if val := u.Get(k); val != nil {
+			t.Fatalf("unexpected key= %s, value =%v ,Expecting key= %s, value = nil", k, val, k)
+		}
+		kk := fmt.Sprintf("key_%d", i+1)
+		testUserDataGet(t, &u, []byte(kk), i+1)
+	}
+	for i := 0; i < 10; i++ {
+		key := fmt.Sprintf("key_new_%d", i)
+		u.Set(key, i)
+		testUserDataGet(t, &u, []byte(key), i)
+	}
+
+}