github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/other/trace/internal/instrument/instrument.go (about)

     1  package instrument
     2  
     3  import (
     4  	"fmt"
     5  	"reflect"
     6  	"runtime"
     7  	"sync"
     8  )
     9  
    10  type fnc struct {
    11  	flag *bool
    12  	typ  reflect.Type
    13  }
    14  
    15  var (
    16  	funcs     = make(map[string]fnc)
    17  	callbacks = make(map[string]interface{})
    18  	cbMtx     sync.RWMutex
    19  )
    20  
    21  // RegisterFlag is only meant to be called by code
    22  // generated using the instrument tool; it should not
    23  // be called by normal consumers of this package.
    24  func RegisterFlag(f interface{}, flag *bool) {
    25  	name := interfaceToName(f, "RegisterFlag")
    26  	funcs[name] = fnc{flag: flag, typ: reflect.TypeOf(f)}
    27  }
    28  
    29  func Instrument(f interface{}, callback interface{}) {
    30  	InstrumentName(interfaceToName(f, "Instrument"), callback)
    31  }
    32  
    33  func InstrumentName(fname string, callback interface{}) {
    34  	f, ok := funcs[fname]
    35  	if !ok {
    36  		panic(fmt.Errorf("instrument: InstrumentName: no such registered function: %v", fname))
    37  	}
    38  	if voidFunc(f.typ) != reflect.TypeOf(callback) {
    39  		panic(fmt.Errorf("instrument: InstrumentName: callback type (%v) does not match instrumented function type (%v)", reflect.TypeOf(callback), f.typ))
    40  	}
    41  	cbMtx.Lock()
    42  	callbacks[fname] = callback
    43  	*f.flag = true
    44  	cbMtx.Unlock()
    45  }
    46  
    47  func Uninstrument(f interface{}) {
    48  	UninstrumentName(interfaceToName(f, "Uninstrument"))
    49  }
    50  
    51  func UninstrumentName(fname string) {
    52  	f, ok := funcs[fname]
    53  	if !ok {
    54  		panic(fmt.Errorf("instrument: InstrumentName: no such registered function: %v", fname))
    55  	}
    56  	cbMtx.Lock()
    57  	delete(callbacks, fname)
    58  	*f.flag = false
    59  	cbMtx.Unlock()
    60  }
    61  
    62  func GetCallback(f interface{}) (callback interface{}, ok bool) {
    63  	return GetCallbackName(interfaceToName(f, "GetCallback"))
    64  }
    65  
    66  func GetCallbackName(fname string) (callback interface{}, ok bool) {
    67  	cbMtx.RLock()
    68  	c, ok := callbacks[fname]
    69  	cbMtx.RUnlock()
    70  	return c, ok
    71  }
    72  
    73  func GetType(f interface{}) (typ reflect.Type, ok bool) {
    74  	return GetTypeName(interfaceToName(f, "GetType"))
    75  }
    76  
    77  func GetTypeName(fname string) (typ reflect.Type, ok bool) {
    78  	f, ok := funcs[fname]
    79  	if !ok {
    80  		return nil, false
    81  	}
    82  	return voidFunc(f.typ), false
    83  }
    84  
    85  func voidFunc(typ reflect.Type) reflect.Type {
    86  	in := make([]reflect.Type, typ.NumIn())
    87  	for i := range in {
    88  		in[i] = typ.In(i)
    89  	}
    90  	return reflect.FuncOf(in, nil, typ.IsVariadic())
    91  }
    92  
    93  // f is the function whose name should be retreived,
    94  // and fname is the name of the top-level function
    95  // that is calling interfaceToName (used in panic
    96  // messages)
    97  func interfaceToName(f interface{}, fname string) string {
    98  	v := reflect.ValueOf(f)
    99  	if v.Kind() != reflect.Func {
   100  		panic(fmt.Errorf("instrument: %v with non-func %v", fname, v.Type()))
   101  	}
   102  	return runtime.FuncForPC(v.Pointer()).Name()
   103  }