github.com/tetratelabs/wazero@v1.7.1/internal/testing/hammer/hammer.go (about)

     1  package hammer
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"testing"
     7  )
     8  
     9  // Hammer invokes a test concurrently in P goroutines N times per goroutine.
    10  //
    11  // Here's an example:
    12  //
    13  //	P := 8               // max count of goroutines
    14  //	N := 1000            // work per goroutine
    15  //	if testing.Short() { // Adjust down if `-test.short`
    16  //		P = 4
    17  //		N = 100
    18  //	}
    19  //
    20  //	hammer.NewHammer(t, P, N).Run(func(name string) {
    21  //		// Do test using name if something needs to be unique.
    22  //	}, nil)
    23  //
    24  //	if t.Failed() {
    25  //		return // At least one test failed, so return now.
    26  //	}
    27  //
    28  // See /RATIONALE.md
    29  type Hammer interface {
    30  	// Run invokes a concurrency test, as described in /RATIONALE.md.
    31  	//
    32  	// * test is concurrently run in P goroutines, each looping N times.
    33  	//   * name is unique within the hammer.
    34  	// * onRunning is any function to run after all goroutines are running, but before test executes.
    35  	//
    36  	// On completion, return early if there's a failure like this:
    37  	//	if t.Failed() {
    38  	//		return
    39  	//	}
    40  	Run(test func(p, n int), onRunning func())
    41  }
    42  
    43  // NewHammer returns a Hammer initialized to indicated count of goroutines (P) and iterations per goroutine (N).
    44  // As discussed in /RATIONALE.md, optimize for Hammer.Run completing in .1 second on a modern laptop.
    45  func NewHammer(t *testing.T, P, N int) Hammer {
    46  	return &hammer{t: t, P: P, N: N}
    47  }
    48  
    49  // hammer implements Hammer
    50  type hammer struct {
    51  	// t is the calling test
    52  	t *testing.T
    53  	// P is the max count of goroutines
    54  	P int
    55  	// N is the work per goroutine
    56  	N int
    57  }
    58  
    59  // Run implements Hammer.Run
    60  func (h *hammer) Run(test func(p, n int), onRunning func()) {
    61  	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(h.P / 2)) // Ensure goroutines have to switch cores.
    62  
    63  	// running track
    64  	running := make(chan int)
    65  	// unblock needs to happen atomically, so we need to use a WaitGroup
    66  	var unblocked sync.WaitGroup
    67  	finished := make(chan int)
    68  
    69  	unblocked.Add(h.P) // P goroutines will be unblocked by the current goroutine.
    70  	for p := 0; p < h.P; p++ {
    71  		p := p // pin p, so it is stable inside the goroutine.
    72  
    73  		go func() { // Launch goroutine 'p'
    74  			defer func() { // Ensure each require.XX failure is visible on hammer test fail.
    75  				if recovered := recover(); recovered != nil {
    76  					// Has been seen to be string, runtime.errorString, and it may be others. Let
    77  					// printing take care of conversion in a generic way.
    78  					h.t.Error(recovered)
    79  				}
    80  				finished <- 1
    81  			}()
    82  			running <- 1
    83  
    84  			unblocked.Wait()           // Wait to be unblocked
    85  			for n := 0; n < h.N; n++ { // Invoke one test
    86  				test(p, n)
    87  			}
    88  		}()
    89  	}
    90  
    91  	// Block until P goroutines are running.
    92  	for i := 0; i < h.P; i++ {
    93  		<-running
    94  	}
    95  
    96  	if onRunning != nil {
    97  		onRunning()
    98  	}
    99  
   100  	// Release all goroutines at the same time.
   101  	unblocked.Add(-h.P)
   102  
   103  	// Block until P goroutines finish.
   104  	for i := 0; i < h.P; i++ {
   105  		<-finished
   106  	}
   107  }