Rework connection liveness checking and connection reuse

README.md

@@ -196,6 +196,7 @@ The following connection settings are available in both DSN strings and the `cli
 * **max_open_conns** - Maximum number of open connections to the database (default: MaxIdleConns + 5)
 * **max_idle_conns** - Maximum number of idle connections in the pool (default: 5)
 * **conn_max_lifetime** - Maximum amount of time a connection may be reused (default: 1h)
+* **conn_idle_ping_threshold** - How long a pooled connection may sit idle before it is verified with a protocol ping when acquired. Detects half-open connections (e.g. silently dropped by a load balancer or firewall) that a socket-level check cannot see. Connections used more recently than this skip the ping entirely (default: 1m; a negative value disables idle pings)
 
 ### Connection Strategy
 * **connection_open_strategy** - Strategy for selecting servers from the connection pool:

clickhouse.go

@@ -90,6 +90,8 @@ type nativeTransport interface {
 	asyncInsert(ctx context.Context, query string, wait bool, args ...any) error
 	ping(context.Context) error
 	isBad() bool
+	markUnverified()
+	revalidateIdle(ctx context.Context) error
 	connID() int
 	connectedAtTime() time.Time
 	isReleased() bool
@@ -342,13 +344,16 @@ func (ch *clickhouse) acquire(ctx context.Context) (conn nativeTransport, err er
 	}
 
 	if err == nil && conn != nil {
-		if !conn.isBad() {
+		if conn.isBad() {
+			conn.close()
+		} else if rerr := conn.revalidateIdle(ctx); rerr != nil {
+			conn.getLogger().Debug("pooled connection failed liveness revalidation", slog.Any("error", rerr))
+			conn.close()
+		} else {
 			conn.setReleased(false)
 			conn.getLogger().Debug("connection acquired from pool")
 			return conn, nil
 		}
-
-		conn.close()
 	}
 
 	if conn, err = ch.dial(ctx); err != nil {
@@ -383,10 +388,18 @@ func (ch *clickhouse) release(conn nativeTransport, err error) {
 	}
 
 	if err != nil {
-		conn.getLogger().Debug("connection closed due to error", slog.Any("error", err))
-		conn.close()
-		return
-	} else if time.Since(conn.connectedAtTime()) >= ch.opt.ConnMaxLifetime {
+		var exc *Exception
+		if errors.As(err, &exc) {
+			conn.getLogger().Debug("connection kept after server exception", slog.Int("code", int(exc.Code)))
+			conn.markUnverified()
+		} else {
+			conn.getLogger().Debug("connection closed due to error", slog.Any("error", err))
+			conn.close()
+			return
+		}
+	}
+
+	if time.Since(conn.connectedAtTime()) >= ch.opt.ConnMaxLifetime {
 		conn.getLogger().Debug("connection closed: lifetime expired",
 			slog.Duration("age", time.Since(conn.connectedAtTime())),
 			slog.Duration("max_lifetime", ch.opt.ConnMaxLifetime))

clickhouse_options.go

@@ -149,19 +149,23 @@ type Options struct {
 	// instead of Logger.
 	Logger *slog.Logger
 
-	Settings             Settings
-	Compression          *Compression
-	DialTimeout          time.Duration // default 30 second
-	MaxOpenConns         int           // default MaxIdleConns + 5
-	MaxIdleConns         int           // default 5
-	ConnMaxLifetime      time.Duration // default 1 hour
-	ConnOpenStrategy     ConnOpenStrategy
-	FreeBufOnConnRelease bool              // drop preserved memory buffer after each query
-	HttpHeaders          map[string]string // set additional headers on HTTP requests
-	HttpUrlPath          string            // set additional URL path for HTTP requests
-	HttpMaxConnsPerHost  int               // MaxConnsPerHost for http.Transport
-	BlockBufferSize      uint8             // default 2 - can be overwritten on query
-	MaxCompressionBuffer int               // default 10485760 - measured in bytes  i.e.
+	Settings        Settings
+	Compression     *Compression
+	DialTimeout     time.Duration // default 30 second
+	MaxOpenConns    int           // default MaxIdleConns + 5
+	MaxIdleConns    int           // default 5
+	ConnMaxLifetime time.Duration // default 1 hour
+	// ConnIdlePingThreshold is how long a pooled connection may sit idle
+	// before it is verified with a protocol ping on acquire.
+	// Default is 1 minute, negative value disables.
+	ConnIdlePingThreshold time.Duration
+	ConnOpenStrategy      ConnOpenStrategy
+	FreeBufOnConnRelease  bool              // drop preserved memory buffer after each query
+	HttpHeaders           map[string]string // set additional headers on HTTP requests
+	HttpUrlPath           string            // set additional URL path for HTTP requests
+	HttpMaxConnsPerHost   int               // MaxConnsPerHost for http.Transport
+	BlockBufferSize       uint8             // default 2 - can be overwritten on query
+	MaxCompressionBuffer  int               // default 10485760 - measured in bytes  i.e.
 
 	// HTTPProxy specifies an HTTP proxy URL to use for requests made by the client.
 	HTTPProxyURL *url.URL
@@ -327,6 +331,12 @@ func (o *Options) fromDSN(in string) error {
 				return fmt.Errorf("conn_max_lifetime invalid value: %w", err)
 			}
 			o.ConnMaxLifetime = connMaxLifetime
+		case "conn_idle_ping_threshold":
+			connIdlePingThreshold, err := time.ParseDuration(params.Get(v))
+			if err != nil {
+				return fmt.Errorf("conn_idle_ping_threshold invalid value: %w", err)
+			}
+			o.ConnIdlePingThreshold = connIdlePingThreshold
 		case "username":
 			o.Auth.Username = params.Get(v)
 		case "password":
@@ -418,6 +428,9 @@ func (o Options) setDefaults() *Options {
 	if o.ConnMaxLifetime == 0 {
 		o.ConnMaxLifetime = time.Hour
 	}
+	if o.ConnIdlePingThreshold == 0 {
+		o.ConnIdlePingThreshold = time.Minute
+	}
 	if o.BlockBufferSize <= 0 {
 		o.BlockBufferSize = 2
 	}

clickhouse_options_test.go

@@ -502,6 +502,21 @@ func TestParseDSN(t *testing.T) {
 			},
 			"",
 		},
+		{
+			"client connection idle ping threshold",
+			"clickhouse://127.0.0.1/test_database?conn_idle_ping_threshold=30s",
+			&Options{
+				Protocol:              Native,
+				ConnIdlePingThreshold: 30 * time.Second,
+				Addr:                  []string{"127.0.0.1"},
+				Settings:              Settings{},
+				Auth: Auth{
+					Database: "test_database",
+				},
+				scheme: "clickhouse",
+			},
+			"",
+		},
 		{
 			"http protocol with proxy",
 			"http://127.0.0.1/?http_proxy=http%3A%2F%2Fproxy.example.com%3A3128",

clickhouse_std.go

@@ -181,6 +181,20 @@ func (std *stdDriver) Open(dsn string) (_ driver.Conn, err error) {
 
 var _ driver.Driver = (*stdDriver)(nil)
 
+func (std *stdDriver) release(_ nativeTransport, err error) {
+	if err == nil {
+		return
+	}
+
+	var exc *Exception
+	if errors.As(err, &exc) {
+		return
+	}
+
+	std.logger.Debug("closing connection after unclean stream termination", slog.Any("error", err))
+	_ = std.conn.close()
+}
+
 func (std *stdDriver) ResetSession(ctx context.Context) error {
 	if std.conn.isBad() {
 		std.logger.Debug("resetting session because connection is bad")
@@ -281,7 +295,7 @@ func (std *stdDriver) QueryContext(ctx context.Context, query string, args []dri
 		return nil, driver.ErrBadConn
 	}
 
-	r, err := std.conn.query(ctx, func(nativeTransport, error) {}, query, rebind(args)...)
+	r, err := std.conn.query(ctx, std.release, query, rebind(args)...)
 	if isConnBrokenError(err) {
 		std.logger.Error("query context got a fatal error, resetting connection", slog.Any("error", err))
 		return nil, driver.ErrBadConn
@@ -306,7 +320,7 @@ func (std *stdDriver) PrepareContext(ctx context.Context, query string) (driver.
 		return nil, driver.ErrBadConn
 	}
 
-	batch, err := std.conn.prepareBatch(ctx, func(nativeTransport, error) {}, func(context.Context) (nativeTransport, error) { return nil, nil }, query, chdriver.PrepareBatchOptions{})
+	batch, err := std.conn.prepareBatch(ctx, std.release, func(context.Context) (nativeTransport, error) { return nil, nil }, query, chdriver.PrepareBatchOptions{})
 	if err != nil {
 		if isConnBrokenError(err) {
 			std.logger.Error("prepare context got a fatal error, resetting connection", slog.Any("error", err))

conn.go

@@ -22,6 +22,15 @@ import (
 	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
 )
 
+// errUnexpectedRead is reported when an idle connection has
+// readable bytes. The native protocol should have no server-initiated packets between
+// queries, so this means the previous response was not fully drained.
+var errUnexpectedRead = errors.New("unexpected read from socket")
+
+// connLivenessWindow is how long a liveness proof is trusted
+// before the socket is probed again.
+const connLivenessWindow = 1 * time.Second
+
 func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, error) {
 	var (
 		err  error
@@ -78,6 +87,7 @@ func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, er
 			structMap:            &structMap{},
 			compression:          compression,
 			connectedAt:          time.Now(),
+			lastLiveAt:           time.Now(),
 			compressor:           compressor,
 			readTimeout:          opt.ReadTimeout,
 			blockBufferSize:      opt.BlockBufferSize,
@@ -131,6 +141,8 @@ type connect struct {
 	structMap            *structMap
 	compression          CompressionMethod
 	connectedAt          time.Time
+	lastLiveAt           time.Time
+	unverified           bool
 	compressor           *compress.Writer
 	readTimeout          time.Duration
 	blockBufferSize      uint8
@@ -192,19 +204,62 @@ func (c *connect) isBad() bool {
 		return true
 	}
 
+	if time.Since(c.lastLiveAt) < connLivenessWindow {
+		return false
+	}
+
 	if err := c.connCheck(); err != nil {
+		if errors.Is(err, errUnexpectedRead) {
+			c.logger.Warn("closing connection: unread server data on idle connection",
+				slog.Int("conn_id", c.id))
+		}
+
 		return true
 	}
 
+	c.lastLiveAt = time.Now()
+
 	return false
 }
 
+func (c *connect) markUnverified() {
+	c.unverified = true
+}
+
+func (c *connect) revalidateIdle(ctx context.Context) error {
+	if !c.unverified {
+		threshold := c.opt.ConnIdlePingThreshold
+		if threshold < 0 || time.Since(c.lastLiveAt) < threshold {
+			return nil
+		}
+	}
+
+	pingCtx := ctx
+	// Limit ping deadline to allow time for re-dial
+	if deadline, ok := ctx.Deadline(); ok {
+		var cancel context.CancelFunc
+		pingCtx, cancel = context.WithDeadline(ctx, time.Now().Add(time.Until(deadline)/2))
+		defer cancel()
+	}
+
+	if err := c.ping(pingCtx); err != nil {
+		return err
+	}
+
+	c.unverified = false
+	c.lastLiveAt = time.Now()
+	return nil
+}
+
 func (c *connect) isReleased() bool {
 	return c.released
 }
 
 func (c *connect) setReleased(released bool) {
 	c.released = released
+	if released {
+		c.lastLiveAt = time.Now()
+	}
 }
 
 func (c *connect) isClosed() bool {

conn_batch.go

@@ -22,6 +22,7 @@ var columnMatch = regexp.MustCompile(`INSERT INTO .+\s\((?P<Columns>.+)\)$`)
 func (c *connect) prepareBatch(ctx context.Context, release nativeTransportRelease, acquire nativeTransportAcquire, query string, opts driver.PrepareBatchOptions) (driver.Batch, error) {
 	query, _, queryColumns, verr := extractNormalizedInsertQueryAndColumns(query)
 	if verr != nil {
+		release(c, nil)
 		return nil, verr
 	}
 
@@ -205,6 +206,7 @@ func (b *batch) Send() (err error) {
 		// close TCP connection on context cancel. There is no other way simple way to interrupt underlying operations.
 		// as verified in the test, this is safe to do and cleanups resources later on
 		if b.conn != nil {
+			b.conn.setClosed()
 			_ = b.conn.conn.Close()
 		}
 	})

conn_check.go

@@ -5,7 +5,6 @@ package clickhouse
 
 import (
 	"crypto/tls"
-	"errors"
 	"io"
 	"syscall"
 )
@@ -28,12 +27,12 @@ func (c *connect) connCheck() error {
 
 	err = rawConn.Read(func(fd uintptr) bool {
 		var buf [1]byte
-		n, err := syscall.Read(int(fd), buf[:])
+		n, _, err := syscall.Recvfrom(int(fd), buf[:], syscall.MSG_PEEK)
 		switch {
 		case n == 0 && err == nil:
 			sysErr = io.EOF
 		case n > 0:
-			sysErr = errors.New("unexpected read from socket")
+			sysErr = errUnexpectedRead
 		case err == syscall.EAGAIN || err == syscall.EWOULDBLOCK:
 			sysErr = nil
 		default:

conn_http.go

@@ -317,6 +317,13 @@ func (h *httpConnect) isBad() bool {
 	return h.client == nil
 }
 
+func (h *httpConnect) markUnverified() {
+}
+
+func (h *httpConnect) revalidateIdle(ctx context.Context) error {
+	return nil
+}
+
 func (h *httpConnect) queryHello(ctx context.Context, release nativeTransportRelease) (proto.ServerHandshake, error) {
 	h.logger.Debug("querying server info via HTTP")
 	ctx = Context(ctx, ignoreExternalTables())

conn_pool_test.go

@@ -474,6 +474,7 @@ type mockTransport struct {
 	released      bool
 	closed        bool
 	bad           bool
+	unverified    bool
 	bufferFreed   bool
 	debugMessages []string
 	mu            sync.Mutex
@@ -513,6 +514,16 @@ func (m *mockTransport) isBad() bool {
 	return m.bad
 }
 
+func (m *mockTransport) markUnverified() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.unverified = true
+}
+
+func (m *mockTransport) revalidateIdle(ctx context.Context) error {
+	return nil
+}
+
 func (m *mockTransport) connID() int {
 	return m.id
 }

conn_query.go

@@ -17,7 +17,7 @@ func (c *connect) query(ctx context.Context, release nativeTransportRelease, que
 
 	if err != nil {
 		c.logger.Error("failed to bind query parameters", slog.Any("error", err))
-		release(c, err)
+		release(c, nil)
 		return nil, err
 	}
 
@@ -48,13 +48,13 @@ func (c *connect) query(ctx context.Context, release nativeTransportRelease, que
 			stream <- b
 		}
 		err := c.process(ctx, onProcess)
+		release(c, err)
 		if err != nil {
 			c.logger.Error("query processing failed", slog.Any("error", err))
 			errors <- err
 		}
 		close(stream)
 		close(errors)
-		release(c, err)
 	}()
 
 	return &rows{