diff --git a/pkg/server/conn.go b/pkg/server/conn.go index cd3ec3740c3ae..99c53f546bf6c 100644 --- a/pkg/server/conn.go +++ b/pkg/server/conn.go @@ -169,7 +169,7 @@ type clientConn struct { lastActive time.Time // last active time authPlugin string // default authentication plugin isUnixSocket bool // connection is Unix Socket file - isClosed int32 // atomic variable to track whether the connection is closed + closeOnce sync.Once // closeOnce is used to make sure clientConn closes only once rsEncoder *column.ResultEncoder // rsEncoder is used to encode the string result to different charsets inputDecoder *util2.InputDecoder // inputDecoder is used to decode the different charsets of incoming strings to utf-8 socketCredUID uint32 // UID from the other end of the Unix Socket @@ -353,29 +353,28 @@ func (cc *clientConn) Close() error { // closeConn is idempotent and thread-safe. // It will be called on the same `clientConn` more than once to avoid connection leak. func closeConn(cc *clientConn, connections int) error { - if !atomic.CompareAndSwapInt32(&cc.isClosed, 0, 1) { - return nil - } - - metrics.ConnGauge.Set(float64(connections)) - if cc.connectionID > 0 { - cc.server.dom.ReleaseConnID(cc.connectionID) - cc.connectionID = 0 - } - if cc.bufReadConn != nil { - err := cc.bufReadConn.Close() - if err != nil { - // We need to expect connection might have already disconnected. - // This is because closeConn() might be called after a connection read-timeout. - logutil.Logger(context.Background()).Debug("could not close connection", zap.Error(err)) + var err error + cc.closeOnce.Do(func() { + metrics.ConnGauge.Set(float64(connections)) + if cc.connectionID > 0 { + cc.server.dom.ReleaseConnID(cc.connectionID) + cc.connectionID = 0 + } + if cc.bufReadConn != nil { + err := cc.bufReadConn.Close() + if err != nil { + // We need to expect connection might have already disconnected. + // This is because closeConn() might be called after a connection read-timeout. + logutil.Logger(context.Background()).Debug("could not close connection", zap.Error(err)) + } } - } - // Close statements and session - // This will release advisory locks, row locks, etc. - if ctx := cc.getCtx(); ctx != nil { - return ctx.Close() - } - return nil + // Close statements and session + // This will release advisory locks, row locks, etc. + if ctx := cc.getCtx(); ctx != nil { + err = ctx.Close() + } + }) + return err } func (cc *clientConn) closeWithoutLock() error { diff --git a/pkg/server/conn_test.go b/pkg/server/conn_test.go index 5a830659bba62..02e9aa3b6dbb7 100644 --- a/pkg/server/conn_test.go +++ b/pkg/server/conn_test.go @@ -2091,5 +2091,4 @@ func TestCloseConn(t *testing.T) { }() } wg.Wait() - require.Equalf(t, int32(1), cc.isClosed, "Expected isClosed to be 1, got 0") }