Merge pull request #303 from ClickHouse/fix/rework-client-close

fix(client): rework client close

Author: ernado

.golangci.yml

@@ -34,15 +34,15 @@ linters-settings:
       - unnecessaryBlock
       - redundantSprint
       - octalLiteral
+      - ifElseChain
 
 linters:
   disable-all: true
   enable:
-    - depguard
     - dogsled
     - errcheck
     - goconst
-#    - gocritic # go1.18
+    - gocritic
     - gofmt
     - goimports
     - revive
@@ -67,6 +67,7 @@ linters:
   # - maligned  (same as prealloc)
   # - funlen    (gocyclo is enough)
   # - gochecknoglobals (we know when it is ok to use globals)
+  # - depguard  (broken)
 
 issues:
   exclude-use-default: false

client.go

@@ -29,14 +29,16 @@ import (
 type Client struct {
 	lg       *zap.Logger
 	conn     net.Conn
-	mux      sync.Mutex
 	buf      *proto.Buffer
 	reader   *proto.Reader
 	info     proto.ClientHello
 	server   proto.ServerHello
 	version  clientVersion
 	quotaKey string
 
+	mux    sync.Mutex
+	closed bool
+
 	// Single packet read timeout.
 	readTimeout time.Duration
 
@@ -80,14 +82,14 @@ var ErrClosed = errors.New("client is closed")
 // Close closes underlying connection and frees all resources,
 // rendering Client to unusable state.
 func (c *Client) Close() error {
-	if c.conn == nil {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	if c.closed {
 		return ErrClosed
 	}
-	defer func() {
-		c.buf = nil
-		c.reader = nil
-		c.conn = nil
-	}()
+
+	c.closed = true
 	if err := c.conn.Close(); err != nil {
 		return errors.Wrap(err, "conn")
 	}
@@ -97,7 +99,9 @@ func (c *Client) Close() error {
 
 // IsClosed indicates that connection is closed.
 func (c *Client) IsClosed() bool {
-	return c.conn == nil
+	c.mux.Lock()
+	defer c.mux.Unlock()
+	return c.closed
 }
 
 // Exception is server-side error.
@@ -251,8 +255,6 @@ func (c *Client) packet(ctx context.Context) (proto.ServerCode, error) {
 }
 
 func (c *Client) flushBuf(ctx context.Context, b *proto.Buffer) error {
-	c.mux.Lock()
-	defer c.mux.Unlock()
 	if err := ctx.Err(); err != nil {
 		return errors.Wrap(err, "context")
 	}

otel_test.go

@@ -13,7 +13,7 @@ type MapPair[V comparable] struct {
 	Values *proto.ColArr[V]
 }
 
-func (m MapPair[V]) Row(i int) ([]string, []V) {
+func (m MapPair[V]) Row(i int) (keys []string, values []V) {
 	return m.Keys.Row(i), m.Values.Row(i)
 }
 

ping.go

@@ -15,7 +15,7 @@ import (
 //
 // Do not call concurrently with Do.
 func (c *Client) Ping(ctx context.Context) (err error) {
-	if c.conn == nil {
+	if c.IsClosed() {
 		return ErrClosed
 	}
 	if c.otel {

proto/col_low_cardinality.go

@@ -13,34 +13,6 @@ var (
 	_ Column    = (*ColLowCardinality[string])(nil)
 )
 
-// DecodeState implements StateDecoder, ensuring state for index column.
-func (c *ColLowCardinality[T]) DecodeState(r *Reader) error {
-	keySerialization, err := r.Int64()
-	if err != nil {
-		return errors.Wrap(err, "version")
-	}
-	if keySerialization != int64(sharedDictionariesWithAdditionalKeys) {
-		return errors.Errorf("got version %d, expected %d",
-			keySerialization, sharedDictionariesWithAdditionalKeys,
-		)
-	}
-	if s, ok := c.index.(StateDecoder); ok {
-		if err := s.DecodeState(r); err != nil {
-			return errors.Wrap(err, "index state")
-		}
-	}
-	return nil
-}
-
-// EncodeState implements StateEncoder, ensuring state for index column.
-func (c ColLowCardinality[T]) EncodeState(b *Buffer) {
-	// Writing key serialization version.
-	b.PutInt64(int64(sharedDictionariesWithAdditionalKeys))
-	if s, ok := c.index.(StateEncoder); ok {
-		s.EncodeState(b)
-	}
-}
-
 //go:generate go run github.com/dmarkham/enumer -type CardinalityKey -trimprefix Key -output col_low_cardinality_enum.go
 
 // CardinalityKey is integer type of ColLowCardinality.Keys column.
@@ -119,6 +91,34 @@ type ColLowCardinality[T comparable] struct {
 	keys []int
 }
 
+// DecodeState implements StateDecoder, ensuring state for index column.
+func (c *ColLowCardinality[T]) DecodeState(r *Reader) error {
+	keySerialization, err := r.Int64()
+	if err != nil {
+		return errors.Wrap(err, "version")
+	}
+	if keySerialization != int64(sharedDictionariesWithAdditionalKeys) {
+		return errors.Errorf("got version %d, expected %d",
+			keySerialization, sharedDictionariesWithAdditionalKeys,
+		)
+	}
+	if s, ok := c.index.(StateDecoder); ok {
+		if err := s.DecodeState(r); err != nil {
+			return errors.Wrap(err, "index state")
+		}
+	}
+	return nil
+}
+
+// EncodeState implements StateEncoder, ensuring state for index column.
+func (c ColLowCardinality[T]) EncodeState(b *Buffer) {
+	// Writing key serialization version.
+	b.PutInt64(int64(sharedDictionariesWithAdditionalKeys))
+	if s, ok := c.index.(StateEncoder); ok {
+		s.EncodeState(b)
+	}
+}
+
 func (c *ColLowCardinality[T]) DecodeColumn(r *Reader, rows int) error {
 	if rows == 0 {
 		// Skipping entirely of no rows.

query.go

@@ -75,7 +75,7 @@ func (c *Client) sendQuery(ctx context.Context, q Query) error {
 			zap.String("query_id", q.QueryID),
 		)
 	}
-	if c.conn == nil {
+	if c.IsClosed() {
 		return ErrClosed
 	}
 	c.encode(proto.Query{
@@ -537,7 +537,7 @@ func (c *Client) handlePacket(ctx context.Context, p proto.ServerCode, q Query)
 
 // Do performs Query on ClickHouse server.
 func (c *Client) Do(ctx context.Context, q Query) (err error) {
-	if c.conn == nil {
+	if c.IsClosed() {
 		return ErrClosed
 	}
 	if len(q.Parameters) > 0 && !proto.FeatureParameters.In(c.protocolVersion) {