From 520e3ea9e67a8b86fac8660fe16395ea25450169 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Mon, 13 Feb 2017 23:29:34 +0100
Subject: [PATCH] signal.semaphore

---
 common/signal/semaphore.go          | 23 +++++++++++++++++++++++
 transport/internet/internal/pool.go | 13 ++++++-------
 2 files changed, 29 insertions(+), 7 deletions(-)
 create mode 100644 common/signal/semaphore.go

diff --git a/common/signal/semaphore.go b/common/signal/semaphore.go
new file mode 100644
index 00000000..034a4ee7
--- /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 efa87724..4269400d 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:
 	}