From a03554c296f839861396fcde6e5a76721ebe7417 Mon Sep 17 00:00:00 2001
From: QihengZhou <qiheng.zhou@bytedance.com>
Date: Thu, 11 Jul 2024 16:50:56 +0800
Subject: [PATCH] feat: add GetValueToMap that sets value to the input map
 (#212)

---
 cloud/metainfo/info.go      | 22 ++++++++++++
 cloud/metainfo/info_test.go | 69 +++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/cloud/metainfo/info.go b/cloud/metainfo/info.go
index efda6e38..b11f08a2 100644
--- a/cloud/metainfo/info.go
+++ b/cloud/metainfo/info.go
@@ -73,6 +73,28 @@ func GetAllValues(ctx context.Context) (m map[string]string) {
 	return
 }
 
+// GetValueToMap retrieves the value set into the context by the given key and set the value to the input map.
+// Only use this function when you want to get a small set of values instead of GetAllValues.
+// The logic of getting value follows GetAllValues, transient value has higher priority if key is same.
+func GetValueToMap(ctx context.Context, m map[string]string, keys ...string) {
+	if m == nil || len(keys) == 0 {
+		return
+	}
+	n := getNode(ctx)
+	if n == nil {
+		return
+	}
+	for _, k := range keys {
+		if idx, ok := search(n.transient, k); ok {
+			m[k] = n.transient[idx].val
+			continue
+		}
+		if idx, ok := search(n.stale, k); ok {
+			m[k] = n.stale[idx].val
+		}
+	}
+}
+
 // RangeValues calls f sequentially for each transient kv.
 // If f returns false, range stops the iteration.
 func RangeValues(ctx context.Context, f func(k, v string) bool) {
diff --git a/cloud/metainfo/info_test.go b/cloud/metainfo/info_test.go
index d9869605..093b7954 100644
--- a/cloud/metainfo/info_test.go
+++ b/cloud/metainfo/info_test.go
@@ -17,6 +17,7 @@ package metainfo_test
 import (
 	"context"
 	"fmt"
+	"math"
 	"testing"
 
 	"github.com/bytedance/gopkg/cloud/metainfo"
@@ -461,12 +462,21 @@ func benchmark(b *testing.B, api string, count int) {
 		for i := 0; i < b.N; i++ {
 			_, _ = metainfo.GetValue(ctx, keys[i%len(keys)])
 		}
+	case "GetValueToMap":
+		b.ReportAllocs()
+		b.ResetTimer()
+		m := make(map[string]string, len(keys))
+		for i := 0; i < b.N; i++ {
+			metainfo.GetValueToMap(ctx, m, keys...)
+		}
 	case "GetAllValues":
 		b.ReportAllocs()
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			_ = metainfo.GetAllValues(ctx)
 		}
+	case "GetValueWithKeys":
+		benchmarkGetValueWithKeys(b, ctx, keys)
 	case "RangeValues":
 		b.ReportAllocs()
 		b.ResetTimer()
@@ -563,6 +573,42 @@ func benchmark(b *testing.B, api string, count int) {
 	}
 }
 
+func benchmarkGetValueWithKeys(b *testing.B, ctx context.Context, keys []string) {
+	selectedRatio := 0.1
+	selectedKeyLength := uint64(math.Round(selectedRatio * float64(len(keys))))
+	selectedKeys := keys[:selectedKeyLength]
+
+	b.Run("GetValue", func(b *testing.B) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		m := make(map[string]string, len(selectedKeys))
+		for i := 0; i < b.N; i++ {
+			for j := 0; j < len(selectedKeys); j++ {
+				key := selectedKeys[j]
+				v, _ := metainfo.GetValue(ctx, key)
+				m[key] = v
+			}
+		}
+	})
+
+	b.Run("GetValueToMap", func(b *testing.B) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		m := make(map[string]string, len(selectedKeys))
+		for i := 0; i < b.N; i++ {
+			metainfo.GetValueToMap(ctx, m, selectedKeys...)
+		}
+	})
+
+	b.Run("GetAllValue", func(b *testing.B) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = metainfo.GetAllValues(ctx)
+		}
+	})
+}
+
 func benchmarkParallel(b *testing.B, api string, count int) {
 	ctx, keys, vals := initMetaInfo(count)
 	switch api {
@@ -584,6 +630,17 @@ func benchmarkParallel(b *testing.B, api string, count int) {
 				i++
 			}
 		})
+	case "GetValueToMap":
+		b.ReportAllocs()
+		b.ResetTimer()
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				m := make(map[string]string, len(keys))
+				for i := 0; i < b.N; i++ {
+					metainfo.GetValueToMap(ctx, m, keys...)
+				}
+			}
+		})
 	case "GetAllValues":
 		b.ReportAllocs()
 		b.ResetTimer()
@@ -727,7 +784,9 @@ func BenchmarkAll(b *testing.B) {
 	APIs := []string{
 		"TransferForward",
 		"GetValue",
+		"GetValueToMap",
 		"GetAllValues",
+		"GetValueWithKeys",
 		"WithValue",
 		"WithValues",
 		"WithValueAcc",
@@ -754,6 +813,7 @@ func BenchmarkAllParallel(b *testing.B) {
 	APIs := []string{
 		"TransferForward",
 		"GetValue",
+		"GetValueToMap",
 		"GetAllValues",
 		"WithValue",
 		"WithValues",
@@ -858,3 +918,12 @@ func TestPersistentValuesCount(t *testing.T) {
 		})
 	}
 }
+
+func TestGetValueToMap(t *testing.T) {
+	ctx := context.Background()
+	k, v := "key", "value"
+	ctx = metainfo.WithValue(ctx, k, v)
+	m := make(map[string]string, 1)
+	metainfo.GetValueToMap(ctx, m, k)
+	assert(t, m[k] == v)
+}