feat: add RangePersistentValues (bytedance#151)

* feat: add RangePersistentValues Change-Id: I60f02e8a0769ea855d775cbfc1980915bb76dfff Change-Id: I1d9a81d81b7fec7ee90396ebdeaa2e8e6e213cfa Co-authored-by: caimufu <[email protected]>
trim21 · Aug 24, 2022 · beb9000 · beb9000
1 parent b879a72
commit beb9000
Show file tree

Hide file tree

Showing 3 changed files with 128 additions and 0 deletions.
diff --git a/cloud/metainfo/README-CN.md b/cloud/metainfo/README-CN.md
@@ -48,6 +48,8 @@ metainfo 包提供了几个常量字符串前缀，用于无 context（例如网
     - 从 context 里获取指定 key 的 transient 数据（包括 transient-upstream 数据）。
 - `GetAllValues(ctx context.Context) map[string]string`
     - 从 context 里获取所有 transient 数据（包括 transient-upstream 数据）。
+- `RangeValues(ctx context.Context, f func(k, v string) bool)`
+    - 从 context 里基于 f 过滤获取 transient 数据（包括 transient-upstream 数据）。
 - `WithValue(ctx context.Context, k string, v string) context.Context`
     - 向 context 里添加一个 transient 数据。
 - `DelValue(ctx context.Context, k string) context.Context`
@@ -56,6 +58,8 @@ metainfo 包提供了几个常量字符串前缀，用于无 context（例如网
     - 从 context 里获取指定 key 的 persistent 数据。
 - `GetAllPersistentValues(ctx context.Context) map[string]string`
     - 从 context 里获取所有 persistent 数据。
+- `RangePersistentValues(ctx context.Context, f func(k, v string) bool)`
+    - 从 context 里基于 f 过滤获取 persistent 数据。
 - `WithPersistentValue(ctx context.Context, k string, v string) context.Context`
     - 向 context 里添加一个 persistent 数据。
 - `DelPersistentValue(ctx context.Context, k string) context.Context`

diff --git a/cloud/metainfo/info.go b/cloud/metainfo/info.go
@@ -73,6 +73,31 @@ func GetAllValues(ctx context.Context) (m map[string]string) {
 	return
 }
 
+// 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) {
+	n := getNode(ctx)
+	if n == nil {
+		return
+	}
+
+	if cnt := len(n.stale) + len(n.transient); cnt == 0 {
+		return
+	}
+
+	for _, kv := range n.stale {
+		if !f(kv.key, kv.val) {
+			return
+		}
+	}
+
+	for _, kv := range n.transient {
+		if !f(kv.key, kv.val) {
+			return
+		}
+	}
+}
+
 // WithValue sets the value into the context by the given key.
 // This value will be propagated to the next service/endpoint through an RPC call.
 //
@@ -131,6 +156,21 @@ func GetAllPersistentValues(ctx context.Context) (m map[string]string) {
 	return
 }
 
+// RangePersistentValues calls f sequentially for each persistent kv.
+// If f returns false, range stops the iteration.
+func RangePersistentValues(ctx context.Context, f func(k, v string) bool) {
+	n := getNode(ctx)
+	if n == nil {
+		return
+	}
+
+	for _, kv := range n.persistent {
+		if !f(kv.key, kv.val) {
+			break
+		}
+	}
+}
+
 // WithPersistentValue sets the value info the context by the given key.
 // This value will be propagated to the services along the RPC call chain.
 func WithPersistentValue(ctx context.Context, k, v string) context.Context {

diff --git a/cloud/metainfo/info_test.go b/cloud/metainfo/info_test.go
@@ -88,6 +88,29 @@ func TestGetAll(t *testing.T) {
 	}
 }
 
+func TestRangeValues(t *testing.T) {
+	ctx := context.Background()
+
+	ss := []string{"1", "2", "3"}
+	for _, k := range ss {
+		ctx = metainfo.WithValue(ctx, "key"+k, "val"+k)
+	}
+
+	m := make(map[string]string, 3)
+	f := func(k, v string) bool {
+		m[k] = v
+		return true
+	}
+
+	metainfo.RangeValues(ctx, f)
+	assert(t, m != nil)
+	assert(t, len(m) == len(ss))
+
+	for _, k := range ss {
+		assert(t, m["key"+k] == "val"+k)
+	}
+}
+
 func TestGetAll2(t *testing.T) {
 	ctx := context.Background()
 
@@ -180,6 +203,29 @@ func TestGetAllPersistent(t *testing.T) {
 	}
 }
 
+func TestRangePersistent(t *testing.T) {
+	ctx := context.Background()
+
+	ss := []string{"1", "2", "3"}
+	for _, k := range ss {
+		ctx = metainfo.WithPersistentValue(ctx, "key"+k, "val"+k)
+	}
+
+	m := make(map[string]string, 3)
+	f := func(k, v string) bool {
+		m[k] = v
+		return true
+	}
+
+	metainfo.RangePersistentValues(ctx, f)
+	assert(t, m != nil)
+	assert(t, len(m) == len(ss))
+
+	for _, k := range ss {
+		assert(t, m["key"+k] == "val"+k)
+	}
+}
+
 func TestGetAllPersistent2(t *testing.T) {
 	ctx := context.Background()
 
@@ -344,6 +390,14 @@ func benchmark(b *testing.B, api string, count int) {
 		for i := 0; i < b.N; i++ {
 			_ = metainfo.GetAllValues(ctx)
 		}
+	case "RangeValues":
+		b.ReportAllocs()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			metainfo.RangeValues(ctx, func(_, _ string) bool {
+				return true
+			})
+		}
 	case "WithValue":
 		b.ReportAllocs()
 		b.ResetTimer()
@@ -374,6 +428,14 @@ func benchmark(b *testing.B, api string, count int) {
 		for i := 0; i < b.N; i++ {
 			_ = metainfo.GetAllPersistentValues(ctx)
 		}
+	case "RangePersistentValues":
+		b.ReportAllocs()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			metainfo.RangePersistentValues(ctx, func(_, _ string) bool {
+				return true
+			})
+		}
 	case "WithPersistentValue":
 		b.ReportAllocs()
 		b.ResetTimer()
@@ -441,6 +503,16 @@ func benchmarkParallel(b *testing.B, api string, count int) {
 				_ = metainfo.GetAllValues(ctx)
 			}
 		})
+	case "RangeValues":
+		b.ReportAllocs()
+		b.ResetTimer()
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				metainfo.RangeValues(ctx, func(_, _ string) bool {
+					return true
+				})
+			}
+		})
 	case "WithValue":
 		b.ReportAllocs()
 		b.ResetTimer()
@@ -486,6 +558,16 @@ func benchmarkParallel(b *testing.B, api string, count int) {
 				_ = metainfo.GetAllPersistentValues(ctx)
 			}
 		})
+	case "RangePersistentValues":
+		b.ReportAllocs()
+		b.ResetTimer()
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				metainfo.RangePersistentValues(ctx, func(_, _ string) bool {
+					return true
+				})
+			}
+		})
 	case "WithPersistentValue":
 		b.ReportAllocs()
 		b.ResetTimer()
@@ -546,6 +628,7 @@ func BenchmarkAll(b *testing.B) {
 		"DelValue",
 		"GetPersistentValue",
 		"GetAllPersistentValues",
+		"RangePersistentValues",
 		"WithPersistentValue",
 		"WithPersistentValueAcc",
 		"DelPersistentValue",
@@ -570,6 +653,7 @@ func BenchmarkAllParallel(b *testing.B) {
 		"DelValue",
 		"GetPersistentValue",
 		"GetAllPersistentValues",
+		"RangePersistentValues",
 		"WithPersistentValue",
 		"WithPersistentValueAcc",
 		"DelPersistentValue",