context.go raw

   1  // Copyright The OpenTelemetry Authors
   2  // SPDX-License-Identifier: Apache-2.0
   3  
   4  package baggage // import "go.opentelemetry.io/otel/internal/baggage"
   5  
   6  import "context"
   7  
   8  type baggageContextKeyType int
   9  
  10  const baggageKey baggageContextKeyType = iota
  11  
  12  // SetHookFunc is a callback called when storing baggage in the context.
  13  type SetHookFunc func(context.Context, List) context.Context
  14  
  15  // GetHookFunc is a callback called when getting baggage from the context.
  16  type GetHookFunc func(context.Context, List) List
  17  
  18  type baggageState struct {
  19  	list List
  20  
  21  	setHook SetHookFunc
  22  	getHook GetHookFunc
  23  }
  24  
  25  // ContextWithSetHook returns a copy of parent with hook configured to be
  26  // invoked every time ContextWithBaggage is called.
  27  //
  28  // Passing nil SetHookFunc creates a context with no set hook to call.
  29  func ContextWithSetHook(parent context.Context, hook SetHookFunc) context.Context {
  30  	var s baggageState
  31  	if v, ok := parent.Value(baggageKey).(baggageState); ok {
  32  		s = v
  33  	}
  34  
  35  	s.setHook = hook
  36  	return context.WithValue(parent, baggageKey, s)
  37  }
  38  
  39  // ContextWithGetHook returns a copy of parent with hook configured to be
  40  // invoked every time FromContext is called.
  41  //
  42  // Passing nil GetHookFunc creates a context with no get hook to call.
  43  func ContextWithGetHook(parent context.Context, hook GetHookFunc) context.Context {
  44  	var s baggageState
  45  	if v, ok := parent.Value(baggageKey).(baggageState); ok {
  46  		s = v
  47  	}
  48  
  49  	s.getHook = hook
  50  	return context.WithValue(parent, baggageKey, s)
  51  }
  52  
  53  // ContextWithList returns a copy of parent with baggage. Passing nil list
  54  // returns a context without any baggage.
  55  func ContextWithList(parent context.Context, list List) context.Context {
  56  	var s baggageState
  57  	if v, ok := parent.Value(baggageKey).(baggageState); ok {
  58  		s = v
  59  	}
  60  
  61  	s.list = list
  62  	ctx := context.WithValue(parent, baggageKey, s)
  63  	if s.setHook != nil {
  64  		ctx = s.setHook(ctx, list)
  65  	}
  66  
  67  	return ctx
  68  }
  69  
  70  // ListFromContext returns the baggage contained in ctx.
  71  func ListFromContext(ctx context.Context) List {
  72  	switch v := ctx.Value(baggageKey).(type) {
  73  	case baggageState:
  74  		if v.getHook != nil {
  75  			return v.getHook(ctx, v.list)
  76  		}
  77  		return v.list
  78  	default:
  79  		return nil
  80  	}
  81  }
  82