From d6e5034bba6b439353ccde76899649996ced08c5 Mon Sep 17 00:00:00 2001
From: Winni Neessen <wn@neessen.net>
Date: Mon, 23 Sep 2024 10:09:38 +0200
Subject: [PATCH] Add new error handling and connection management in connpool

Introduce ErrClosed and ErrNilConn errors for better error handling. Implement Close and MarkUnusable methods for improved connection lifecycle management. Add put method to return connections to the pool or close them if necessary.
---
 connpool.go | 50 ++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 48 insertions(+), 2 deletions(-)

diff --git a/connpool.go b/connpool.go
index 23ed5d0..b22019c 100644
--- a/connpool.go
+++ b/connpool.go
@@ -17,11 +17,13 @@ import (
 // concurrency template.
 
 var (
+	// ErrClosed is returned when an operation is attempted on a closed connection pool.
+	ErrClosed = errors.New("connection pool is closed")
+	// ErrNilConn is returned when a nil connection is passed back to the connection pool.
+	ErrNilConn = errors.New("connection is nil")
 	// ErrPoolInvalidCap is returned when the connection pool's capacity settings are
 	// invalid (e.g., initial capacity is negative).
 	ErrPoolInvalidCap = errors.New("invalid connection pool capacity settings")
-	// ErrClosed is returned when an operation is attempted on a closed connection pool.
-	ErrClosed = errors.New("connection pool is closed")
 )
 
 // Pool interface describes a connection pool implementation. A Pool is
@@ -66,6 +68,28 @@ type PoolConn struct {
 	unusable bool
 }
 
+// Close puts a given pool connection back to the pool instead of closing it.
+func (c *PoolConn) Close() error {
+	c.mutex.RLock()
+	defer c.mutex.RUnlock()
+
+	if c.unusable {
+		if c.Conn != nil {
+			return c.Conn.Close()
+		}
+		return nil
+	}
+	return c.pool.put(c.Conn)
+}
+
+// MarkUnusable marks the connection not usable any more, to let the pool close it instead
+// of returning it to pool.
+func (c *PoolConn) MarkUnusable() {
+	c.mutex.Lock()
+	c.unusable = true
+	c.mutex.Unlock()
+}
+
 // NewConnPool returns a new pool based on buffered channels with an initial
 // capacity and maximum capacity. The DialContextFunc is used when the initial
 // capacity is greater than zero to fill the pool. A zero initialCap doesn't
@@ -167,6 +191,28 @@ func (p *connPool) getConnsAndDialContext() (context.Context, chan net.Conn, Dia
 	return ctx, conns, dialCtxFunc
 }
 
+// put puts a passed connection back into the pool. If the pool is full or closed,
+// conn is simply closed. A nil conn will be rejected with an error.
+func (p *connPool) put(conn net.Conn) error {
+	if conn == nil {
+		return ErrNilConn
+	}
+
+	p.mutex.RLock()
+	defer p.mutex.RUnlock()
+
+	if p.conns == nil {
+		return conn.Close()
+	}
+
+	select {
+	case p.conns <- conn:
+		return nil
+	default:
+		return conn.Close()
+	}
+}
+
 // wrapConn wraps a given net.Conn with a PoolConn, modifying the net.Conn's Close() method.
 func (p *connPool) wrapConn(conn net.Conn) net.Conn {
 	poolconn := &PoolConn{pool: p}