github.com/glide-im/glide@v1.6.0/pkg/rpc/context.go (about)

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"github.com/smallnest/rpcx/share"
     6  )
     7  
     8  type ExtraContext struct {
     9  	context.Context
    10  }
    11  
    12  func NewContextFrom(c context.Context) *ExtraContext {
    13  	return &ExtraContext{c}
    14  }
    15  
    16  func NewContext() *ExtraContext {
    17  	return NewContextFrom(context.Background())
    18  }
    19  
    20  func (c *ExtraContext) PutReqExtra(k string, v string) *ExtraContext {
    21  	mate := c.Context.Value(share.ReqMetaDataKey)
    22  	if mate == nil {
    23  		mate = map[string]string{}
    24  		c.Context = context.WithValue(c.Context, share.ReqMetaDataKey, mate)
    25  	}
    26  	m := c.Context.Value(share.ReqMetaDataKey).(map[string]string)
    27  	m[k] = v
    28  	return c
    29  }
    30  
    31  func (c *ExtraContext) PutResExtra(k string, v string) *ExtraContext {
    32  	mate := c.Context.Value(share.ResMetaDataKey)
    33  	if mate == nil {
    34  		mate = map[string]string{}
    35  		c.Context = context.WithValue(c.Context, share.ResMetaDataKey, mate)
    36  	}
    37  	m := c.Context.Value(share.ResMetaDataKey).(map[string]string)
    38  	m[k] = v
    39  	return c
    40  }
    41  
    42  func (c *ExtraContext) GetReqExtra(k string) (string, bool) {
    43  	mate := c.Context.Value(share.ReqMetaDataKey)
    44  	if mate == nil {
    45  		return "", false
    46  	}
    47  	m := c.Context.Value(share.ReqMetaDataKey).(map[string]string)
    48  	v, ok := m[k]
    49  	return v, ok
    50  }
    51  
    52  func (c *ExtraContext) GetResExtra(k string) (string, bool) {
    53  	mate := c.Context.Value(share.ResMetaDataKey)
    54  	if mate == nil {
    55  		return "", false
    56  	}
    57  	m := c.Context.Value(share.ResMetaDataKey).(map[string]string)
    58  	v, ok := m[k]
    59  	return v, ok
    60  }