fix(metainfo): old kvs missed after WithValues() (bytedance#186)

* fix(metainfo): old kvs missed after `WithValues()` * opt: add Len() API for better performance when dumping KVS
joway · Apr 17, 2024 · 16eac62 · 16eac62
1 parent b2af779
commit 16eac62
Show file tree

Hide file tree

Showing 2 changed files with 121 additions and 4 deletions.
diff --git 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.stale) + 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 {
@@ -228,7 +246,7 @@ func WithPersistentValues(ctx context.Context, kvs ...string) context.Context {
 	if m := getNode(ctx); m != nil {
 		nn := *m
 		n = &nn
-		n.persistent = make([]kv, 0, len(m.persistent)+kvLen)
+		n.persistent = make([]kv, len(m.persistent), len(m.persistent)+kvLen)
 		copy(n.persistent, m.persistent)
 	} else {
 		n = &node{
@@ -276,7 +294,7 @@ func WithValues(ctx context.Context, kvs ...string) context.Context {
 	if m := getNode(ctx); m != nil {
 		nn := *m
 		n = &nn
-		n.transient = make([]kv, 0, len(m.transient)+kvLen)
+		n.transient = make([]kv, len(m.transient), len(m.transient)+kvLen)
 		copy(n.transient, m.transient)
 	} else {
 		n = &node{

diff --git a/cloud/metainfo/info_test.go b/cloud/metainfo/info_test.go
@@ -37,20 +37,37 @@ func TestWithValue(t *testing.T) {
 func TestWithValues(t *testing.T) {
 	ctx := context.Background()
 
-	k, v := "Key", "Value"
+	k, v := "Key-0", "Value-0"
 	ctx = metainfo.WithValue(ctx, k, v)
 
 	kvs := []string{"Key-1", "Value-1", "Key-2", "Value-2", "Key-3", "Value-3"}
 	ctx = metainfo.WithValues(ctx, kvs...)
 	assert(t, ctx != nil)
 
-	for i := 1; i <= 3; i++ {
+	for i := 0; i <= 3; i++ {
 		x, ok := metainfo.GetValue(ctx, fmt.Sprintf("Key-%d", i))
 		assert(t, ok)
 		assert(t, x == fmt.Sprintf("Value-%d", i))
 	}
 }
 
+func TestWithPersistValues(t *testing.T) {
+	ctx := context.Background()
+
+	k, v := "Key-0", "Value-0"
+	ctx = metainfo.WithPersistentValue(ctx, k, v)
+
+	kvs := []string{"Key-1", "Value-1", "Key-2", "Value-2", "Key-3", "Value-3"}
+	ctx = metainfo.WithPersistentValues(ctx, kvs...)
+	assert(t, ctx != nil)
+
+	for i := 0; i <= 3; i++ {
+		x, ok := metainfo.GetPersistentValue(ctx, fmt.Sprintf("Key-%d", i))
+		assert(t, ok)
+		assert(t, x == fmt.Sprintf("Value-%d", i))
+	}
+}
+
 func TestWithEmpty(t *testing.T) {
 	ctx := context.Background()
 
@@ -759,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)
+			}
+		})
+	}
+}