github.com/songzhibin97/gkit@v1.2.13/trace/transport.go (about)

     1  package trace
     2  
     3  import (
     4  	"context"
     5  	"net/url"
     6  )
     7  
     8  // package transport 链路追踪
     9  
    10  const (
    11  	KindGRPC Kind = "GRPC"
    12  	KindHTTP Kind = "HTTP"
    13  )
    14  
    15  type (
    16  	Kind string
    17  
    18  	// Server 链路追踪的服务
    19  	Server interface {
    20  		Start(ctx context.Context) error
    21  		Shutdown(ctx context.Context) error
    22  	}
    23  
    24  	// EndPointer 注册点
    25  	EndPointer interface {
    26  		Endpoint() (*url.URL, error)
    27  	}
    28  
    29  	// Header 抽象获取协议头部的动作行为
    30  	Header interface {
    31  		Get(key string) string
    32  		Set(key string, value string)
    33  		Keys() []string
    34  	}
    35  
    36  	// Transporter 链路追踪的上下文
    37  	Transporter interface {
    38  		// Kind 返回 KindGRPC or KindHTTP 用于区分协议调用
    39  		Kind() Kind
    40  
    41  		// Endpoint
    42  		// Server Transporter: grpc://127.0.0.1:9000
    43  		// Client Transporter: discovery://provider-demo
    44  		Endpoint() string
    45  
    46  		// Operation protobuf 行为
    47  		Operation() string
    48  
    49  		// RequestHeader 返回请求头
    50  		RequestHeader() Header
    51  
    52  		// ResponseHeader 返回响应头
    53  		ResponseHeader() Header
    54  	}
    55  )
    56  
    57  func (k Kind) String() string { return string(k) }
    58  
    59  // transportKey 保证context.WithValue唯一性
    60  type serverTransportKey struct{}
    61  type clientTransportKey struct{}
    62  
    63  // NewServerTransportContext 新建服务端链路上下文
    64  func NewServerTransportContext(ctx context.Context, tp Transporter) context.Context {
    65  	return context.WithValue(ctx, serverTransportKey{}, tp)
    66  }
    67  
    68  // FromServerTransportContext 通过服务端context获取链路信息
    69  func FromServerTransportContext(ctx context.Context) (Transporter, bool) {
    70  	v, ok := ctx.Value(serverTransportKey{}).(Transporter)
    71  	return v, ok
    72  }
    73  
    74  // NewClientTransportContext 新建客户端链路上下文
    75  func NewClientTransportContext(ctx context.Context, tp Transporter) context.Context {
    76  	return context.WithValue(ctx, clientTransportKey{}, tp)
    77  }
    78  
    79  // FromClientTransportContext 通过客户端context获取链路信息
    80  func FromClientTransportContext(ctx context.Context) (Transporter, bool) {
    81  	v, ok := ctx.Value(clientTransportKey{}).(Transporter)
    82  	return v, ok
    83  }