Merge 5b8f883 into 751ab03

core/handler/tracing_handler.go

@@ -251,9 +251,12 @@ func (t *TracingConsumerHandler) Name() string {
 
 func (t *TracingConsumerHandler) getTracer(i *invocation.Invocation) opentracing.Tracer {
 	caller := common.DefaultValue
-	if c, ok := i.Metadata[common.CallerKey].(string); ok && c != "" {
-		caller = c + ":" + runtime.HostName
+	if i.Metadata != nil {
+		if c, ok := i.Metadata[common.CallerKey].(string); ok && c != "" {
+			caller = c + ":" + runtime.HostName
+		}
 	}
+
 	return tracing.GetTracer(caller)
 }
 

core/invocation/invocation.go

@@ -36,19 +36,20 @@ type Invocation struct {
 	Args               interface{}
 	URLPathFormat      string
 	Reply              interface{}
-	Ctx                context.Context        //ctx can save protocol header
+	Ctx                context.Context        //ctx can save protocol headers
 	Metadata           map[string]interface{} //local scope data
 	RouteTags          utiltags.Tags          //route tags is decided in router handler
 	Strategy           string                 //load balancing strategy
 	Filters            []string
 }
 
-// CreateConsumerInvocation create invocation
-func CreateConsumerInvocation() *Invocation {
-	return &Invocation{
+// New create invocation
+func New(ctx context.Context) *Invocation {
+	inv := &Invocation{
 		SourceServiceID: runtime.ServiceID,
-		Metadata:        make(map[string]interface{}),
+		Ctx:             ctx,
 	}
+	return inv
 }
 
 //GetSessionID return session id
@@ -65,5 +66,23 @@ func (inv *Invocation) SetSessionID(value string) {
 
 //SetMetadata local scope params
 func (inv *Invocation) SetMetadata(key string, value interface{}) {
+	if inv.Metadata == nil {
+		inv.Metadata = make(map[string]interface{})
+	}
 	inv.Metadata[key] = value
 }
+
+//SetHeader set headers, the client and server plugins should use them in protocol headers
+func (inv *Invocation) SetHeader(k, v string) {
+	if inv.Ctx.Value(common.ContextHeaderKey{}) == nil {
+		inv.Ctx = context.WithValue(inv.Ctx, common.ContextHeaderKey{}, map[string]string{})
+	}
+	m := inv.Ctx.Value(common.ContextHeaderKey{}).(map[string]string)
+
+	m[k] = v
+}
+
+//Headers return a map that protocol plugin should deliver in transport
+func (inv *Invocation) Headers() map[string]string {
+	return inv.Ctx.Value(common.ContextHeaderKey{}).(map[string]string)
+}

core/invocation/invocation_test.go

@@ -3,13 +3,21 @@ package invocation_test
 import (
 	"context"
 	"github.com/ServiceComb/go-chassis/core/invocation"
+	"github.com/stretchr/testify/assert"
 	"testing"
 )
 
 func TestChain(t *testing.T) {
 	var inv = new(invocation.Invocation)
 	inv.Endpoint = "1.2.3.4"
 }
+func TestInvocation_Headers(t *testing.T) {
+	inv := invocation.New(context.TODO())
+	inv.SetMetadata("a", "1")
+	inv.SetHeader("asd", "123")
+	assert.Equal(t, "123", inv.Headers()["asd"])
+	assert.Equal(t, "1", inv.Metadata["a"])
+}
 
 /*
 

core/options.go

@@ -137,6 +137,9 @@ func wrapInvocationWithOpts(i *invocation.Invocation, opts InvokeOptions) {
 	i.Protocol = opts.Protocol
 	i.Strategy = opts.StrategyFunc
 	i.Filters = opts.Filters
-	i.Metadata = opts.Metadata
+	if opts.Metadata != nil {
+		i.Metadata = opts.Metadata
+	}
+
 	i.RouteTags = opts.RouteTags
 }

core/rest_invoker.go

@@ -40,21 +40,17 @@ func (ri *RestInvoker) ContextDo(ctx context.Context, req *rest.Request, options
 
 	resp := rest.NewResponse()
 
-	inv := invocation.CreateConsumerInvocation()
+	inv := invocation.New(ctx)
 	wrapInvocationWithOpts(inv, opts)
 	inv.MicroServiceName = req.GetRequest().Host
 	// TODO load from openAPI schema
 	// inv.SchemaID = schemaID
 	// inv.OperationID = operationID
 	inv.Args = req
 	inv.Reply = resp
-	inv.Ctx = ctx
 	inv.URLPathFormat = req.Req.URL.Path
 
-	if inv.Metadata == nil {
-		inv.Metadata = make(map[string]interface{})
-	}
-	inv.Metadata[common.RestMethod] = req.GetMethod()
+	inv.SetMetadata(common.RestMethod, req.GetMethod())
 
 	err := ri.invoke(inv)
 	return resp, err

core/rpc_invoker.go

@@ -37,13 +37,12 @@ func (ri *RPCInvoker) Invoke(ctx context.Context, microServiceName, schemaID, op
 		opts.Protocol = common.ProtocolHighway
 	}
 
-	i := invocation.CreateConsumerInvocation()
+	i := invocation.New(ctx)
 	wrapInvocationWithOpts(i, opts)
 	i.MicroServiceName = microServiceName
 	i.SchemaID = schemaID
 	i.OperationID = operationID
 	i.Args = arg
 	i.Reply = reply
-	i.Ctx = ctx
 	return ri.invoke(i)
 }

server/restful/context.go

@@ -95,3 +95,8 @@ func (bs *Context) ReadRestfulRequest() *restful.Request {
 func (bs *Context) ReadResponseWriter() http.ResponseWriter {
 	return bs.resp.ResponseWriter
 }
+
+//ReadRestfulResponse return a native go-restful Response
+func (bs *Context) ReadRestfulResponse() *restful.Response {
+	return bs.resp
+}

server/restful/restful_server.go

@@ -140,7 +140,8 @@ func (r *restfulServer) Register(schema interface{}, options ...server.RegisterO
 				lager.Logger.Errorf(err, "transfer http request to invocation failed")
 				return
 			}
-			bs := NewBaseServer(context.TODO())
+			//give inv.ctx to user handlers, it may include headers
+			bs := NewBaseServer(inv.Ctx)
 			bs.req = req
 			bs.resp = rep
 			c.Next(inv, func(ir *invocation.Response) error {