github.com/wasilibs/wazerox@v0.0.0-20240124024944-4923be63ab5f/internal/testing/hammer/hammer.go (about)

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