go.undefinedlabs.com/scopeagent@v0.4.2/init.go (about)

     1  package scopeagent // import "go.undefinedlabs.com/scopeagent"
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"os/signal"
     8  	"reflect"
     9  	"runtime"
    10  	"sync"
    11  	"syscall"
    12  	"testing"
    13  
    14  	"go.undefinedlabs.com/scopeagent/agent"
    15  	"go.undefinedlabs.com/scopeagent/env"
    16  	"go.undefinedlabs.com/scopeagent/instrumentation"
    17  	"go.undefinedlabs.com/scopeagent/instrumentation/logging"
    18  	scopetesting "go.undefinedlabs.com/scopeagent/instrumentation/testing"
    19  )
    20  
    21  var (
    22  	runningMutex sync.RWMutex
    23  	running      bool
    24  )
    25  
    26  // Helper function to run a `testing.M` object and gracefully stopping the agent afterwards
    27  func Run(m *testing.M, opts ...agent.Option) int {
    28  	if getRunningFlag() {
    29  		return m.Run()
    30  	}
    31  	setRunningFlag(true)
    32  	defer setRunningFlag(false)
    33  	opts = append(opts, agent.WithTestingModeEnabled())
    34  	newAgent, err := agent.NewAgent(opts...)
    35  	if err != nil {
    36  		res := m.Run()
    37  		if env.ScopeDebug.Value {
    38  			fmt.Printf("\n%v\n", err)
    39  		}
    40  		return res
    41  	}
    42  
    43  	logging.PatchStandardLogger()
    44  
    45  	scopetesting.Init(m)
    46  
    47  	// Handle SIGINT and SIGTERM
    48  	sigs := make(chan os.Signal, 1)
    49  	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
    50  	go func() {
    51  		<-sigs
    52  		if newAgent != nil {
    53  			instrumentation.Logger().Println("Terminating agent, sending partial results...")
    54  			newAgent.Stop()
    55  		}
    56  		os.Exit(1)
    57  	}()
    58  
    59  	return newAgent.Run(m)
    60  }
    61  
    62  // Instruments the given test, returning a `Test` object that can be used to extend the test trace
    63  func StartTest(t *testing.T, opts ...scopetesting.Option) *scopetesting.Test {
    64  	pc, _, _, _ := runtime.Caller(1)
    65  	return scopetesting.StartTestFromCaller(t, pc, opts...)
    66  }
    67  
    68  // Gets the *Test from a *testing.T
    69  func GetTest(t *testing.T) *scopetesting.Test {
    70  	return scopetesting.GetTest(t)
    71  }
    72  
    73  // Gets the context from a test
    74  func GetContextFromTest(t *testing.T) context.Context {
    75  	test := GetTest(t)
    76  	if test != nil {
    77  		return test.Context()
    78  	}
    79  	return context.TODO()
    80  }
    81  
    82  // Sets the test code from the caller of this func
    83  func SetTestCodeFromCaller(t *testing.T) {
    84  	SetTestCodeFromCallerSkip(t, 1)
    85  }
    86  
    87  // Sets the test code from the caller of this func
    88  func SetTestCodeFromCallerSkip(t *testing.T, skip int) {
    89  	test := GetTest(t)
    90  	if test == nil {
    91  		return
    92  	}
    93  	pc, _, _, _ := runtime.Caller(skip + 1)
    94  	test.SetTestCode(pc)
    95  }
    96  
    97  // Sets the test code from a func
    98  func SetTestCodeFromFunc(t *testing.T, fn interface{}) {
    99  	test := GetTest(t)
   100  	if test == nil {
   101  		return
   102  	}
   103  	value := reflect.ValueOf(fn)
   104  	pc := value.Pointer()
   105  	test.SetTestCode(pc)
   106  }
   107  
   108  // Gets the *Benchmark from a *testing.B
   109  func GetBenchmark(b *testing.B) *scopetesting.Benchmark {
   110  	return scopetesting.GetBenchmark(b)
   111  }
   112  
   113  // Instruments the given benchmark func
   114  //
   115  // Example of usage:
   116  //
   117  // func factorial(value int) int {
   118  //		if value == 1 {
   119  //			return 1
   120  //		}
   121  //		return value * factorial(value-1)
   122  // }
   123  //
   124  // func BenchmarkFactorial(b *testing.B) {
   125  // 		scopeagent.StartBenchmark(b, func(b *testing.B) {
   126  // 			for i := 0; i < b.N; i++ {
   127  //				_ = factorial(25)
   128  //			}
   129  //		}
   130  // }
   131  //
   132  // func BenchmarkFactorialWithSubBenchmarks(b *testing.B) {
   133  //		res := 0
   134  //
   135  //		b.Run("25", func(b *testing.B) {
   136  //			scopeagent.StartBenchmark(b, func(b *testing.B) {
   137  //				for i := 0; i < b.N; i++ {
   138  //					res = factorial(25)
   139  //				}
   140  //			})
   141  //		})
   142  //
   143  //		b.Run("50", func(b *testing.B) {
   144  //			scopeagent.StartBenchmark(b, func(b *testing.B) {
   145  //				for i := 0; i < b.N; i++ {
   146  //					res = factorial(50)
   147  //				}
   148  //			})
   149  //		})
   150  //
   151  //		b.Run("75", func(b *testing.B) {
   152  //			scopeagent.StartBenchmark(b, func(b *testing.B) {
   153  //				for i := 0; i < b.N; i++ {
   154  //					res = factorial(75)
   155  //				}
   156  //			})
   157  //		})
   158  //
   159  //		b.Run("100", func(b *testing.B) {
   160  //			scopeagent.StartBenchmark(b, func(b *testing.B) {
   161  //				for i := 0; i < b.N; i++ {
   162  //					res = factorial(100)
   163  //				}
   164  //			})
   165  //		})
   166  //
   167  //		_ = res
   168  // }
   169  func StartBenchmark(b *testing.B, benchFunc func(b *testing.B)) {
   170  	pc, _, _, _ := runtime.Caller(1)
   171  	scopetesting.StartBenchmark(b, pc, benchFunc)
   172  }
   173  
   174  func setRunningFlag(value bool) {
   175  	runningMutex.Lock()
   176  	defer runningMutex.Unlock()
   177  	running = value
   178  }
   179  func getRunningFlag() bool {
   180  	runningMutex.RLock()
   181  	defer runningMutex.RUnlock()
   182  	return running
   183  }