github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/plugins/conn_limiter/limiter_test.go (about)

     1  package conn_limiter
     2  
     3  import (
     4  	"github.com/nyan233/littlerpc/core/middle/plugin"
     5  	"github.com/stretchr/testify/assert"
     6  	"sync"
     7  	"sync/atomic"
     8  	"testing"
     9  )
    10  
    11  func TestConnLimiter(t *testing.T) {
    12  	const (
    13  		GoroutineSize = 2 * 512
    14  	)
    15  	p := NewServer(2, 512).(*Limiter)
    16  	done := make(chan int, 1)
    17  	var wg sync.WaitGroup
    18  	wg.Add(GoroutineSize + 100)
    19  	var falseCount atomic.Int64
    20  	for i := 0; i < GoroutineSize+100; i++ {
    21  		go func() {
    22  			select {
    23  			case <-done:
    24  				if !p.Event4S(plugin.OnOpen) {
    25  					falseCount.Add(1)
    26  				}
    27  				wg.Done()
    28  			}
    29  		}()
    30  	}
    31  	close(done)
    32  	wg.Wait()
    33  	assert.Equal(t, falseCount.Load(), int64(100))
    34  	wg.Add(GoroutineSize)
    35  	for i := 0; i < GoroutineSize; i++ {
    36  		go func() {
    37  			defer wg.Done()
    38  			p.Event4S(plugin.OnClose)
    39  			p.Event4S(plugin.OnMessage)
    40  		}()
    41  	}
    42  	wg.Wait()
    43  	assert.Equal(t, p.counter.Load(), int64(0))
    44  }