github.com/wangyougui/gf/v2@v2.6.5/os/gctx/gctx.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/wangyougui/gf.
     6  
     7  // Package gctx wraps context.Context and provides extra context features.
     8  package gctx
     9  
    10  import (
    11  	"context"
    12  	"os"
    13  	"strings"
    14  
    15  	"go.opentelemetry.io/otel"
    16  	"go.opentelemetry.io/otel/propagation"
    17  
    18  	"github.com/wangyougui/gf/v2/net/gtrace"
    19  )
    20  
    21  type (
    22  	Ctx    = context.Context // Ctx is short name alias for context.Context.
    23  	StrKey string            // StrKey is a type for warps basic type string as context key.
    24  )
    25  
    26  var (
    27  	// initCtx is the context initialized from process environment.
    28  	initCtx context.Context
    29  )
    30  
    31  func init() {
    32  	// All environment key-value pairs.
    33  	m := make(map[string]string)
    34  	i := 0
    35  	for _, s := range os.Environ() {
    36  		i = strings.IndexByte(s, '=')
    37  		if i == -1 {
    38  			continue
    39  		}
    40  		m[s[0:i]] = s[i+1:]
    41  	}
    42  	// OpenTelemetry from environments.
    43  	initCtx = otel.GetTextMapPropagator().Extract(
    44  		context.Background(),
    45  		propagation.MapCarrier(m),
    46  	)
    47  	initCtx = WithCtx(initCtx)
    48  }
    49  
    50  // New creates and returns a context which contains context id.
    51  func New() context.Context {
    52  	return WithCtx(context.Background())
    53  }
    54  
    55  // WithCtx creates and returns a context containing context id upon given parent context `ctx`.
    56  func WithCtx(ctx context.Context) context.Context {
    57  	if CtxId(ctx) != "" {
    58  		return ctx
    59  	}
    60  	var span *gtrace.Span
    61  	ctx, span = gtrace.NewSpan(ctx, "gctx.WithCtx")
    62  	defer span.End()
    63  	return ctx
    64  }
    65  
    66  // CtxId retrieves and returns the context id from context.
    67  func CtxId(ctx context.Context) string {
    68  	return gtrace.GetTraceID(ctx)
    69  }
    70  
    71  // SetInitCtx sets custom initialization context.
    72  // Note that this function cannot be called in multiple goroutines.
    73  func SetInitCtx(ctx context.Context) {
    74  	initCtx = ctx
    75  }
    76  
    77  // GetInitCtx returns the initialization context.
    78  // Initialization context is used in `main` or `init` functions.
    79  func GetInitCtx() context.Context {
    80  	return initCtx
    81  }