diff --git a/common/signal/semaphore.go b/common/signal/semaphore.go new file mode 100644 index 000000000..034a4ee77 --- /dev/null +++ b/common/signal/semaphore.go @@ -0,0 +1,23 @@ +package signal + +type Semaphore struct { + token chan bool +} + +func NewSemaphore(n int) *Semaphore { + s := &Semaphore{ + token: make(chan bool, n), + } + for i := 0; i < n; i++ { + s.token <- true + } + return s +} + +func (s *Semaphore) Wait() <-chan bool { + return s.token +} + +func (s *Semaphore) Signal() { + s.token <- true +} diff --git a/transport/internet/internal/pool.go b/transport/internet/internal/pool.go index efa877246..4269400dd 100644 --- a/transport/internet/internal/pool.go +++ b/transport/internet/internal/pool.go @@ -4,6 +4,8 @@ import ( "net" "sync" "time" + + "v2ray.com/core/common/signal" ) // ConnectionRecyler is the interface for recycling connections. @@ -31,16 +33,15 @@ func (ec *ExpiringConnection) Expired() bool { type Pool struct { sync.Mutex connsByDest map[ConnectionID][]*ExpiringConnection - cleanupToken chan bool + cleanupToken *signal.Semaphore } // NewConnectionPool creates a new Pool. func NewConnectionPool() *Pool { p := &Pool{ connsByDest: make(map[ConnectionID][]*ExpiringConnection), - cleanupToken: make(chan bool, 1), + cleanupToken: signal.NewSemaphore(1), } - p.cleanupToken <- true return p } @@ -74,9 +75,7 @@ func (p *Pool) Get(id ConnectionID) net.Conn { } func (p *Pool) cleanup() { - defer func() { - p.cleanupToken <- true - }() + defer p.cleanupToken.Signal() for len(p.connsByDest) > 0 { time.Sleep(time.Second * 5) @@ -121,7 +120,7 @@ func (p *Pool) Put(id ConnectionID, conn net.Conn) { p.connsByDest[id] = list select { - case <-p.cleanupToken: + case <-p.cleanupToken.Wait(): go p.cleanup() default: }