github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/local/local.go (about)

     1  package local
     2  
     3  import (
     4  	"fmt"
     5  	"runtime"
     6  )
     7  
     8  type local []interface{}
     9  
    10  var callbacks []Callbacks
    11  var defaults []interface{}
    12  
    13  // Token represents a handle on a particular
    14  // local variable; it is used to distinguish
    15  // a local variable from others so that different
    16  // packages can use goroutine-local storage
    17  // simultaneously without interfering with
    18  // each other.
    19  type Token token
    20  
    21  // make sure that Token is an opaque type
    22  // so that the caller cannot manipulate it,
    23  // and we can change it later and not break
    24  // external code
    25  type token int
    26  
    27  // Callbacks is a collection of functions that
    28  // should be used for modifying local values.
    29  type Callbacks struct {
    30  	// LocalForSpawn takes the current goroutine's
    31  	// local and produces a local that should be
    32  	// set in the spawned goroutine. If LocalForSpawn
    33  	// is nil, the current goroutine's local is used
    34  	// as the initial value in the spawned goroutine.
    35  	LocalForSpawn func(local interface{}) interface{}
    36  }
    37  
    38  // Register registers a new local variable whose
    39  // initial value and callbacks are set from the
    40  // arguments. The returned Token can be used to
    41  // identify the variable in future calls.
    42  //
    43  // Register should ONLY be called during initialization,
    44  // and in the main goroutine.
    45  func Register(initial interface{}, c Callbacks) Token {
    46  	defaults = append(defaults, initial)
    47  	callbacks = append(callbacks, c)
    48  	return Token(len(callbacks) - 1)
    49  }
    50  
    51  // GetSpawnCallback returns a function which should
    52  // be called in the spawned goroutine in order to
    53  // set the local variables appropriately.
    54  //
    55  // NOTE: This should only be called by code generated
    56  // with the associated rewrite tool.
    57  func GetSpawnCallback() func() {
    58  	l := getLocal()
    59  	newl := make(local, len(l))
    60  	for i, c := range callbacks {
    61  		f := c.LocalForSpawn
    62  		if f == nil {
    63  			newl[i] = l[i]
    64  		} else {
    65  			newl[i] = f(l[i])
    66  		}
    67  	}
    68  
    69  	return func() {
    70  		runtime.SetLocal(newl)
    71  	}
    72  }
    73  
    74  // GetLocal returns the local variable associated
    75  // with the given Token.
    76  func GetLocal(t Token) interface{} {
    77  	l := getLocal()
    78  	return l[t]
    79  }
    80  
    81  // SetLocal sets the local variable associated
    82  // with the given Token.
    83  func SetLocal(t Token, l interface{}) {
    84  	ll := getLocal()
    85  	ll[t] = l
    86  }
    87  
    88  // getLocal retreives this goroutine's
    89  // locals slice, initializing it to
    90  // a copy of defaults if none exists
    91  func getLocal() local {
    92  	l, ok := runtime.GetLocal().(local)
    93  	if !ok {
    94  		l = make(local, len(defaults))
    95  		copy(l, defaults)
    96  		runtime.SetLocal(l)
    97  	}
    98  	if len(l) != len(defaults) {
    99  		panic(fmt.Errorf("unexpected number of locals: got %v; want %v", len(l), len(defaults)))
   100  	}
   101  	return l
   102  }