github.com/anycable/anycable-go@v1.5.1/utils/gopool_test.go (about)

     1  package utils
     2  
     3  import (
     4  	"bytes"
     5  	"runtime"
     6  	"strconv"
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  func TestWorkerRespawn(t *testing.T) {
    14  	pool := NewGoPool("test", 1)
    15  
    16  	var wg sync.WaitGroup
    17  
    18  	resChan := make(chan uint64, 2)
    19  	dataChan := make(chan struct{}, workerRespawnThreshold+2)
    20  
    21  	wg.Add(1)
    22  
    23  	pool.Schedule(func() {
    24  		resChan <- getGID()
    25  	})
    26  
    27  	for i := 0; i < workerRespawnThreshold+2; i++ {
    28  		pool.Schedule(func() {
    29  			dataChan <- struct{}{}
    30  		})
    31  	}
    32  
    33  	pool.Schedule(func() {
    34  		resChan <- getGID()
    35  		wg.Done()
    36  	})
    37  
    38  	initial := <-resChan
    39  	current := <-resChan
    40  
    41  	assert.NotEqual(t, initial, current)
    42  	assert.Equal(t, workerRespawnThreshold+2, len(dataChan))
    43  }
    44  
    45  // Get current goroutine ID
    46  // Source: https://blog.sgmansfield.com/2015/12/goroutine-ids/
    47  func getGID() uint64 {
    48  	b := make([]byte, 64)
    49  	b = b[:runtime.Stack(b, false)]
    50  	b = bytes.TrimPrefix(b, []byte("goroutine "))
    51  	b = b[:bytes.IndexByte(b, ' ')]
    52  	n, _ := strconv.ParseUint(string(b), 10, 64)
    53  	return n
    54  }