From ce0647ce255e811434715296e5ef0130740289b0 Mon Sep 17 00:00:00 2001
From: "duanyi.aster" <duanyi.aster@bytedance.com>
Date: Thu, 27 Jul 2023 19:24:21 +0800
Subject: [PATCH] opt: add Len() API for better performance when dumping KVS

---
 cloud/metainfo/info.go      | 18 ++++++++
 cloud/metainfo/info_test.go | 82 +++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/cloud/metainfo/info.go b/cloud/metainfo/info.go
index 427900ee..297f5523 100644
--- a/cloud/metainfo/info.go
+++ b/cloud/metainfo/info.go
@@ -211,6 +211,24 @@ func getValue(kvs []string, i int) string {
 	return kvs[i*2+1]
 }
 
+// CountPersistentValues counts the length of persisten KV pairs
+func CountPersistentValues(ctx context.Context) int {
+	if n := getNode(ctx); n == nil {
+		return 0
+	} else {
+		return len(n.persistent)
+	}
+}
+
+// CountValues counts the length of transient KV pairs
+func CountValues(ctx context.Context) int {
+	if n := getNode(ctx); n == nil {
+		return 0
+	} else {
+		return len(n.transient)
+	}
+}
+
 // WithPersistentValues sets the values into the context by the given keys.
 // This value will be propagated to the services along the RPC call chain.
 func WithPersistentValues(ctx context.Context, kvs ...string) context.Context {
diff --git a/cloud/metainfo/info_test.go b/cloud/metainfo/info_test.go
index 7244766e..d9869605 100644
--- a/cloud/metainfo/info_test.go
+++ b/cloud/metainfo/info_test.go
@@ -776,3 +776,85 @@ func BenchmarkAllParallel(b *testing.B) {
 		}
 	}
 }
+
+func TestValuesCount(t *testing.T) {
+	ctx := context.Background()
+	type args struct {
+		ctx context.Context
+	}
+	tests := []struct {
+		name string
+		args args
+		want int
+	}{
+		{
+			name: "0",
+			args: args{
+				ctx: ctx,
+			},
+			want: 0,
+		},
+		{
+			name: "0",
+			args: args{
+				ctx: metainfo.WithPersistentValues(ctx, "1", "1", "2", "2"),
+			},
+			want: 0,
+		},
+		{
+			name: "2",
+			args: args{
+				ctx: metainfo.WithValues(ctx, "1", "1", "2", "2"),
+			},
+			want: 2,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := metainfo.CountValues(tt.args.ctx); got != tt.want {
+				t.Errorf("ValuesCount() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestPersistentValuesCount(t *testing.T) {
+	ctx := context.Background()
+	type args struct {
+		ctx context.Context
+	}
+	tests := []struct {
+		name string
+		args args
+		want int
+	}{
+		{
+			name: "0",
+			args: args{
+				ctx: ctx,
+			},
+			want: 0,
+		},
+		{
+			name: "0",
+			args: args{
+				ctx: metainfo.WithValues(ctx, "1", "1", "2", "2"),
+			},
+			want: 0,
+		},
+		{
+			name: "2",
+			args: args{
+				ctx: metainfo.WithPersistentValues(ctx, "1", "1", "2", "2"),
+			},
+			want: 2,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := metainfo.CountPersistentValues(tt.args.ctx); got != tt.want {
+				t.Errorf("ValuesCount() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}