From 8a2c4c511e27445a64ee473380644811baffdbd5 Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Thu, 21 Jan 2021 18:35:31 +0300
Subject: [PATCH] metadata: add iterator method

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 metadata/context.go       | 49 ++++++++++++++++++++++++++
 metadata/metadata.go      | 73 ++++++++++++++++-----------------------
 metadata/metadata_test.go | 16 +++++++++
 3 files changed, 95 insertions(+), 43 deletions(-)
 create mode 100644 metadata/context.go

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")