perf(metainfo): optimized SetMetaInfoFromMap

goos: darwin goarch: arm64 pkg: github.com/bytedance/gopkg/cloud/metainfo │ ./old.txt │ ./new.txt │ │ sec/op │ sec/op vs base │ SetMetaInfoFromMap-12 1184.5n ± 3% 889.0n ± 2% -24.94% (p=0.002 n=6) │ ./old.txt │ ./new.txt │ │ B/op │ B/op vs base │ SetMetaInfoFromMap-12 704.0 ± 0% 709.0 ± 0% +0.71% (p=0.002 n=6) │ ./old.txt │ ./new.txt │ │ allocs/op │ allocs/op vs base │ SetMetaInfoFromMap-12 3.000 ± 0% 3.000 ± 0% ~ (p=1.000 n=6) ¹ ¹ all samples are equal
bytedance · Oct 29, 2024 · 9141909 · 9141909
1 parent 780ca9e
commit 9141909
Show file tree

Hide file tree

Showing 4 changed files with 171 additions and 39 deletions.
diff --git a/.gitignore b/.gitignore
@@ -8,6 +8,10 @@
 # Test binary, built with `go test -c`
 *.test
 
+# Benchmark comparision
+old.txt
+new.txt
+
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
 

diff --git a/cloud/metainfo/info.go b/cloud/metainfo/info.go
@@ -31,6 +31,8 @@ const (
 	lenPP  = len(PrefixPersistent)
 	lenB   = len(PrefixBackward)
 	lenBD  = len(PrefixBackwardDownstream)
+
+	lenU = lenPTU - lenPT // UPSTREAM_
 )
 
 // **Using empty string as key or value is not support.**

diff --git a/cloud/metainfo/utils.go b/cloud/metainfo/utils.go
@@ -16,7 +16,9 @@ package metainfo
 
 import (
 	"context"
+	"fmt"
 	"strings"
+	"sync"
 )
 
 // HasMetaInfo detects whether the given context contains metainfo.
@@ -37,47 +39,14 @@ func SetMetaInfoFromMap(ctx context.Context, m map[string]string) context.Contex
 		// fast path
 		return newCtxFromMap(ctx, m)
 	}
-	// inherit from node
-	mapSize := len(m)
-	persistent := newKVStore(mapSize)
-	transient := newKVStore(mapSize)
-	stale := newKVStore(mapSize)
-	sliceToMap(nd.persistent, persistent)
-	sliceToMap(nd.transient, transient)
-	sliceToMap(nd.stale, stale)
-
-	// insert new kvs from m to node
-	for k, v := range m {
-		if len(k) == 0 || len(v) == 0 {
-			continue
-		}
-		switch {
-		case strings.HasPrefix(k, PrefixTransientUpstream):
-			if len(k) > lenPTU { // do not move this condition to the case statement to prevent a PTU matches PT
-				stale[k[lenPTU:]] = v
-			}
-		case strings.HasPrefix(k, PrefixTransient):
-			if len(k) > lenPT {
-				transient[k[lenPT:]] = v
-			}
-		case strings.HasPrefix(k, PrefixPersistent):
-			if len(k) > lenPP {
-				persistent[k[lenPP:]] = v
-			}
-		}
-	}
 
-	// return original ctx if no invalid key in map
-	if (persistent.size() + transient.size() + stale.size()) == 0 {
+	p := poolKVMerge.Get().(*kvMerge)
+	defer poolKVMerge.Put(p)
+	if p.Load(m) == 0 {
+		//  no new kv added?
 		return ctx
 	}
-
-	// make new node, and transfer map to list
-	nd = newNodeFromMaps(persistent, transient, stale)
-	persistent.recycle()
-	transient.recycle()
-	stale.recycle()
-	return withNode(ctx, nd)
+	return withNode(ctx, p.Merge(nd))
 }
 
 func newCtxFromMap(ctx context.Context, m map[string]string) context.Context {
@@ -145,3 +114,129 @@ func sliceToMap(slice []kv, kvs kvstore) {
 		kvs[kv.key] = kv.val
 	}
 }
+
+var poolKVMerge = sync.Pool{
+	New: func() interface{} {
+		p := &kvMerge{}
+		p.dup = make(map[string]bool, 8)
+		p.persistent = make([]kv, 0, 8)
+		p.transient = make([]kv, 0, 8)
+		p.stale = make([]kv, 0, 8)
+		return p
+	},
+}
+
+type kvMerge struct {
+	dup        map[string]bool
+	persistent []kv // PrefixPersistent
+	transient  []kv // PrefixTransient
+	stale      []kv // PrefixTransientUpstream
+}
+
+func (p *kvMerge) String() string {
+	return fmt.Sprintf("persistent:%v, transient:%v, stale:%v",
+		p.persistent, p.transient, p.stale)
+}
+
+func (p *kvMerge) resetdup() {
+	for k := range p.dup {
+		delete(p.dup, k)
+	}
+}
+
+func (p *kvMerge) Load(m map[string]string) int {
+	p.persistent = p.persistent[:0]
+	p.transient = p.transient[:0]
+	p.stale = p.stale[:0]
+	for k, v := range m {
+		if len(k) == 0 || len(v) == 0 {
+			continue
+		}
+		switch {
+		case strings.HasPrefix(k, PrefixTransient):
+			if len(k) <= lenPT {
+				continue
+			}
+			k = k[lenPT:]
+			if strings.HasPrefix(k, "UPSTREAM_") { // PrefixTransientUpstream {
+				if len(k) > lenU {
+					p.stale = append(p.stale, kv{key: k[lenU:], val: v})
+				}
+			} else {
+				p.transient = append(p.transient, kv{key: k, val: v})
+			}
+
+		case strings.HasPrefix(k, PrefixPersistent):
+			if len(k) > lenPP {
+				p.persistent = append(p.persistent, kv{key: k[lenPP:], val: v})
+			}
+		}
+	}
+	return len(p.stale) + len(p.transient) + len(p.persistent)
+}
+
+func (p *kvMerge) Merge(old *node) *node {
+	// this method assumes that
+	// keys in old.persistent, old.transient, and old.stale are unique
+
+	if len(p.persistent) != 0 {
+		p.resetdup()
+		for i := range p.persistent {
+			p.dup[p.persistent[i].key] = true
+		}
+		for j := range old.persistent {
+			if !p.dup[old.persistent[j].key] {
+				p.persistent = append(p.persistent, old.persistent[j])
+			}
+		}
+	}
+	if len(p.transient) != 0 {
+		p.resetdup()
+		for i := range p.transient {
+			p.dup[p.transient[i].key] = true
+		}
+		for j := range old.transient {
+			if !p.dup[old.transient[j].key] {
+				p.transient = append(p.transient, old.transient[j])
+			}
+		}
+	}
+
+	if len(p.stale) != 0 {
+		p.resetdup()
+		for i := range p.stale {
+			p.dup[p.stale[i].key] = true
+		}
+		for j := range old.stale {
+			if !p.dup[old.transient[j].key] {
+				p.stale = append(p.stale, old.stale[j])
+			}
+		}
+	}
+
+	// copy to ret
+
+	ret := &node{}
+	kvs := make([]kv, len(p.persistent)+len(p.transient)+len(p.stale))
+	if n := len(p.persistent); n != 0 {
+		copy(kvs, p.persistent)
+		ret.persistent = kvs[:n:n]
+		kvs = kvs[n:]
+	} else {
+		ret.persistent = old.persistent
+	}
+	if n := len(p.transient); n != 0 {
+		copy(kvs, p.transient)
+		ret.transient = kvs[:n:n]
+		kvs = kvs[n:]
+	} else {
+		ret.transient = old.transient
+	}
+	if len(p.stale) != 0 { // last one, use kvs directly
+		copy(kvs, p.stale)
+		ret.stale = kvs
+	} else {
+		ret.stale = old.stale
+	}
+	return ret
+}
diff --git a/cloud/metainfo/utils_test.go b/cloud/metainfo/utils_test.go
@@ -90,6 +90,34 @@ func TestSetMetaInfoFromMap(t *testing.T) {
 	assert(t, !ok5)
 	assert(t, ok6)
 	assert(t, v3 == "v6")
+
+	// Overwrites
+	m[k4] = "v4+"
+	m[k5] = "v5+"
+	m[k6] = "v6+"
+	ctx3 := metainfo.SetMetaInfoFromMap(ctx, m)
+	assert(t, ctx3 != ctx)
+	ctx = ctx3
+
+	t.Log(metainfo.GetAllValues(ctx))
+
+	v1, ok1 = metainfo.GetValue(ctx, "k4")
+	v2, ok2 = metainfo.GetValue(ctx, "k5")
+	_, ok3 = metainfo.GetValue(ctx, "k6")
+	assert(t, ok1)
+	assert(t, ok2)
+	assert(t, !ok3)
+	assert(t, v1 == "v4+")
+	assert(t, v2 == "v5+")
+
+	_, ok4 = metainfo.GetPersistentValue(ctx, "k4")
+	_, ok5 = metainfo.GetPersistentValue(ctx, "k5")
+	v3, ok6 = metainfo.GetPersistentValue(ctx, "k6")
+	assert(t, !ok4)
+	assert(t, !ok5)
+	assert(t, ok6)
+	assert(t, v3 == "v6+")
+
 }
 
 func TestSetMetaInfoFromMapKeepPreviousData(t *testing.T) {
@@ -151,9 +179,12 @@ func TestSaveMetaInfoToMap(t *testing.T) {
 func BenchmarkSetMetaInfoFromMap(b *testing.B) {
 	ctx := metainfo.WithPersistentValue(context.Background(), "key", "val")
 	m := map[string]string{}
-	for i := 0; i < 32; i++ {
+	for i := 0; i < 16; i++ {
 		m[fmt.Sprintf("key-%d", i)] = fmt.Sprintf("val-%d", i)
 	}
+	for i := 0; i < 16; i++ {
+		m[fmt.Sprintf("%s_key_%d", metainfo.PrefixPersistent, i)] = fmt.Sprintf("val_%d", i)
+	}
 	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {