rework metadata

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
2024-03-09 17:21:20 +03:00
parent 4b261083d6
commit 2f8c0d9f9d
12 changed files with 152 additions and 279 deletions
--- a/metadata/context.go
+++ b/metadata/context.go
@@ -17,11 +17,11 @@ func FromIncomingContext(ctx context.Context) (Metadata, bool) {
 	if ctx == nil {
 		return nil, false
 	}
-	md, ok := ctx.Value(mdIncomingKey{}).(*rawMetadata)
-	if !ok || md.md == nil {
+	md, ok := ctx.Value(mdIncomingKey{}).(Metadata)
+	if !ok || md == nil {
 		return nil, false
 	}
-	return md.md, ok
+	return md, ok
 }

 // FromOutgoingContext returns metadata from outgoing ctx
@@ -30,11 +30,11 @@ func FromOutgoingContext(ctx context.Context) (Metadata, bool) {
 	if ctx == nil {
 		return nil, false
 	}
-	md, ok := ctx.Value(mdOutgoingKey{}).(*rawMetadata)
-	if !ok || md.md == nil {
+	md, ok := ctx.Value(mdOutgoingKey{}).(Metadata)
+	if !ok || md == nil {
 		return nil, false
 	}
-	return md.md, ok
+	return md, ok
 }

 // FromContext returns metadata from the given context
@@ -43,11 +43,11 @@ func FromContext(ctx context.Context) (Metadata, bool) {
 	if ctx == nil {
 		return nil, false
 	}
-	md, ok := ctx.Value(mdKey{}).(*rawMetadata)
-	if !ok || md.md == nil {
+	md, ok := ctx.Value(mdKey{}).(Metadata)
+	if !ok || md == nil {
 		return nil, false
 	}
-	return md.md, ok
+	return md, ok
 }

 // NewContext creates a new context with the given metadata
@@ -55,45 +55,16 @@ func NewContext(ctx context.Context, md Metadata) context.Context {
 	if ctx == nil {
 		ctx = context.Background()
 	}
-	ctx = context.WithValue(ctx, mdKey{}, &rawMetadata{md})
-	ctx = context.WithValue(ctx, mdIncomingKey{}, &rawMetadata{})
-	ctx = context.WithValue(ctx, mdOutgoingKey{}, &rawMetadata{})
+	ctx = context.WithValue(ctx, mdKey{}, md)
 	return ctx
 }

-// SetOutgoingContext modify outgoing context with given metadata
-func SetOutgoingContext(ctx context.Context, md Metadata) bool {
-	if ctx == nil {
-		return false
-	}
-	if omd, ok := ctx.Value(mdOutgoingKey{}).(*rawMetadata); ok {
-		omd.md = md
-		return true
-	}
-	return false
-}
-
-// SetIncomingContext modify incoming context with given metadata
-func SetIncomingContext(ctx context.Context, md Metadata) bool {
-	if ctx == nil {
-		return false
-	}
-	if omd, ok := ctx.Value(mdIncomingKey{}).(*rawMetadata); ok {
-		omd.md = md
-		return true
-	}
-	return false
-}
-
 // NewIncomingContext creates a new context with incoming metadata attached
 func NewIncomingContext(ctx context.Context, md Metadata) context.Context {
 	if ctx == nil {
 		ctx = context.Background()
 	}
-	ctx = context.WithValue(ctx, mdIncomingKey{}, &rawMetadata{md})
-	if v, ok := ctx.Value(mdOutgoingKey{}).(*rawMetadata); !ok || v == nil {
-		ctx = context.WithValue(ctx, mdOutgoingKey{}, &rawMetadata{})
-	}
+	ctx = context.WithValue(ctx, mdIncomingKey{}, md)
 	return ctx
 }

@@ -102,41 +73,28 @@ func NewOutgoingContext(ctx context.Context, md Metadata) context.Context {
 	if ctx == nil {
 		ctx = context.Background()
 	}
-	ctx = context.WithValue(ctx, mdOutgoingKey{}, &rawMetadata{md})
-	if v, ok := ctx.Value(mdIncomingKey{}).(*rawMetadata); !ok || v == nil {
-		ctx = context.WithValue(ctx, mdIncomingKey{}, &rawMetadata{})
-	}
+	ctx = context.WithValue(ctx, mdOutgoingKey{}, md)
 	return ctx
 }

 // AppendOutgoingContext apends new md to context
 func AppendOutgoingContext(ctx context.Context, kv ...string) context.Context {
-	md, ok := Pairs(kv...)
-	if !ok {
-		return ctx
-	}
+	md := Pairs(kv...)
 	omd, ok := FromOutgoingContext(ctx)
 	if !ok {
 		return NewOutgoingContext(ctx, md)
 	}
-	for k, v := range md {
-		omd.Set(k, v)
-	}
-	return NewOutgoingContext(ctx, omd)
+	nmd := Merge(omd, md, true)
+	return NewOutgoingContext(ctx, nmd)
 }

 // AppendIncomingContext apends new md to context
 func AppendIncomingContext(ctx context.Context, kv ...string) context.Context {
-	md, ok := Pairs(kv...)
-	if !ok {
-		return ctx
-	}
+	md := Pairs(kv...)
 	omd, ok := FromIncomingContext(ctx)
 	if !ok {
 		return NewIncomingContext(ctx, md)
 	}
-	for k, v := range md {
-		omd.Set(k, v)
-	}
-	return NewIncomingContext(ctx, omd)
+	nmd := Merge(omd, md, true)
+	return NewIncomingContext(ctx, nmd)
 }
--- a/metadata/context_test.go
+++ b/metadata/context_test.go
@@ -24,7 +24,7 @@ func TestNewNilContext(t *testing.T) {
 }

 func TestFromContext(t *testing.T) {
-	ctx := context.WithValue(context.TODO(), mdKey{}, &rawMetadata{New(0)})
+	ctx := context.WithValue(context.TODO(), mdKey{}, New(0))

 	c, ok := FromContext(ctx)
 	if c == nil || !ok {
@@ -42,7 +42,7 @@ func TestNewContext(t *testing.T) {
 }

 func TestFromIncomingContext(t *testing.T) {
-	ctx := context.WithValue(context.TODO(), mdIncomingKey{}, &rawMetadata{New(0)})
+	ctx := context.WithValue(context.TODO(), mdIncomingKey{}, New(0))

 	c, ok := FromIncomingContext(ctx)
 	if c == nil || !ok {
@@ -51,7 +51,7 @@ func TestFromIncomingContext(t *testing.T) {
 }

 func TestFromOutgoingContext(t *testing.T) {
-	ctx := context.WithValue(context.TODO(), mdOutgoingKey{}, &rawMetadata{New(0)})
+	ctx := context.WithValue(context.TODO(), mdOutgoingKey{}, New(0))

 	c, ok := FromOutgoingContext(ctx)
 	if c == nil || !ok {
@@ -59,36 +59,6 @@ func TestFromOutgoingContext(t *testing.T) {
 	}
 }

-func TestSetIncomingContext(t *testing.T) {
-	md := New(1)
-	md.Set("key", "val")
-	ctx := context.WithValue(context.TODO(), mdIncomingKey{}, &rawMetadata{})
-	if !SetIncomingContext(ctx, md) {
-		t.Fatal("SetIncomingContext not works")
-	}
-	md, ok := FromIncomingContext(ctx)
-	if md == nil || !ok {
-		t.Fatal("SetIncomingContext not works")
-	} else if v, ok := md.Get("key"); !ok || v != "val" {
-		t.Fatal("SetIncomingContext not works")
-	}
-}
-
-func TestSetOutgoingContext(t *testing.T) {
-	md := New(1)
-	md.Set("key", "val")
-	ctx := context.WithValue(context.TODO(), mdOutgoingKey{}, &rawMetadata{})
-	if !SetOutgoingContext(ctx, md) {
-		t.Fatal("SetOutgoingContext not works")
-	}
-	md, ok := FromOutgoingContext(ctx)
-	if md == nil || !ok {
-		t.Fatal("SetOutgoingContext not works")
-	} else if v, ok := md.Get("key"); !ok || v != "val" {
-		t.Fatal("SetOutgoingContext not works")
-	}
-}
-
 func TestNewIncomingContext(t *testing.T) {
 	md := New(1)
 	md.Set("key", "val")
--- a/metadata/metadata.go
+++ b/metadata/metadata.go
@@ -1,9 +1,9 @@
 // Package metadata is a way of defining message headers
-package metadata // import "go.unistack.org/micro/v4/metadata"
+package metadata

 import (
 	"net/textproto"
-	"sort"
+	"strings"
 )

 var (
@@ -24,47 +24,7 @@ var (
 )

 // Metadata is our way of representing request headers internally.
-// They're used at the RPC level and translate back and forth
-// from Transport headers.
-type Metadata map[string]string
-
-type rawMetadata struct {
-	md Metadata
-}
-
-// defaultMetadataSize used when need to init new Metadata
-var defaultMetadataSize = 2
-
-// Iterator used to iterate over metadata with order
-type Iterator struct {
-	md   Metadata
-	keys []string
-	cur  int
-	cnt  int
-}
-
-// Next advance iterator to next element
-func (iter *Iterator) Next(k, v *string) bool {
-	if iter.cur+1 > iter.cnt {
-		return false
-	}
-
-	*k = iter.keys[iter.cur]
-	*v = iter.md[*k]
-	iter.cur++
-	return true
-}
-
-// Iterator returns the itarator for metadata in sorted order
-func (md Metadata) Iterator() *Iterator {
-	iter := &Iterator{md: md, cnt: len(md)}
-	iter.keys = make([]string, 0, iter.cnt)
-	for k := range md {
-		iter.keys = append(iter.keys, k)
-	}
-	sort.Strings(iter.keys)
-	return iter
-}
+type Metadata map[string][]string

 // Get returns value from metadata by key
 func (md Metadata) Get(key string) (string, bool) {
@@ -74,7 +34,7 @@ func (md Metadata) Get(key string) (string, bool) {
 		// slow path
 		val, ok = md[textproto.CanonicalMIMEHeaderKey(key)]
 	}
-	return val, ok
+	return strings.Join(val, ","), ok
 }

 // Set is used to store value in metadata
@@ -83,10 +43,19 @@ func (md Metadata) Set(kv ...string) {
 		kv = kv[:len(kv)-1]
 	}
 	for idx := 0; idx < len(kv); idx += 2 {
-		md[textproto.CanonicalMIMEHeaderKey(kv[idx])] = kv[idx+1]
+		md[textproto.CanonicalMIMEHeaderKey(kv[idx])] = []string{kv[idx+1]}
 	}
 }

+// Append is used to append value in metadata
+func (md Metadata) Append(k string, v ...string) {
+	if len(v) == 0 {
+		return
+	}
+	k = textproto.CanonicalMIMEHeaderKey(k)
+	md[k] = append(md[k], v...)
+}
+
 // Del is used to remove value from metadata
 func (md Metadata) Del(keys ...string) {
 	for _, key := range keys {
@@ -99,9 +68,9 @@ func (md Metadata) Del(keys ...string) {

 // Copy makes a copy of the metadata
 func Copy(md Metadata) Metadata {
-	nmd := New(len(md))
-	for key, val := range md {
-		nmd.Set(key, val)
+	nmd := make(Metadata, len(md))
+	for k, v := range md {
+		nmd[k] = v
 	}
 	return nmd
 }
@@ -109,35 +78,40 @@ func Copy(md Metadata) Metadata {
 // New return new sized metadata
 func New(size int) Metadata {
 	if size == 0 {
-		size = defaultMetadataSize
+		size = 2
 	}
 	return make(Metadata, size)
 }

 // Merge merges metadata to existing metadata, overwriting if specified
 func Merge(omd Metadata, mmd Metadata, overwrite bool) Metadata {
-	var ok bool
 	nmd := Copy(omd)
 	for key, val := range mmd {
-		_, ok = nmd[key]
+		nval, ok := nmd[key]
 		switch {
+		case ok && overwrite:
+			nmd[key] = nval
+			continue
 		case ok && !overwrite:
 			continue
-		case val != "":
-			nmd.Set(key, val)
-		case ok && val == "":
-			nmd.Del(key)
+		case !ok:
+			for _, v := range val {
+				if v != "" {
+					nval = append(nval, v)
+				}
+			}
+			nmd[key] = nval
 		}
 	}
 	return nmd
 }

 // Pairs from which metadata created
-func Pairs(kv ...string) (Metadata, bool) {
+func Pairs(kv ...string) Metadata {
 	if len(kv)%2 == 1 {
-		return nil, false
+		kv = kv[:len(kv)-1]
 	}
-	md := New(len(kv) / 2)
+	md := make(Metadata, len(kv)/2)
 	md.Set(kv...)
-	return md, true
+	return md
 }
--- a/metadata/metadata_test.go
+++ b/metadata/metadata_test.go
@@ -33,20 +33,25 @@ func TestAppend(t *testing.T) {
 }

 func TestPairs(t *testing.T) {
-	md, ok := Pairs("key1", "val1", "key2", "val2")
-	if !ok {
-		t.Fatal("odd number of kv")
-	}
-	if _, ok = md.Get("key1"); !ok {
+	md := Pairs("key1", "val1", "key2", "val2")
+
+	if _, ok := md.Get("key1"); !ok {
 		t.Fatal("key1 not found")
 	}
 }

-func testCtx(ctx context.Context) {
-	md := New(2)
-	md.Set("Key1", "Val1_new")
-	md.Set("Key3", "Val3")
-	SetOutgoingContext(ctx, md)
+func testIncomingCtx(ctx context.Context) {
+	if md, ok := FromIncomingContext(ctx); ok && md != nil {
+		md.Set("Key1", "Val1_new")
+		md.Set("Key3", "Val3")
+	}
+}
+
+func testOutgoingCtx(ctx context.Context) {
+	if md, ok := FromOutgoingContext(ctx); ok && md != nil {
+		md.Set("Key1", "Val1_new")
+		md.Set("Key3", "Val3")
+	}
 }

 func TestPassing(t *testing.T) {
@@ -55,8 +60,8 @@ func TestPassing(t *testing.T) {
 	md1.Set("Key1", "Val1")
 	md1.Set("Key2", "Val2")

-	ctx = NewIncomingContext(ctx, md1)
-	testCtx(ctx)
+	ctx = NewOutgoingContext(ctx, md1)
+	testOutgoingCtx(ctx)
 	md, ok := FromOutgoingContext(ctx)
 	if !ok {
 		t.Fatalf("missing metadata from outgoing context")
@@ -68,10 +73,10 @@ func TestPassing(t *testing.T) {

 func TestMerge(t *testing.T) {
 	omd := Metadata{
-		"key1": "val1",
+		"key1": []string{"val1"},
 	}
 	mmd := Metadata{
-		"key2": "val2",
+		"key2": []string{"val2"},
 	}

 	nmd := Merge(omd, mmd, true)
@@ -80,21 +85,6 @@ func TestMerge(t *testing.T) {
 	}
 }

-func TestIterator(t *testing.T) {
-	md := Metadata{
-		"1Last":   "last",
-		"2First":  "first",
-		"3Second": "second",
-	}
-
-	iter := md.Iterator()
-	var k, v string
-
-	for iter.Next(&k, &v) {
-		// fmt.Printf("k: %s, v: %s\n", k, v)
-	}
-}
-
 func TestMedataCanonicalKey(t *testing.T) {
 	md := New(1)
 	md.Set("x-request-id", "12345")
@@ -134,10 +124,7 @@ func TestMetadataSet(t *testing.T) {
 }

 func TestMetadataDelete(t *testing.T) {
-	md := Metadata{
-		"Foo": "bar",
-		"Baz": "empty",
-	}
+	md := Pairs("Foo", "bar", "Baz", "empty")

 	md.Del("Baz")
 	_, ok := md.Get("Baz")
@@ -156,24 +143,19 @@ func TestNilContext(t *testing.T) {
 }

 func TestMetadataCopy(t *testing.T) {
-	md := Metadata{
-		"Foo": "bar",
-		"Bar": "baz",
-	}
+	md := Pairs("Foo", "bar", "Bar", "baz")

 	cp := Copy(md)

 	for k, v := range md {
-		if cv := cp[k]; cv != v {
+		if cv := cp[k]; len(cv) != len(v) {
 			t.Fatalf("Got %s:%s for %s:%s", k, cv, k, v)
 		}
 	}
 }

 func TestMetadataContext(t *testing.T) {
-	md := Metadata{
-		"Foo": "bar",
-	}
+	md := Pairs("Foo", "bar")

 	ctx := NewContext(context.TODO(), md)

@@ -182,7 +164,7 @@ func TestMetadataContext(t *testing.T) {
 		t.Errorf("Unexpected error retrieving metadata, got %t", ok)
 	}

-	if emd["Foo"] != md["Foo"] {
+	if len(emd["Foo"]) != len(md["Foo"]) {
 		t.Errorf("Expected key: %s val: %s, got key: %s val: %s", "Foo", md["Foo"], "Foo", emd["Foo"])
 	}