From 3532f2a414726ed96bb295795a9c62383f6a6adb Mon Sep 17 00:00:00 2001
From: Monkey <golang@88.com>
Date: Tue, 14 Feb 2023 18:01:53 +0800
Subject: [PATCH] fix: limit the number of connections created (#2441)

* fix: limit the number of connections created

Signed-off-by: monkey92t <golang@88.com>
---
 internal/pool/pool.go      | 30 +++++++++++++++++++-----------
 internal/pool/pool_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/internal/pool/pool.go b/internal/pool/pool.go
index bd17a028..bb9b14be 100644
--- a/internal/pool/pool.go
+++ b/internal/pool/pool.go
@@ -112,18 +112,25 @@ func (p *ConnPool) checkMinIdleConns() {
 		return
 	}
 	for p.poolSize < p.cfg.PoolSize && p.idleConnsLen < p.cfg.MinIdleConns {
-		p.poolSize++
-		p.idleConnsLen++
+		select {
+		case p.queue <- struct{}{}:
+			p.poolSize++
+			p.idleConnsLen++
 
-		go func() {
-			err := p.addIdleConn()
-			if err != nil && err != ErrClosed {
-				p.connsMu.Lock()
-				p.poolSize--
-				p.idleConnsLen--
-				p.connsMu.Unlock()
-			}
-		}()
+			go func() {
+				err := p.addIdleConn()
+				if err != nil && err != ErrClosed {
+					p.connsMu.Lock()
+					p.poolSize--
+					p.idleConnsLen--
+					p.connsMu.Unlock()
+				}
+
+				p.freeTurn()
+			}()
+		default:
+			return
+		}
 	}
 }
 
@@ -401,6 +408,7 @@ func (p *ConnPool) removeConn(cn *Conn) {
 			break
 		}
 	}
+	atomic.AddUint32(&p.stats.StaleConns, 1)
 }
 
 func (p *ConnPool) closeConn(cn *Conn) error {
diff --git a/internal/pool/pool_test.go b/internal/pool/pool_test.go
index 92fed14a..76dec996 100644
--- a/internal/pool/pool_test.go
+++ b/internal/pool/pool_test.go
@@ -327,4 +327,30 @@ var _ = Describe("race", func() {
 			}
 		})
 	})
+
+	It("limit the number of connections", func() {
+		opt := &pool.Options{
+			Dialer: func(ctx context.Context) (net.Conn, error) {
+				return &net.TCPConn{}, nil
+			},
+			PoolSize:     1000,
+			MinIdleConns: 50,
+			PoolTimeout:  3 * time.Second,
+		}
+		p := pool.NewConnPool(opt)
+
+		var wg sync.WaitGroup
+		for i := 0; i < opt.PoolSize; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				_, _ = p.Get(ctx)
+			}()
+		}
+		wg.Wait()
+
+		stats := p.Stats()
+		Expect(stats.IdleConns).To(Equal(uint32(0)))
+		Expect(stats.TotalConns).To(Equal(uint32(opt.PoolSize)))
+	})
 })