From e23ea028bd95ae0f8a76a59ea53b7557e0c19fe5 Mon Sep 17 00:00:00 2001
From: Nikolay Vorobev <na.vorobyev@gmail.com>
Date: Wed, 20 Sep 2023 14:55:23 +0300
Subject: [PATCH] Added MaxActiveConns (#2646)

* Added the ability to set a connection growth limit when there are not enough connections in the pool using MaxActiveConns

* fix comment

* fix

* fix

---------

Co-authored-by: ofekshenawa <104765379+ofekshenawa@users.noreply.github.com>
---
 internal/pool/pool.go | 19 +++++++++++++------
 options.go            |  8 +++++++-
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/internal/pool/pool.go b/internal/pool/pool.go
index bb9b14beb..f391cedae 100644
--- a/internal/pool/pool.go
+++ b/internal/pool/pool.go
@@ -15,6 +15,10 @@ var (
 	// ErrClosed performs any operation on the closed client will return this error.
 	ErrClosed = errors.New("redis: client is closed")
 
+	// ErrPoolExhausted is returned from a pool connection method
+	// when the maximum number of database connections in the pool has been reached.
+	ErrPoolExhausted = errors.New("redis: connection pool exhausted")
+
 	// ErrPoolTimeout timed out waiting to get a connection from the connection pool.
 	ErrPoolTimeout = errors.New("redis: connection pool timeout")
 )
@@ -61,6 +65,7 @@ type Options struct {
 	PoolTimeout     time.Duration
 	MinIdleConns    int
 	MaxIdleConns    int
+	MaxActiveConns  int
 	ConnMaxIdleTime time.Duration
 	ConnMaxLifetime time.Duration
 }
@@ -159,6 +164,14 @@ func (p *ConnPool) NewConn(ctx context.Context) (*Conn, error) {
 }
 
 func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) {
+	if p.closed() {
+		return nil, ErrClosed
+	}
+
+	if p.cfg.MaxActiveConns > 0 && p.poolSize >= p.cfg.MaxActiveConns {
+		return nil, ErrPoolExhausted
+	}
+
 	cn, err := p.dialConn(ctx, pooled)
 	if err != nil {
 		return nil, err
@@ -167,12 +180,6 @@ func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) {
 	p.connsMu.Lock()
 	defer p.connsMu.Unlock()
 
-	// It is not allowed to add new connections to the closed connection pool.
-	if p.closed() {
-		_ = cn.Close()
-		return nil, ErrClosed
-	}
-
 	p.conns = append(p.conns, cn)
 	if pooled {
 		// If pool is full remove the cn on next Put.
diff --git a/options.go b/options.go
index f10bad38b..ba65defd2 100644
--- a/options.go
+++ b/options.go
@@ -98,8 +98,10 @@ type Options struct {
 	// Note that FIFO has slightly higher overhead compared to LIFO,
 	// but it helps closing idle connections faster reducing the pool size.
 	PoolFIFO bool
-	// Maximum number of socket connections.
+	// Base number of socket connections.
 	// Default is 10 connections per every available CPU as reported by runtime.GOMAXPROCS.
+	// If there is not enough connections in the pool, new connections will be allocated in excess of PoolSize,
+	// you can limit it through MaxActiveConns
 	PoolSize int
 	// Amount of time client waits for connection if all connections
 	// are busy before returning an error.
@@ -112,6 +114,9 @@ type Options struct {
 	// Maximum number of idle connections.
 	// Default is 0. the idle connections are not closed by default.
 	MaxIdleConns int
+	// Maximum number of connections allocated by the pool at a given time.
+	// When zero, there is no limit on the number of connections in the pool.
+	MaxActiveConns int
 	// ConnMaxIdleTime is the maximum amount of time a connection may be idle.
 	// Should be less than server's timeout.
 	//
@@ -502,6 +507,7 @@ func newConnPool(
 		PoolTimeout:     opt.PoolTimeout,
 		MinIdleConns:    opt.MinIdleConns,
 		MaxIdleConns:    opt.MaxIdleConns,
+		MaxActiveConns:  opt.MaxActiveConns,
 		ConnMaxIdleTime: opt.ConnMaxIdleTime,
 		ConnMaxLifetime: opt.ConnMaxLifetime,
 	})