From b05859dc00d24f7ba4c033f47aa7147c4403c096 Mon Sep 17 00:00:00 2001 From: Rafi Shamim Date: Sat, 11 Mar 2023 18:02:34 -0500 Subject: [PATCH] sql: fix check for closing connExecutor during draining This fixes a minor bug in which the connection would not get closed at the earliest possible time during server shutdown. The connection is supposed to be closed as soon as we handle a Sync message when the conn_executor is in the draining state and not in a transaction. Since the transaction state was checked before state transitions occurred, this would cause the connection to remain open for an extra bit of time. Release note: None --- pkg/sql/conn_executor.go | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/pkg/sql/conn_executor.go b/pkg/sql/conn_executor.go index 9dfd019faa65..9facd8a25f35 100644 --- a/pkg/sql/conn_executor.go +++ b/pkg/sql/conn_executor.go @@ -1911,6 +1911,7 @@ func (ex *connExecutor) run( var err error if err = ex.execCmd(); err != nil { + // Both of these errors are normal ways for the connExecutor to exit. if errors.IsAny(err, io.EOF, errDrainingComplete) { return nil } @@ -1934,7 +1935,7 @@ var errDrainingComplete = fmt.Errorf("draining done. this is a good time to fini // Returns drainingComplete if the session should finish because draining is // complete (i.e. we received a DrainRequest - possibly previously - and the // connection is found to be idle). -func (ex *connExecutor) execCmd() error { +func (ex *connExecutor) execCmd() (retErr error) { ctx := ex.Ctx() cmd, pos, err := ex.stmtBuf.CurCmd() if err != nil { @@ -2121,16 +2122,16 @@ func (ex *connExecutor) execCmd() error { // Note that the Sync result will flush results to the network connection. res = ex.clientComm.CreateSyncResult(pos) if ex.draining { - // If we're draining, check whether this is a good time to finish the + // If we're draining, then after handing the Sync connExecutor state + // transition, check whether this is a good time to finish the // connection. If we're not inside a transaction, we stop processing // now. If we are inside a transaction, we'll check again the next time // a Sync is processed. - if ex.idleConn() { - // If we're about to close the connection, close res in order to flush - // now, as we won't have an opportunity to do it later. - res.Close(ctx, stateToTxnStatusIndicator(ex.machine.CurState())) - return errDrainingComplete - } + defer func() { + if ex.idleConn() { + retErr = errDrainingComplete + } + }() } case CopyIn: ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionQueryReceived, tcmd.TimeReceived)