feat(chpool): Add health-check function

Author: Neven Miculinic

chpool/client.go

@@ -28,7 +28,21 @@ func (c *Client) Release() {
 		return
 	}
 
-	c.res.Release()
+	if time.Since(c.res.Value().lastHealthCheckTimestamp) > c.p.options.HealthCheckPeriod && c.p.options.HealthCheckFunc != nil {
+		go func() {
+			now := time.Now()
+			ctx, cancel := context.WithTimeout(context.Background(), c.p.options.HealthCheckTimeout)
+			defer cancel()
+			if err := c.p.options.HealthCheckFunc(ctx, c.res.Value().client); err != nil {
+				c.res.Destroy()
+				return
+			}
+			c.res.Value().lastHealthCheckTimestamp = now
+			c.res.Release()
+		}()
+	} else {
+		c.res.Release()
+	}
 }
 
 func (c *Client) Do(ctx context.Context, q ch.Query) (err error) {

chpool/conn.go

@@ -1,14 +1,17 @@
 package chpool
 
 import (
+	"time"
+
 	"github.com/jackc/puddle/v2"
 
 	"github.com/ClickHouse/ch-go"
 )
 
 type connResource struct {
-	client  *ch.Client
-	clients []Client
+	lastHealthCheckTimestamp time.Time
+	client                   *ch.Client
+	clients                  []Client
 }
 
 func (cr *connResource) getConn(p *Pool, res *puddle.Resource[*connResource]) *Client {

chpool/pool.go

@@ -23,19 +23,26 @@ type Pool struct {
 
 // Options for Pool.
 type Options struct {
-	ClientOptions     ch.Options
-	MaxConnLifetime   time.Duration
-	MaxConnIdleTime   time.Duration
-	MaxConns          int32
-	MinConns          int32
-	HealthCheckPeriod time.Duration
+	ClientOptions      ch.Options
+	MaxConnLifetime    time.Duration
+	MaxConnIdleTime    time.Duration
+	MaxConns           int32
+	MinConns           int32
+	HealthCheckPeriod  time.Duration
+	HealthCheckFunc    func(ctx context.Context, client *ch.Client) error
+	HealthCheckTimeout time.Duration
+}
+
+func DefaultHealthCheckFunc(ctx context.Context, client *ch.Client) error {
+	return client.Ping(ctx)
 }
 
 // Defaults for pool.
 const (
-	DefaultMaxConnLifetime   = time.Hour
-	DefaultMaxConnIdleTime   = time.Minute * 30
-	DefaultHealthCheckPeriod = time.Minute
+	DefaultMaxConnLifetime    = time.Hour
+	DefaultMaxConnIdleTime    = time.Minute * 30
+	DefaultHealthCheckPeriod  = time.Minute
+	DefaultHealthCheckTimeout = time.Second
 )
 
 func (o *Options) setDefaults() {
@@ -51,6 +58,12 @@ func (o *Options) setDefaults() {
 	if o.HealthCheckPeriod == 0 {
 		o.HealthCheckPeriod = DefaultHealthCheckPeriod
 	}
+	if o.HealthCheckFunc == nil {
+		o.HealthCheckFunc = DefaultHealthCheckFunc
+	}
+	if o.HealthCheckTimeout == 0 {
+		o.HealthCheckTimeout = DefaultHealthCheckTimeout
+	}
 }
 
 // Dial returns a pool of connections to ClickHouse.
@@ -163,15 +176,34 @@ func (p *Pool) checkIdleConnsHealth() {
 	resources := p.pool.AcquireAllIdle()
 
 	now := time.Now()
+	wg := sync.WaitGroup{}
 	for _, res := range resources {
-		if now.Sub(res.CreationTime()) > p.options.MaxConnLifetime {
-			res.Destroy()
-		} else if res.IdleDuration() > p.options.MaxConnIdleTime {
-			res.Destroy()
-		} else {
+		res := res
+		wg.Add(1)
+		go func() {
+			if now.Sub(res.CreationTime()) > p.options.MaxConnLifetime {
+				res.Destroy()
+				return
+			}
+
+			if res.IdleDuration() > p.options.MaxConnIdleTime {
+				res.Destroy()
+				return
+			}
+
+			if p.options.HealthCheckFunc != nil {
+				ctx, cancel := context.WithTimeout(context.Background(), p.options.HealthCheckTimeout)
+				defer cancel()
+				if err := p.options.HealthCheckFunc(ctx, res.Value().client); err != nil {
+					res.Destroy()
+					return
+				}
+				res.Value().lastHealthCheckTimestamp = now
+			}
 			res.ReleaseUnused()
-		}
+		}()
 	}
+	wg.Wait()
 }
 
 func (p *Pool) checkMinConns() {