From e02c16214bb98372093a3c7817287c208801de7f Mon Sep 17 00:00:00 2001
From: Cr <631807682@qq.com>
Date: Thu, 19 Oct 2023 11:37:03 +0800
Subject: [PATCH] feat: support peek multi value for query args (#684)

Co-authored-by: kinggo <lilong.21@bytedance.com>
---
 pkg/protocol/args.go      | 11 +++++++++++
 pkg/protocol/args_test.go | 23 +++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/pkg/protocol/args.go b/pkg/protocol/args.go
index f29aefead..fa2a670e5 100644
--- a/pkg/protocol/args.go
+++ b/pkg/protocol/args.go
@@ -379,6 +379,17 @@ func (a *Args) PeekExists(key string) (string, bool) {
 	return peekArgStrExists(a.args, key)
 }
 
+// PeekAll returns all the arg values for the given key.
+func (a *Args) PeekAll(key string) [][]byte {
+	var values [][]byte
+	a.VisitAll(func(k, v []byte) {
+		if bytesconv.B2s(k) == key {
+			values = append(values, v)
+		}
+	})
+	return values
+}
+
 func visitArgs(args []argsKV, f func(k, v []byte)) {
 	for i, n := 0, len(args); i < n; i++ {
 		kv := &args[i]
diff --git a/pkg/protocol/args_test.go b/pkg/protocol/args_test.go
index 6a613dc19..0e89f30cc 100644
--- a/pkg/protocol/args_test.go
+++ b/pkg/protocol/args_test.go
@@ -129,3 +129,26 @@ func TestArgsVisitAll(t *testing.T) {
 	})
 	assert.DeepEqual(t, []string{"cloudwego", "hertz", "hello", "world"}, s)
 }
+
+func TestArgsPeekMulti(t *testing.T) {
+	var a Args
+	a.Add("cloudwego", "hertz")
+	a.Add("cloudwego", "kitex")
+	a.Add("cloudwego", "")
+	a.Add("hello", "world")
+
+	vv := a.PeekAll("cloudwego")
+	expectedVV := [][]byte{
+		[]byte("hertz"),
+		[]byte("kitex"),
+		[]byte(nil),
+	}
+	assert.DeepEqual(t, expectedVV, vv)
+
+	vv = a.PeekAll("aaaa")
+	assert.DeepEqual(t, 0, len(vv))
+
+	vv = a.PeekAll("hello")
+	expectedVV = [][]byte{[]byte("world")}
+	assert.DeepEqual(t, expectedVV, vv)
+}