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

     1  package conn_limiter
     2  
     3  import (
     4  	"github.com/nyan233/littlerpc/core/middle/plugin"
     5  	"sync/atomic"
     6  )
     7  
     8  type Limiter struct {
     9  	plugin.AbstractServer
    10  	max     int
    11  	counter atomic.Int64
    12  }
    13  
    14  func NewServer(concurrentSize int, clientSize int) plugin.ServerPlugin {
    15  	return &Limiter{
    16  		max: concurrentSize * clientSize,
    17  	}
    18  }
    19  
    20  func (l *Limiter) Event4S(ev plugin.Event) (next bool) {
    21  	switch ev {
    22  	case plugin.OnOpen:
    23  		_ = l.counter.Add(1)
    24  		if l.counter.Load() > int64(l.max) {
    25  			l.counter.Add(-1)
    26  			return false
    27  		}
    28  		return true
    29  	case plugin.OnClose:
    30  		l.counter.Add(-1)
    31  		return true
    32  	default:
    33  		return true
    34  	}
    35  }