Skip to content

Commit

Permalink
server: exit tidb-server directly when all connections not in txn (pi…
Browse files Browse the repository at this point in the history
  • Loading branch information
2 people authored and GitHub Enterprise committed Jan 2, 2024
1 parent 9f7b51f commit e176921
Show file tree
Hide file tree
Showing 5 changed files with 44 additions and 12 deletions.
7 changes: 7 additions & 0 deletions server/conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -1163,6 +1163,13 @@ func (cc *clientConn) Run(ctx context.Context) {
return
}

// Should check InTxn() to avoid execute `begin` stmt.
if cc.server.inShutdownMode.Load() {
if !cc.ctx.GetSessionVars().InTxn() {
return
}
}

if !atomic.CompareAndSwapInt32(&cc.status, connStatusReading, connStatusDispatching) {
return
}
Expand Down
30 changes: 29 additions & 1 deletion server/conn_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -764,7 +764,7 @@ func TestConnExecutionTimeout(t *testing.T) {
}

func TestShutDown(t *testing.T) {
store := testkit.CreateMockStore(t)
store, dom := testkit.CreateMockStoreAndDomain(t)

cc := &clientConn{}
se, err := session.CreateSession4Test(store)
Expand All @@ -776,6 +776,34 @@ func TestShutDown(t *testing.T) {
// assert ErrQueryInterrupted
err = cc.handleQuery(context.Background(), "select 1")
require.Equal(t, executor.ErrQueryInterrupted, err)

cfg := newTestConfig()
cfg.Port = 0
cfg.Status.StatusPort = 0
drv := NewTiDBDriver(store)
srv, err := NewServer(cfg, drv)
require.NoError(t, err)
srv.SetDomain(dom)

cc = &clientConn{server: srv}
cc.setCtx(tc)

// test in txn
srv.clients[cc.connectionID+1] = cc
cc.getCtx().GetSessionVars().SetInTxn(true)

waitTime := 100 * time.Millisecond
begin := time.Now()
srv.DrainClients(waitTime, waitTime)
require.Greater(t, time.Since(begin), waitTime)

// test not in txn
srv.clients[cc.connectionID+2] = cc
cc.getCtx().GetSessionVars().SetInTxn(false)

begin = time.Now()
srv.DrainClients(waitTime, waitTime)
require.Less(t, time.Since(begin), waitTime)
}

type snapshotCache interface {
Expand Down
9 changes: 3 additions & 6 deletions server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ import (
"crypto/tls"
"fmt"
"io"
"math/rand"
"net"
"net/http" //nolint:goimports
// For pprof
Expand Down Expand Up @@ -334,9 +333,6 @@ func NewServer(cfg *config.Config, driver IDriver) (*Server, error) {
}
}

// Init rand seed for randomBuf()
rand.Seed(time.Now().UTC().UnixNano())

variable.RegisterStatistics(s)

return s, nil
Expand Down Expand Up @@ -556,8 +552,6 @@ func (s *Server) closeListener() {
metrics.ServerEventCounter.WithLabelValues(metrics.EventClose).Inc()
}

var gracefulCloseConnectionsTimeout = 15 * time.Second

// Close closes the server.
func (s *Server) Close() {
s.startShutdown()
Expand Down Expand Up @@ -891,6 +885,9 @@ func (s *Server) DrainClients(drainWait time.Duration, cancelWait time.Duration)
go func() {
defer close(allDone)
for _, conn := range conns {
if !conn.getCtx().GetSessionVars().InTxn() {
continue
}
select {
case <-conn.quit:
case <-quitWaitingForConns:
Expand Down
6 changes: 3 additions & 3 deletions tidb-server/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -224,9 +224,9 @@ func main() {
terror.RegisterFinish()

exited := make(chan struct{})
signal.SetupSignalHandler(func(graceful bool) {
signal.SetupSignalHandler(func() {
svr.Close()
cleanup(svr, storage, dom, graceful)
cleanup(svr, storage, dom)
cpuprofile.StopCPUProfiler()
close(exited)
})
Expand Down Expand Up @@ -839,7 +839,7 @@ func closeDomainAndStorage(storage kv.Storage, dom *domain.Domain) {
// We should better provider a dynamic way to set this value.
var gracefulCloseConnectionsTimeout = 15 * time.Second

func cleanup(svr *server.Server, storage kv.Storage, dom *domain.Domain, _ bool) {
func cleanup(svr *server.Server, storage kv.Storage, dom *domain.Domain) {
dom.StopAutoAnalyze()

drainClientWait := gracefulCloseConnectionsTimeout
Expand Down
4 changes: 2 additions & 2 deletions util/signal/signal_posix.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import (
)

// SetupSignalHandler setup signal handler for TiDB Server
func SetupSignalHandler(shutdownFunc func(bool)) {
func SetupSignalHandler(shutdownFunc func()) {
usrDefSignalChan := make(chan os.Signal, 1)

signal.Notify(usrDefSignalChan, syscall.SIGUSR1)
Expand All @@ -52,6 +52,6 @@ func SetupSignalHandler(shutdownFunc func(bool)) {
go func() {
sig := <-closeSignalChan
logutil.BgLogger().Info("got signal to exit", zap.Stringer("signal", sig))
shutdownFunc(sig != syscall.SIGHUP)
shutdownFunc()
}()
}

0 comments on commit e176921

Please sign in to comment.