diff --git a/metadata/context.go b/metadata/context.go new file mode 100644 index 00000000..06d126ce --- /dev/null +++ b/metadata/context.go @@ -0,0 +1,49 @@ +// Package metadata is a way of defining message headers +package metadata + +import ( + "context" +) + +// FromContext returns metadata from the given context +func FromContext(ctx context.Context) (Metadata, bool) { + if ctx == nil { + return nil, false + } + md, ok := ctx.Value(metadataKey{}).(Metadata) + if !ok { + return nil, ok + } + nmd := Copy(md) + return nmd, ok +} + +// NewContext creates a new context with the given metadata +func NewContext(ctx context.Context, md Metadata) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, metadataKey{}, Copy(md)) +} + +// MergeContext merges metadata to existing metadata, overwriting if specified +func MergeContext(ctx context.Context, pmd Metadata, overwrite bool) context.Context { + if ctx == nil { + ctx = context.Background() + } + md, ok := FromContext(ctx) + if !ok { + return context.WithValue(ctx, metadataKey{}, Copy(pmd)) + } + nmd := Copy(md) + for key, val := range pmd { + if _, ok := nmd[key]; ok && !overwrite { + // skip + } else if val != "" { + nmd.Set(key, val) + } else { + nmd.Del(key) + } + } + return context.WithValue(ctx, metadataKey{}, nmd) +} diff --git a/metadata/metadata.go b/metadata/metadata.go index 1b467247..528c03a8 100644 --- a/metadata/metadata.go +++ b/metadata/metadata.go @@ -4,6 +4,7 @@ package metadata import ( "context" "net/textproto" + "sort" ) type metadataKey struct{} @@ -18,6 +19,35 @@ var ( DefaultMetadataSize = 6 ) +type Iterator struct { + cur int + cnt int + keys []string + md Metadata +} + +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 +} + +// Iterate returns run user func with map key, val sorted by key +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 +} + // Get returns value from metadata by key func (md Metadata) Get(key string) (string, bool) { // fast path @@ -83,19 +113,6 @@ func Get(ctx context.Context, key string) (string, bool) { return md.Get(key) } -// FromContext returns metadata from the given context -func FromContext(ctx context.Context) (Metadata, bool) { - if ctx == nil { - return nil, false - } - md, ok := ctx.Value(metadataKey{}).(Metadata) - if !ok { - return nil, ok - } - nmd := Copy(md) - return nmd, ok -} - // New return new sized metadata func New(size int) Metadata { if size == 0 { @@ -103,33 +120,3 @@ func New(size int) Metadata { } return make(Metadata, size) } - -// NewContext creates a new context with the given metadata -func NewContext(ctx context.Context, md Metadata) context.Context { - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, metadataKey{}, Copy(md)) -} - -// MergeContext merges metadata to existing metadata, overwriting if specified -func MergeContext(ctx context.Context, pmd Metadata, overwrite bool) context.Context { - if ctx == nil { - ctx = context.Background() - } - md, ok := FromContext(ctx) - if !ok { - return context.WithValue(ctx, metadataKey{}, Copy(pmd)) - } - nmd := Copy(md) - for key, val := range pmd { - if _, ok := nmd[key]; ok && !overwrite { - // skip - } else if val != "" { - nmd.Set(key, val) - } else { - nmd.Del(key) - } - } - return context.WithValue(ctx, metadataKey{}, nmd) -} diff --git a/metadata/metadata_test.go b/metadata/metadata_test.go index 56f56e54..94d51bb2 100644 --- a/metadata/metadata_test.go +++ b/metadata/metadata_test.go @@ -2,10 +2,26 @@ package metadata import ( "context" + "fmt" "reflect" "testing" ) +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) { ctx := Set(context.TODO(), "x-request-id", "12345") v, ok := Get(ctx, "x-request-id")