From b802a3e5e3f7fb5b25f2839ca17e3118dcc9b84c Mon Sep 17 00:00:00 2001
From: Kent Wang <kentwang@tencent.com>
Date: Wed, 3 Jul 2024 21:21:56 +0800
Subject: [PATCH] Add ConnChecker to customize conn health check

---
 internal/pool/pool.go | 7 +++++++
 options.go            | 4 ++++
 2 files changed, 11 insertions(+)

diff --git a/internal/pool/pool.go b/internal/pool/pool.go
index 2125f3e1..d60edc4b 100644
--- a/internal/pool/pool.go
+++ b/internal/pool/pool.go
@@ -68,6 +68,7 @@ type Options struct {
 	MaxActiveConns  int
 	ConnMaxIdleTime time.Duration
 	ConnMaxLifetime time.Duration
+	ConnChecker     func(net.Conn) error
 }
 
 type lastDialErrorWrap struct {
@@ -513,6 +514,12 @@ func (p *ConnPool) isHealthyConn(cn *Conn) bool {
 		return false
 	}
 
+	if p.cfg.ConnChecker != nil {
+		if err := p.cfg.ConnChecker(cn.netConn); err != nil {
+			return false
+		}
+	}
+
 	cn.SetUsedAt(now)
 	return true
 }
diff --git a/options.go b/options.go
index 6ed693a0..01205bec 100644
--- a/options.go
+++ b/options.go
@@ -139,6 +139,9 @@ type Options struct {
 	// Default is to not close idle connections.
 	ConnMaxLifetime time.Duration
 
+	// ConnChecker checks the health of a connection before returning it to the client.
+	ConnChecker func(net.Conn) error
+
 	// TLS Config to use. When set, TLS will be negotiated.
 	TLSConfig *tls.Config
 
@@ -520,5 +523,6 @@ func newConnPool(
 		MaxActiveConns:  opt.MaxActiveConns,
 		ConnMaxIdleTime: opt.ConnMaxIdleTime,
 		ConnMaxLifetime: opt.ConnMaxLifetime,
+		ConnChecker:     opt.ConnChecker,
 	})
 }