trpc.group/trpc-go/trpc-go@v1.0.3/rpcz/context.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package rpcz
    15  
    16  import "context"
    17  
    18  type (
    19  	spanKey    struct{}
    20  	notSpanKey struct{}
    21  )
    22  
    23  // SpanFromContext returns Span from ctx.
    24  // If no Span isn't currently set in ctx, returns GlobalRPCZ.
    25  func SpanFromContext(ctx context.Context) Span {
    26  	if s, ok := ctx.Value(spanKey{}).(Span); ok {
    27  		return s
    28  	}
    29  	return GlobalRPCZ
    30  }
    31  
    32  // ContextWithSpan returns a copy of parent with span set as the current Span.
    33  func ContextWithSpan(ctx context.Context, span Span) context.Context {
    34  	return context.WithValue(ctx, spanKey{}, span)
    35  }
    36  
    37  // NewSpanContext creates a new span and binds it to a new ctx by combining SpanFromContext, Span.NewChildSpan and
    38  // ContextWithSpan. It returns the original ctx if new span is the same as original one to avoid unnecessary valueCtx.
    39  // Ender used to end this span immediately，and doesn't wait for related work to stop.
    40  // Ender doesn't end context.Context.
    41  func NewSpanContext(ctx context.Context, name string) (Span, Ender, context.Context) {
    42  	span := SpanFromContext(ctx)
    43  	newSpan, end := span.NewChild(name)
    44  	if newSpan == span {
    45  		return newSpan, end, ctx
    46  	}
    47  	return newSpan, end, ContextWithSpan(ctx, newSpan)
    48  }