adm-ods/common.in/utils/context.go

package utils

import (
	"context"

	"adm-ods/common.in/span"
	beegoctx "github.com/astaxie/beego/context"
	"github.com/smallnest/rpcx/share"
)

func NewContextFromBeegoCtx(ctx *beegoctx.Context) context.Context {
	// request metadata
	metadata := make(map[string]string)
	if spanID, err := span.ParseSpanID(ctx.Input.Header("span")); err == nil {
		metadata["span"] = span.NewSpanID(*spanID).String()
	}

	return context.WithValue(context.Background(), share.ReqMetaDataKey, metadata)
}

func NewContextFromRpcCtx(ctx context.Context) context.Context {
	metadata := make(map[string]string)
	metadata["span"] = span.NewRootSpanID().String() // default new a root spanid
	if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
		if mapMD := reqMetadata.(map[string]string); mapMD != nil {
			if v, ok := mapMD["span"]; ok {
				if spanID, err := span.ParseSpanID(v); err == nil {
					metadata["span"] = span.NewSpanID(*spanID).String()
				}
			}
		}
	}

	return context.WithValue(context.Background(), share.ReqMetaDataKey, metadata)
}

func GetSpanFromRpcCtx(ctx context.Context) (spanID span.SpanID) {
	if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
		if mapMD := reqMetadata.(map[string]string); mapMD != nil {
			if v, ok := mapMD["span"]; ok {
				if s, err := span.ParseSpanID(v); err == nil {
					spanID = *s
				}
			}
		}
	}

	return
}

func GetMethodFromRpcCtx(ctx context.Context) (method string) {
	if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
		if mapMD := reqMetadata.(map[string]string); mapMD != nil {
			if v, ok := mapMD["method"]; ok {
				method = v
			}
		}
	}

	return
}

func SetMethodToRpcCtx(ctx context.Context, method string) {

	if reqMetadata := ctx.Value(share.ReqMetaDataKey); reqMetadata != nil {
		if mapMD := reqMetadata.(map[string]string); mapMD != nil {
			mapMD["method"] = method
		}
	}
}