broker/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Broker, bool) {
 	return c, ok
 }
 
+// MustContext returns broker from passed context
+func MustContext(ctx context.Context) Broker {
+	b, ok := FromContext(ctx)
+	if !ok {
+		panic("missing broker")
+	}
+	return b
+}
+
 // NewContext savess broker in context
 func NewContext(ctx context.Context, s Broker) context.Context {
 	if ctx == nil {

client/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Client, bool) {
 	return c, ok
 }
 
+// MustContext get client from context
+func MustContext(ctx context.Context) Client {
+	c, ok := FromContext(ctx)
+	if !ok {
+		panic("missing client")
+	}
+	return c
+}
+
 // NewContext put client in context
 func NewContext(ctx context.Context, c Client) context.Context {
 	if ctx == nil {

codec/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Codec, bool) {
 	return c, ok
 }
 
+// MustContext returns codec from context
+func MustContext(ctx context.Context) Codec {
+	c, ok := FromContext(ctx)
+	if !ok {
+		panic("missing codec")
+	}
+	return c
+}
+
 // NewContext put codec in context
 func NewContext(ctx context.Context, c Codec) context.Context {
 	if ctx == nil {

config/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Config, bool) {
 	return c, ok
 }
 
+// MustContext returns store from context
+func MustContext(ctx context.Context) Config {
+	c, ok := FromContext(ctx)
+	if !ok {
+		panic("missing config")
+	}
+	return c
+}
+
 // NewContext put store in context
 func NewContext(ctx context.Context, c Config) context.Context {
 	if ctx == nil {

flow/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Flow, bool) {
 	return c, ok
 }
 
+// MustContext returns Flow from context
+func MustContext(ctx context.Context) Flow {
+	f, ok := FromContext(ctx)
+	if !ok {
+		panic("missing flow")
+	}
+	return f
+}
+
 // NewContext stores Flow to context
 func NewContext(ctx context.Context, f Flow) context.Context {
 	if ctx == nil {

go.mod

@@ -31,6 +31,7 @@ require (
 	github.com/opencontainers/runtime-spec v1.2.0 // indirect
 	github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/rogpeppe/go-internal v1.13.1 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/stretchr/testify v1.10.0 // indirect
 	go.uber.org/goleak v1.3.0 // indirect

go.sum

@@ -57,8 +57,7 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
 github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
-github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
-github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
+github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
 github.com/silas/dag v0.0.0-20220518035006-a7e85ada93c5 h1:G/FZtUu7a6NTWl3KUHMV9jkLAh/Rvtf03NWMHaEDl+E=
 github.com/silas/dag v0.0.0-20220518035006-a7e85ada93c5/go.mod h1:7RTUFBdIRC9nZ7/3RyRNH1bdqIShrDejd1YbLwgPS+I=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=

logger/context.go

@@ -4,17 +4,6 @@ import "context"
 
 type loggerKey struct{}
 
-// MustContext returns logger from passed context or DefaultLogger if empty
-func MustContext(ctx context.Context) Logger {
-	if ctx == nil {
-		return DefaultLogger.Clone()
-	}
-	if l, ok := ctx.Value(loggerKey{}).(Logger); ok && l != nil {
-		return l
-	}
-	return DefaultLogger.Clone()
-}
-
 // FromContext returns logger from passed context
 func FromContext(ctx context.Context) (Logger, bool) {
 	if ctx == nil {
@@ -24,6 +13,15 @@ func FromContext(ctx context.Context) (Logger, bool) {
 	return l, ok
 }
 
+// MustContext returns logger from passed context or DefaultLogger if empty
+func MustContext(ctx context.Context) Logger {
+	l, ok := FromContext(ctx)
+	if !ok {
+		panic("missing logger")
+	}
+	return l
+}
+
 // NewContext stores logger into passed context
 func NewContext(ctx context.Context, l Logger) context.Context {
 	if ctx == nil {

metadata/context.go

@@ -11,34 +11,6 @@ type (
 	mdKey         struct{}
 )
 
-// MustIncomingContext returns metadata from incoming ctx
-// returned metadata shoud not be modified or race condition happens.
-// If metadata not exists panics.
-func MustIncomingContext(ctx context.Context) Metadata {
-	if ctx == nil {
-		panic("missing metadata")
-	}
-	md, ok := ctx.Value(mdIncomingKey{}).(*rawMetadata)
-	if !ok {
-		panic("missing metadata")
-	}
-	return md.md
-}
-
-// MustOutgoingContext returns metadata from outgoing ctx
-// returned metadata shoud not be modified or race condition happens.
-// If metadata not exists panics.
-func MustOutgoingContext(ctx context.Context) Metadata {
-	if ctx == nil {
-		panic("missing metadata")
-	}
-	md, ok := ctx.Value(mdOutgoingKey{}).(*rawMetadata)
-	if !ok {
-		panic("missing metadata")
-	}
-	return md.md
-}
-
 // FromIncomingContext returns metadata from incoming ctx
 // returned metadata shoud not be modified or race condition happens
 func FromIncomingContext(ctx context.Context) (Metadata, bool) {
@@ -52,6 +24,17 @@ func FromIncomingContext(ctx context.Context) (Metadata, bool) {
 	return md.md, ok
 }
 
+// MustIncomingContext returns metadata from incoming ctx
+// returned metadata shoud not be modified or race condition happens.
+// If metadata not exists panics.
+func MustIncomingContext(ctx context.Context) Metadata {
+	md, ok := FromIncomingContext(ctx)
+	if !ok {
+		panic("missing metadata")
+	}
+	return md
+}
+
 // FromOutgoingContext returns metadata from outgoing ctx
 // returned metadata shoud not be modified or race condition happens
 func FromOutgoingContext(ctx context.Context) (Metadata, bool) {
@@ -65,6 +48,17 @@ func FromOutgoingContext(ctx context.Context) (Metadata, bool) {
 	return md.md, ok
 }
 
+// MustOutgoingContext returns metadata from outgoing ctx
+// returned metadata shoud not be modified or race condition happens.
+// If metadata not exists panics.
+func MustOutgoingContext(ctx context.Context) Metadata {
+	md, ok := FromOutgoingContext(ctx)
+	if !ok {
+		panic("missing metadata")
+	}
+	return md
+}
+
 // FromContext returns metadata from the given context
 // returned metadata shoud not be modified or race condition happens
 func FromContext(ctx context.Context) (Metadata, bool) {
@@ -78,6 +72,16 @@ func FromContext(ctx context.Context) (Metadata, bool) {
 	return md.md, ok
 }
 
+// MustContext returns metadata from the given context
+// returned metadata shoud not be modified or race condition happens
+func MustContext(ctx context.Context) Metadata {
+	md, ok := FromContext(ctx)
+	if !ok {
+		panic("missing metadata")
+	}
+	return md
+}
+
 // NewContext creates a new context with the given metadata
 func NewContext(ctx context.Context, md Metadata) context.Context {
 	if ctx == nil {
@@ -139,7 +143,7 @@ func AppendOutgoingContext(ctx context.Context, kv ...string) context.Context {
 	for k, v := range md {
 		omd.Set(k, v)
 	}
-	return NewOutgoingContext(ctx, omd)
+	return ctx
 }
 
 // AppendIncomingContext apends new md to context
@@ -155,5 +159,21 @@ func AppendIncomingContext(ctx context.Context, kv ...string) context.Context {
 	for k, v := range md {
 		omd.Set(k, v)
 	}
-	return NewIncomingContext(ctx, omd)
+	return ctx
+}
+
+// AppendContext apends new md to context
+func AppendContext(ctx context.Context, kv ...string) context.Context {
+	md, ok := Pairs(kv...)
+	if !ok {
+		return ctx
+	}
+	omd, ok := FromContext(ctx)
+	if !ok {
+		return NewContext(ctx, md)
+	}
+	for k, v := range md {
+		omd.Set(k, v)
+	}
+	return ctx
 }

metadata/metadata_test.go

@@ -32,8 +32,8 @@ func TestMultipleUsage(t *testing.T) {
 		m.Del("key1_3")
 		return ctx
 	}(ctx)
-	t.Logf("imd %#+v", imd)
-	t.Logf("omd %#+v", omd)
+	_ = imd
+	_ = omd
 }
 
 func TestMetadataSetMultiple(t *testing.T) {

meter/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Meter, bool) {
 	return c, ok
 }
 
+// MustContext get meter from context
+func MustContext(ctx context.Context) Meter {
+	m, ok := FromContext(ctx)
+	if !ok {
+		panic("missing meter")
+	}
+	return m
+}
+
 // NewContext put meter in context
 func NewContext(ctx context.Context, c Meter) context.Context {
 	if ctx == nil {

register/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Register, bool) {
 	return c, ok
 }
 
+// MustContext get register from context
+func MustContext(ctx context.Context) Register {
+	r, ok := FromContext(ctx)
+	if !ok {
+		panic("missing register")
+	}
+	return r
+}
+
 // NewContext put register in context
 func NewContext(ctx context.Context, c Register) context.Context {
 	if ctx == nil {

router/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Router, bool) {
 	return c, ok
 }
 
+// MustContext get router from context
+func MustContext(ctx context.Context) Router {
+	r, ok := FromContext(ctx)
+	if !ok {
+		panic("missing router")
+	}
+	return r
+}
+
 // NewContext put router in context
 func NewContext(ctx context.Context, c Router) context.Context {
 	if ctx == nil {

server/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Server, bool) {
 	return c, ok
 }
 
+// MustContext returns Server from context
+func MustContext(ctx context.Context) Server {
+	s, ok := FromContext(ctx)
+	if !ok {
+		panic("missing server")
+	}
+	return s
+}
+
 // NewContext stores Server to context
 func NewContext(ctx context.Context, s Server) context.Context {
 	if ctx == nil {

store/context.go

@@ -15,6 +15,15 @@ func FromContext(ctx context.Context) (Store, bool) {
 	return c, ok
 }
 
+// MustContext get store from context
+func MustContext(ctx context.Context) Store {
+	s, ok := FromContext(ctx)
+	if !ok {
+		panic("missing store")
+	}
+	return s
+}
+
 // NewContext put store in context
 func NewContext(ctx context.Context, c Store) context.Context {
 	if ctx == nil {

tracer/context.go

@@ -18,6 +18,15 @@ func FromContext(ctx context.Context) (Tracer, bool) {
 	return nil, false
 }
 
+// MustContext returns a tracer from context
+func MustContext(ctx context.Context) Tracer {
+	t, ok := FromContext(ctx)
+	if !ok {
+		panic("missing tracer")
+	}
+	return t
+}
+
 // NewContext saves the tracer in the context
 func NewContext(ctx context.Context, tracer Tracer) context.Context {
 	if ctx == nil {