refactor: improve readability and make new node for ctx

joway · Apr 17, 2024 · e33c0e4 · e33c0e4
1 parent 93b0cda
commit e33c0e4
Show file tree

Hide file tree

Showing 5 changed files with 138 additions and 102 deletions.
diff --git a/cloud/metainfo/http.go b/cloud/metainfo/http.go
@@ -76,63 +76,73 @@ func FromHTTPHeader(ctx context.Context, h HTTPHeaderCarrier) context.Context {
 	if ctx == nil || h == nil {
 		return ctx
 	}
-	var persistent kvstore
-	var transient kvstore
+	nd := getNode(ctx)
+	if nd == nil || nd.size() == 0 {
+		return newCtxFromHTTPHeader(ctx, h)
+	}
 
 	// inherit from exist ctx node
-	nd := getNode(ctx)
-	inherit := nd != nil
-	if inherit {
-		// inherit from node
-		persistent = newKVStore()
-		transient = newKVStore()
-		sliceToMap(nd.persistent, persistent)
-		sliceToMap(nd.transient, transient)
-	} else {
-		// make new node
-		nd = &node{
-			persistent: make([]kv, 0, 16), // 32B * 16 = 512B
-			transient:  make([]kv, 0, 16),
-			stale:      make([]kv, 0, 16),
+	persistent := newKVStore()
+	transient := newKVStore()
+	sliceToMap(nd.persistent, persistent)
+	sliceToMap(nd.transient, transient)
+
+	// insert new kvs from http header
+	h.Visit(func(k, v string) {
+		if len(v) == 0 {
+			return
+		}
+		kk := strings.ToLower(k)
+		ln := len(kk)
+		if ln > lenHPT && strings.HasPrefix(kk, HTTPPrefixTransient) {
+			kk = HTTPHeaderToCGIVariable(kk[lenHPT:])
+			transient[kk] = v
+		} else if ln > lenHPP && strings.HasPrefix(kk, HTTPPrefixPersistent) {
+			kk = HTTPHeaderToCGIVariable(kk[lenHPP:])
+			persistent[kk] = v
 		}
+	})
+
+	// return original ctx if no invalid key in http header
+	if (persistent.size() + transient.size()) == 0 {
+		return ctx
 	}
 
+	// make new kvs
+	nd = newNodeFromMaps(persistent, transient, nil)
+	persistent.recycle()
+	transient.recycle()
+	ctx = withNode(ctx, nd)
+	return ctx
+}
+
+func newCtxFromHTTPHeader(ctx context.Context, h HTTPHeaderCarrier) context.Context {
+	nd := &node{
+		persistent: make([]kv, 0, 16), // 32B * 16 = 512B
+		transient:  make([]kv, 0, 16),
+		stale:      []kv{},
+	}
 	// insert new kvs from http header to node
 	h.Visit(func(k, v string) {
 		if len(v) == 0 {
 			return
 		}
-
 		kk := strings.ToLower(k)
 		ln := len(kk)
 		if ln > lenHPT && strings.HasPrefix(kk, HTTPPrefixTransient) {
 			kk = HTTPHeaderToCGIVariable(kk[lenHPT:])
-			if inherit {
-				transient[kk] = v
-			} else {
-				nd.transient = append(nd.transient, kv{key: kk, val: v})
-			}
+			nd.transient = append(nd.transient, kv{key: kk, val: v})
 		} else if ln > lenHPP && strings.HasPrefix(kk, HTTPPrefixPersistent) {
 			kk = HTTPHeaderToCGIVariable(kk[lenHPP:])
-			if inherit {
-				persistent[kk] = v
-			} else {
-				nd.persistent = append(nd.persistent, kv{key: kk, val: v})
-			}
+			nd.persistent = append(nd.persistent, kv{key: kk, val: v})
 		}
 	})
 
-	if nd.size() == 0 { // return original ctx if no invalid key in map
+	// return original ctx if no invalid key in http header
+	if nd.size() == 0 {
 		return ctx
 	}
-	if inherit {
-		nd.transient = transient.toList(nd.transient)
-		nd.persistent = persistent.toList(nd.persistent)
-		transient.recycle()
-		persistent.recycle()
-	}
-	ctx = withNode(ctx, nd)
-	return ctx
+	return withNode(ctx, nd)
 }
 
 // ToHTTPHeader writes all metainfo into the given HTTP header.

diff --git a/cloud/metainfo/http_test.go b/cloud/metainfo/http_test.go
@@ -87,7 +87,7 @@ func TestFromHTTPHeaderKeepPreviousData(t *testing.T) {
 	c1 := metainfo.FromHTTPHeader(c0, metainfo.HTTPHeader(h))
 	assert(t, c0 != c1)
 	vs := metainfo.GetAllValues(c1)
-	assert(t, len(vs) == 3)
+	assert(t, len(vs) == 3, len(vs))
 	assert(t, vs["tk"] == "tv" && vs["uk"] == "uv" && vs["XK"] == "xv")
 	vs = metainfo.GetAllPersistentValues(c1)
 	assert(t, len(vs) == 3)

diff --git a/cloud/metainfo/kv.go b/cloud/metainfo/kv.go
@@ -27,6 +27,33 @@ type kv struct {
 	val string
 }
 
+func newNodeFromMaps(persistent, transient, stale kvstore) *node {
+	ps, ts, sz := persistent.size(), transient.size(), stale.size()
+	// make slices together to reduce malloc cost
+	kvs := make([]kv, ps+ts+sz)
+	nd := new(node)
+	nd.persistent = kvs[:ps]
+	nd.transient = kvs[ps : ps+ts]
+	nd.stale = kvs[ps+ts:]
+
+	i := 0
+	for k, v := range persistent {
+		nd.persistent[i].key, nd.persistent[i].val = k, v
+		i++
+	}
+	i = 0
+	for k, v := range transient {
+		nd.transient[i].key, nd.transient[i].val = k, v
+		i++
+	}
+	i = 0
+	for k, v := range stale {
+		nd.stale[i].key, nd.stale[i].val = k, v
+		i++
+	}
+	return nd
+}
+
 type node struct {
 	persistent []kv
 	transient  []kv

diff --git a/cloud/metainfo/kvstore.go b/cloud/metainfo/kvstore.go
@@ -31,11 +31,8 @@ func newKVStore(size ...int) kvstore {
 	return kvs.(kvstore)
 }
 
-func (store kvstore) toList(kvs []kv) []kv {
-	for k, v := range store {
-		kvs = append(kvs, kv{key: k, val: v})
-	}
-	return kvs
+func (store kvstore) size() int {
+	return len(store)
 }
 
 func (store kvstore) recycle() {

diff --git a/cloud/metainfo/utils.go b/cloud/metainfo/utils.go
@@ -31,30 +31,20 @@ func SetMetaInfoFromMap(ctx context.Context, m map[string]string) context.Contex
 	if ctx == nil || len(m) == 0 {
 		return ctx
 	}
-	var persistent kvstore
-	var transient kvstore
-	var stale kvstore
 
-	// inherit from exist ctx node
-	mapSize := len(m)
 	nd := getNode(ctx)
-	inherit := nd != nil
-	if inherit {
-		// inherit from node
-		persistent = newKVStore(mapSize)
-		transient = newKVStore(mapSize)
-		stale = newKVStore(mapSize)
-		sliceToMap(nd.persistent, persistent)
-		sliceToMap(nd.transient, transient)
-		sliceToMap(nd.stale, stale)
-	} else {
-		// make new node
-		nd = &node{
-			persistent: make([]kv, 0, mapSize),
-			transient:  make([]kv, 0, mapSize),
-			stale:      make([]kv, 0, mapSize),
-		}
+	if nd == nil || nd.size() == 0 {
+		// 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 {
@@ -64,45 +54,67 @@ func SetMetaInfoFromMap(ctx context.Context, m map[string]string) context.Contex
 		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
-				if inherit {
-					stale[k[lenPTU:]] = v
-				} else {
-					nd.stale = append(nd.stale, kv{key: k[lenPTU:], val: v})
-				}
+				stale[k[lenPTU:]] = v
 			}
 		case strings.HasPrefix(k, PrefixTransient):
 			if len(k) > lenPT {
-				if inherit {
-					transient[k[lenPT:]] = v
-				} else {
-					nd.transient = append(nd.transient, kv{key: k[lenPT:], val: v})
-				}
+				transient[k[lenPT:]] = v
 			}
 		case strings.HasPrefix(k, PrefixPersistent):
 			if len(k) > lenPP {
-				if inherit {
-					persistent[k[lenPP:]] = v
-				} else {
-					nd.persistent = append(nd.persistent, kv{key: k[lenPP:], val: v})
-				}
+				persistent[k[lenPP:]] = v
 			}
 		}
 	}
 
-	if nd.size() == 0 { // return original ctx if no invalid key in map
+	// return original ctx if no invalid key in map
+	if (persistent.size() + transient.size() + stale.size()) == 0 {
 		return ctx
 	}
-	// transfer map to list
-	if inherit {
-		nd.stale = stale.toList(nd.stale)
-		nd.transient = transient.toList(nd.transient)
-		nd.persistent = persistent.toList(nd.persistent)
-		stale.recycle()
-		transient.recycle()
-		persistent.recycle()
+
+	// make new node, and transfer map to list
+	nd = newNodeFromMaps(persistent, transient, stale)
+	persistent.recycle()
+	transient.recycle()
+	stale.recycle()
+	return withNode(ctx, nd)
+}
+
+func newCtxFromMap(ctx context.Context, m map[string]string) context.Context {
+	// make new node
+	mapSize := len(m)
+	nd := &node{
+		persistent: make([]kv, 0, mapSize),
+		transient:  make([]kv, 0, mapSize),
+		stale:      make([]kv, 0, mapSize),
 	}
-	ctx = withNode(ctx, nd)
-	return ctx
+
+	// 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
+				nd.stale = append(nd.stale, kv{key: k[lenPTU:], val: v})
+			}
+		case strings.HasPrefix(k, PrefixTransient):
+			if len(k) > lenPT {
+				nd.transient = append(nd.transient, kv{key: k[lenPT:], val: v})
+			}
+		case strings.HasPrefix(k, PrefixPersistent):
+			if len(k) > lenPP {
+				nd.persistent = append(nd.persistent, kv{key: k[lenPP:], val: v})
+			}
+		}
+	}
+
+	// return original ctx if no invalid key in map
+	if nd.size() == 0 {
+		return ctx
+	}
+	return withNode(ctx, nd)
 }
 
 // SaveMetaInfoToMap set key-value pairs from ctx to m while filtering out transient-upstream data.
@@ -126,20 +138,10 @@ func SaveMetaInfoToMap(ctx context.Context, m map[string]string) {
 
 // sliceToMap converts a kv slice to map.
 func sliceToMap(slice []kv, kvs kvstore) {
-	for _, kv := range slice {
-		kvs[kv.key] = kv.val
-	}
-}
-
-// mapToSlice converts a map to a kv slice. If the map is empty, the return value will be nil.
-func mapToSlice(kvs kvstore) (slice []kv) {
-	size := len(kvs)
-	if size == 0 {
+	if len(slice) == 0 {
 		return
 	}
-	slice = make([]kv, 0, size)
-	for k, v := range kvs {
-		slice = append(slice, kv{key: k, val: v})
+	for _, kv := range slice {
+		kvs[kv.key] = kv.val
 	}
-	return
 }