github.com/grailbio/base@v0.0.11/backgroundcontext/backgroundcontext.go (about)

     1  // Copyright 2019 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package backgroundcontext manages the singleton v23 context.  This
     6  // package is not for general use. It it only for packages that (1)
     7  // need to access a background context, and (2) are used both as a
     8  // binary and as a shared library (e.g., in R).
     9  package backgroundcontext
    10  
    11  import (
    12  	"context"
    13  	"sync/atomic"
    14  	"unsafe"
    15  
    16  	v23context "v.io/v23/context"
    17  )
    18  
    19  var ptr unsafe.Pointer
    20  
    21  // Set sets the singleton global context. It should be called at most once,
    22  // usually immediately after a process starts. Calling Set() multiple times will
    23  // cause a panic. Thread safe.
    24  func Set(ctx *v23context.T) {
    25  	if !atomic.CompareAndSwapPointer(&ptr, nil, unsafe.Pointer(ctx)) {
    26  		panic("backgroundcontext.Set called twice")
    27  	}
    28  }
    29  
    30  // T returns the background context set by Set. It panics if Set has not been
    31  // called yet. Thread safe.
    32  func T() *v23context.T {
    33  	p := atomic.LoadPointer(&ptr)
    34  	if p == nil {
    35  		panic("backgroundcontext.Set not yet called")
    36  	}
    37  	return (*v23context.T)(p)
    38  }
    39  
    40  // Get returns a background context: if a v23 context has been set, it is returned;
    41  // otherwise the standard Go background context is returned.
    42  func Get() context.Context {
    43  	if p := (*v23context.T)(atomic.LoadPointer(&ptr)); p != nil {
    44  		return p
    45  	}
    46  	return context.Background()
    47  }
    48  
    49  type wrapped struct {
    50  	context.Context
    51  	v23ctx *v23context.T
    52  }
    53  
    54  func (w *wrapped) Value(key interface{}) interface{} {
    55  	val := w.Context.Value(key)
    56  	if val != nil {
    57  		return val
    58  	}
    59  	return w.v23ctx.Value(key)
    60  }
    61  
    62  // Wrap wraps the provided context, composing it with the defined
    63  // background context, if any. This allows contexts to be wrapped
    64  // so that v23 stubs can get the background context's RPC client.
    65  //
    66  // Cancelations are not forwarded: it is assumed that the set background
    67  // context is never canceled.
    68  //
    69  // BUG: this is very complicated; but it seems required to make
    70  // vanadium play nicely with contexts defined outside of its
    71  // universe.
    72  func Wrap(ctx context.Context) context.Context {
    73  	v23ctx := (*v23context.T)(atomic.LoadPointer(&ptr))
    74  	if v23ctx == nil {
    75  		return ctx
    76  	}
    77  	return &wrapped{ctx, v23ctx}
    78  }