diff --git a/pkg/frontend/mysql_buffer.go b/pkg/frontend/mysql_buffer.go index bd607f21cbec3..0b5d059ff4005 100644 --- a/pkg/frontend/mysql_buffer.go +++ b/pkg/frontend/mysql_buffer.go @@ -202,6 +202,7 @@ func (block *MemBlock) Adjust() { type Conn struct { id uint64 conn net.Conn + livenessProbe *socketLivenessProbe localAddr, remoteAddr string sequenceId uint8 header [4]byte @@ -311,6 +312,7 @@ func NewIOSessionWithOptions( c := &Conn{ conn: conn, + livenessProbe: newSocketLivenessProbe(conn), localAddr: conn.LocalAddr().String(), remoteAddr: conn.RemoteAddr().String(), fixBuf: MemBlock{}, @@ -358,6 +360,7 @@ func (c *Conn) GetSequenceID() uint8 { func (c *Conn) UseConn(conn net.Conn) { c.conn = conn + c.livenessProbe = newSocketLivenessProbe(conn) } func (c *Conn) freeDynamicBuffUnsafe() { diff --git a/pkg/frontend/routine_manager.go b/pkg/frontend/routine_manager.go index eaa791758bc7d..3af64c3738562 100644 --- a/pkg/frontend/routine_manager.go +++ b/pkg/frontend/routine_manager.go @@ -20,7 +20,6 @@ import ( "crypto/x509" "fmt" "math" - "net" "os" "sync" "time" @@ -40,6 +39,8 @@ import ( type RoutineManager struct { mu sync.RWMutex + disconnectProbeMu sync.Mutex + disconnectProbeScratch []activeClientRequest ctx context.Context clients map[*Conn]*Routine workerWG sync.WaitGroup @@ -73,6 +74,11 @@ type activeClientRequest struct { routine *Routine } +var clientDisconnectProbeErrorEvent = logutil.Event{ + Name: "frontend.client-disconnect-probe.error", + Message: "failed to probe active client connection", +} + func NewKillRecord(killtime time.Time, version uint64) KillRecord { return KillRecord{ killTime: killtime, @@ -195,11 +201,14 @@ func (rm *RoutineManager) getRoutineByConnID(id uint32) *Routine { return nil } -func (rm *RoutineManager) longRunningRequests(now time.Time, minimum time.Duration) []activeClientRequest { +func (rm *RoutineManager) appendLongRunningRequests( + requests []activeClientRequest, + now time.Time, + minimum time.Duration, +) []activeClientRequest { rm.mu.RLock() defer rm.mu.RUnlock() nowValue := clientRequestClockValue(now) - var requests []activeClientRequest for conn, routine := range rm.clients { if conn != nil && routine != nil && routine.requestRunningLongerThan(nowValue, minimum) { requests = append(requests, activeClientRequest{conn: conn, routine: routine}) @@ -211,15 +220,27 @@ func (rm *RoutineManager) longRunningRequests(now time.Time, minimum time.Durati func (rm *RoutineManager) cancelDisconnectedRequests( now time.Time, minimum time.Duration, - probe func(net.Conn) (bool, error), + probe func(*Conn) (bool, error), ) { if probe == nil { return } - for _, request := range rm.longRunningRequests(now, minimum) { - closed, err := probe(request.conn.RawConn()) + rm.disconnectProbeMu.Lock() + defer rm.disconnectProbeMu.Unlock() + requests := rm.appendLongRunningRequests(rm.disconnectProbeScratch[:0], now, minimum) + defer func() { + clear(requests) + rm.disconnectProbeScratch = requests[:0] + }() + for _, request := range requests { + closed, err := probe(request.conn) if err != nil { - logutil.Debugf("failed to probe active client connection %s: %v", request.conn.RemoteAddress(), err) + clientDisconnectProbeErrorEvent.DebugLazy(func() []zap.Field { + return []zap.Field{ + zap.String("connection", request.conn.RemoteAddress()), + zap.Error(err), + } + }) continue } if !closed { diff --git a/pkg/frontend/routine_manager_test.go b/pkg/frontend/routine_manager_test.go index 9fc473998a7f7..9c7601063ddda 100644 --- a/pkg/frontend/routine_manager_test.go +++ b/pkg/frontend/routine_manager_test.go @@ -342,9 +342,9 @@ func TestRoutineManagerCancelDisconnectedLongRunningRequests(t *testing.T) { }} probes := 0 - rm.cancelDisconnectedRequests(now, grace, func(conn net.Conn) (bool, error) { + rm.cancelDisconnectedRequests(now, grace, func(conn *Conn) (bool, error) { probes++ - return conn == longServer, nil + return conn.RawConn() == longServer, nil }) require.Equal(t, 1, probes, "only requests beyond the grace period should be probed") @@ -359,13 +359,44 @@ func TestRoutineManagerCancelDisconnectedLongRunningRequests(t *testing.T) { default: } - rm.cancelDisconnectedRequests(now, grace, func(net.Conn) (bool, error) { + rm.cancelDisconnectedRequests(now, grace, func(*Conn) (bool, error) { probes++ return true, nil }) require.Equal(t, 1, probes, "a routine already closing should not be probed again") } +func TestClientDisconnectProbePolicyCoversNewRequests(t *testing.T) { + now := time.Now() + serverConn, clientConn := net.Pipe() + t.Cleanup(func() { + _ = serverConn.Close() + _ = clientConn.Close() + }) + + routine := NewRoutine(context.Background(), &testMysqlWriter{}, &config.FrontendParameters{}) + t.Cleanup(routine.cancelRoutineFunc) + routine.requestStartedAt.Store(clientRequestClockValue(now)) + requestCtx, cancelRequest := context.WithCancel(context.Background()) + t.Cleanup(cancelRequest) + routine.setCancelRequestFunc(cancelRequest) + + conn := &Conn{conn: serverConn, remoteAddr: "new-request"} + rm := &RoutineManager{clients: map[*Conn]*Routine{conn: routine}} + probes := 0 + rm.cancelDisconnectedRequests(now, clientDisconnectProbeGrace, func(*Conn) (bool, error) { + probes++ + return true, nil + }) + + require.Equal(t, 1, probes, "a new active request must be probed without an age grace period") + select { + case <-requestCtx.Done(): + case <-time.After(time.Second): + t.Fatal("a disconnected new request was not canceled") + } +} + func TestRoutineManagerProbeErrorDoesNotCancelRequest(t *testing.T) { now := time.Now() serverConn, clientConn := net.Pipe() @@ -383,7 +414,7 @@ func TestRoutineManagerProbeErrorDoesNotCancelRequest(t *testing.T) { conn := &Conn{conn: serverConn} rm := &RoutineManager{clients: map[*Conn]*Routine{conn: routine}} - rm.cancelDisconnectedRequests(now, 30*time.Second, func(net.Conn) (bool, error) { + rm.cancelDisconnectedRequests(now, 30*time.Second, func(*Conn) (bool, error) { return false, errors.New("probe failed") }) @@ -436,7 +467,8 @@ func BenchmarkRoutineManagerLongRunningRequests(b *testing.B) { b.ReportAllocs() b.ResetTimer() for i := 0; i < b.N; i++ { - _ = rm.longRunningRequests(now, 30*time.Second) + requests := rm.appendLongRunningRequests(nil, now, 30*time.Second) + clear(requests) } }) } diff --git a/pkg/frontend/server.go b/pkg/frontend/server.go index 529537e5cfb22..0ce712a820427 100644 --- a/pkg/frontend/server.go +++ b/pkg/frontend/server.go @@ -64,8 +64,10 @@ var initConnectionID uint32 = 1000 var ConnIDAllocKey = "____server_conn_id" const ( - clientDisconnectProbeInterval = 5 * time.Second - clientDisconnectProbeGrace = 30 * time.Second + // The request handler owns the connection read loop while a statement is + // executing, so probe every active request from the first monitor tick. + clientDisconnectProbeInterval = time.Second + clientDisconnectProbeGrace = 0 ) // MOServer MatrixOne Server diff --git a/pkg/frontend/socket_liveness_alloc_unix_test.go b/pkg/frontend/socket_liveness_alloc_unix_test.go new file mode 100644 index 0000000000000..85ed503fde38c --- /dev/null +++ b/pkg/frontend/socket_liveness_alloc_unix_test.go @@ -0,0 +1,47 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build (darwin || linux) && !race + +package frontend + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// The race runtime adds bookkeeping allocations to the syscall callback path. +// Functional live-socket behavior remains covered under -race by the tests in +// socket_liveness_unix_test.go. +func TestConnectionPeerClosedLivePathDoesNotAllocate(t *testing.T) { + server, client := tcpConnectionPair(t) + t.Cleanup(func() { + _ = server.Close() + _ = client.Close() + }) + conn := &Conn{ + conn: server, + livenessProbe: newSocketLivenessProbe(server), + } + + var closed bool + var err error + allocs := testing.AllocsPerRun(100, func() { + closed, err = connectionPeerClosed(conn) + }) + require.NoError(t, err) + require.False(t, closed) + require.Zero(t, allocs) +} diff --git a/pkg/frontend/socket_liveness_other.go b/pkg/frontend/socket_liveness_other.go index 1c48e66ade910..919f600e3a600 100644 --- a/pkg/frontend/socket_liveness_other.go +++ b/pkg/frontend/socket_liveness_other.go @@ -18,6 +18,18 @@ package frontend import "net" -func connectionPeerClosed(net.Conn) (bool, error) { +type socketLivenessProbe struct{} + +func newSocketLivenessProbe(net.Conn) *socketLivenessProbe { + return &socketLivenessProbe{} +} + +func (*socketLivenessProbe) connectionPeerClosed() (bool, error) { return false, nil } + +func rawConnectionPeerClosed(net.Conn) (bool, error) { + return false, nil +} + +func connectionPeerClosed(*Conn) (bool, error) { return false, nil } diff --git a/pkg/frontend/socket_liveness_unix.go b/pkg/frontend/socket_liveness_unix.go index 772448ea0bf19..b35451d0ecd45 100644 --- a/pkg/frontend/socket_liveness_unix.go +++ b/pkg/frontend/socket_liveness_unix.go @@ -25,62 +25,115 @@ import ( "golang.org/x/sys/unix" ) -// connectionPeerClosed checks the socket read side without consuming protocol -// bytes. For TLS, peeking at the underlying encrypted stream is sufficient: -// only EOF/error is interpreted, never payload. -func connectionPeerClosed(conn net.Conn) (bool, error) { +// socketLivenessProbe caches the non-owning RawConn handle and callback state +// for one frontend connection. RoutineManager serializes monitor passes, so a +// probe is never used concurrently and the hot path needs no per-probe lock or +// allocation. +type socketLivenessProbe struct { + rawConn syscall.RawConn + initErr error + nilConn bool + pollFDs [1]unix.PollFd + peek [1]byte + n int + peerClosed bool + probeErr error + control func(uintptr) +} + +func newSocketLivenessProbe(conn net.Conn) *socketLivenessProbe { + probe := &socketLivenessProbe{nilConn: conn == nil} if conn == nil { - return true, nil + return probe } if tlsConn, ok := conn.(*tls.Conn); ok { conn = tlsConn.NetConn() } syscallConn, ok := conn.(syscall.Conn) if !ok { - return false, nil + return probe } - rawConn, err := syscallConn.SyscallConn() - if err != nil { - return false, err + probe.rawConn, probe.initErr = syscallConn.SyscallConn() + if probe.initErr == nil { + probe.control = probe.probeFD } + return probe +} - var ( - n int - peerClosed bool - probeErr error - ) - if err = rawConn.Control(func(fd uintptr) { - pollFDs := []unix.PollFd{{ - Fd: int32(fd), - Events: unix.POLLIN | unix.POLLHUP | unix.POLLERR | socketReadHangupPollEvent(), - }} - if _, probeErr = unix.Poll(pollFDs, 0); probeErr != nil { - return - } - if pollFDs[0].Revents&(unix.POLLHUP|unix.POLLERR|unix.POLLNVAL|socketReadHangupPollEvent()) != 0 { - peerClosed = true - return - } - var one [1]byte - n, _, probeErr = unix.Recvfrom(int(fd), one[:], unix.MSG_PEEK|unix.MSG_DONTWAIT) - }); err != nil { +// connectionPeerClosed checks the socket read side without consuming protocol +// bytes. For TLS, peeking at the underlying encrypted stream is sufficient: +// only EOF/error is interpreted, never payload. +func (probe *socketLivenessProbe) connectionPeerClosed() (bool, error) { + if probe == nil || probe.nilConn { + return true, nil + } + if probe.initErr != nil { + return false, probe.initErr + } + if probe.rawConn == nil { + return false, nil + } + + probe.n = -1 + probe.peerClosed = false + probe.probeErr = nil + probe.pollFDs[0].Revents = 0 + if err := probe.rawConn.Control(probe.control); err != nil { return false, err } - if peerClosed { + if probe.peerClosed { return true, nil } - if probeErr == nil { - return n == 0, nil + if probe.probeErr == nil { + return probe.n == 0, nil } - if errors.Is(probeErr, unix.EAGAIN) || - errors.Is(probeErr, unix.EWOULDBLOCK) || - errors.Is(probeErr, unix.EINTR) { + if errors.Is(probe.probeErr, unix.EAGAIN) || + errors.Is(probe.probeErr, unix.EWOULDBLOCK) || + errors.Is(probe.probeErr, unix.EINTR) { return false, nil } - if errors.Is(probeErr, unix.ECONNRESET) || - errors.Is(probeErr, unix.ENOTCONN) || - errors.Is(probeErr, unix.EBADF) { + if errors.Is(probe.probeErr, unix.ECONNRESET) || + errors.Is(probe.probeErr, unix.ENOTCONN) || + errors.Is(probe.probeErr, unix.EBADF) { return true, nil } - return false, probeErr + return false, probe.probeErr +} + +func (probe *socketLivenessProbe) probeFD(fd uintptr) { + probe.pollFDs[0] = unix.PollFd{ + Fd: int32(fd), + Events: unix.POLLIN | unix.POLLHUP | unix.POLLERR | socketReadHangupPollEvent(), + } + if _, probe.probeErr = unix.Poll(probe.pollFDs[:], 0); probe.probeErr != nil { + return + } + revents := probe.pollFDs[0].Revents + if revents&(unix.POLLHUP|unix.POLLERR|unix.POLLNVAL|socketReadHangupPollEvent()) != 0 { + probe.peerClosed = true + return + } + // A connected socket with no readable event is the overwhelmingly common + // case. Avoid a second syscall for every active request; recv is only needed + // to distinguish readable payload from EOF. + if revents&unix.POLLIN == 0 { + return + } + probe.n, _, probe.probeErr = unix.Recvfrom( + int(fd), probe.peek[:], unix.MSG_PEEK|unix.MSG_DONTWAIT, + ) +} + +func rawConnectionPeerClosed(conn net.Conn) (bool, error) { + return newSocketLivenessProbe(conn).connectionPeerClosed() +} + +func connectionPeerClosed(conn *Conn) (bool, error) { + if conn == nil { + return true, nil + } + if conn.livenessProbe == nil { + return rawConnectionPeerClosed(conn.RawConn()) + } + return conn.livenessProbe.connectionPeerClosed() } diff --git a/pkg/frontend/socket_liveness_unix_test.go b/pkg/frontend/socket_liveness_unix_test.go index 91a8c61cb4ae4..ab0d64c57da7c 100644 --- a/pkg/frontend/socket_liveness_unix_test.go +++ b/pkg/frontend/socket_liveness_unix_test.go @@ -18,6 +18,7 @@ package frontend import ( "crypto/tls" + "fmt" "net" "testing" "time" @@ -25,7 +26,7 @@ import ( "github.com/stretchr/testify/require" ) -func tcpConnectionPair(t *testing.T) (net.Conn, net.Conn) { +func tcpConnectionPair(t testing.TB) (net.Conn, net.Conn) { t.Helper() listener, err := net.Listen("tcp", "127.0.0.1:0") require.NoError(t, err) @@ -55,6 +56,51 @@ func tcpConnectionPair(t *testing.T) (net.Conn, net.Conn) { return nil, nil } +func tcpConnectionPairs(t testing.TB, count int) ([]net.Conn, []net.Conn) { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = listener.Close() }) + servers := make([]net.Conn, 0, count) + clients := make([]net.Conn, 0, count) + t.Cleanup(func() { + for _, conn := range servers { + _ = conn.Close() + } + for _, conn := range clients { + _ = conn.Close() + } + }) + + accepted := make(chan net.Conn, count) + acceptErr := make(chan error, 1) + go func() { + for range count { + conn, err := listener.Accept() + if err != nil { + acceptErr <- err + return + } + accepted <- conn + } + }() + + for range count { + client, err := net.Dial("tcp", listener.Addr().String()) + require.NoError(t, err) + clients = append(clients, client) + } + for range count { + select { + case server := <-accepted: + servers = append(servers, server) + case err := <-acceptErr: + require.NoError(t, err) + } + } + return servers, clients +} + func TestConnectionPeerClosedDoesNotConsumeProtocolBytes(t *testing.T) { server, client := tcpConnectionPair(t) t.Cleanup(func() { @@ -62,13 +108,13 @@ func TestConnectionPeerClosedDoesNotConsumeProtocolBytes(t *testing.T) { _ = client.Close() }) - closed, err := connectionPeerClosed(server) + closed, err := rawConnectionPeerClosed(server) require.NoError(t, err) require.False(t, closed) _, err = client.Write([]byte{0x2a}) require.NoError(t, err) - closed, err = connectionPeerClosed(server) + closed, err = rawConnectionPeerClosed(server) require.NoError(t, err) require.False(t, closed) @@ -85,7 +131,7 @@ func TestConnectionPeerClosedDetectsDisconnect(t *testing.T) { require.NoError(t, client.Close()) require.Eventually(t, func() bool { - closed, err := connectionPeerClosed(server) + closed, err := rawConnectionPeerClosed(server) return err == nil && closed }, time.Second, time.Millisecond) } @@ -98,7 +144,7 @@ func TestConnectionPeerClosedDetectsDisconnectBehindUnreadBytes(t *testing.T) { require.NoError(t, err) require.NoError(t, client.Close()) require.Eventually(t, func() bool { - closed, err := connectionPeerClosed(server) + closed, err := rawConnectionPeerClosed(server) return err == nil && closed }, time.Second, time.Millisecond) } @@ -110,13 +156,91 @@ func TestConnectionPeerClosedUnwrapsTLS(t *testing.T) { require.NoError(t, client.Close()) require.Eventually(t, func() bool { - closed, err := connectionPeerClosed(tlsServer) + closed, err := rawConnectionPeerClosed(tlsServer) return err == nil && closed }, time.Second, time.Millisecond) } func TestConnectionPeerClosedNilConnection(t *testing.T) { - closed, err := connectionPeerClosed(nil) + closed, err := rawConnectionPeerClosed(nil) require.NoError(t, err) require.True(t, closed) } + +func TestConnectionPeerClosedRefreshesProbeWhenConnChanges(t *testing.T) { + firstServer, firstClient := tcpConnectionPair(t) + secondServer, secondClient := tcpConnectionPair(t) + t.Cleanup(func() { + _ = firstServer.Close() + _ = firstClient.Close() + _ = secondServer.Close() + }) + + conn := &Conn{ + conn: firstServer, + livenessProbe: newSocketLivenessProbe(firstServer), + } + conn.UseConn(secondServer) + require.NoError(t, secondClient.Close()) + require.Eventually(t, func() bool { + closed, err := connectionPeerClosed(conn) + return err == nil && closed + }, time.Second, time.Millisecond) +} + +func BenchmarkRoutineManagerClientDisconnectProbe(b *testing.B) { + server, client := tcpConnectionPair(b) + b.Cleanup(func() { + _ = server.Close() + _ = client.Close() + }) + now := time.Now() + + for _, population := range []struct { + name string + connections int + distinctFDs bool + }{ + {name: "10k-distinct-fds", connections: 10_000, distinctFDs: true}, + {name: "10k-shared-fd", connections: 10_000}, + {name: "100k-shared-fd", connections: 100_000}, + } { + connections := population.connections + probeConnections := []net.Conn{server} + if population.distinctFDs { + // Use distinct live TCP sockets for the realistic population. The + // shared-fd cases isolate the bounded scan and exact syscall count; the + // 100k upper bound cannot use one loopback destination because it does + // not supply 100k distinct ephemeral client ports. + probeConnections, _ = tcpConnectionPairs(b, connections) + } + for _, activePercent := range []int{0, 1, 10, 100} { + b.Run(fmt.Sprintf("%s/active=%d%%", population.name, activePercent), func(b *testing.B) { + active := connections * activePercent / 100 + rm := &RoutineManager{clients: make(map[*Conn]*Routine, connections)} + for i := 0; i < connections; i++ { + routine := &Routine{} + if i < active { + routine.requestStartedAt.Store(clientRequestClockValue(now)) + } + rawConn := probeConnections[i%len(probeConnections)] + rm.clients[&Conn{ + conn: rawConn, + livenessProbe: newSocketLivenessProbe(rawConn), + }] = routine + } + + b.ReportMetric(float64(connections), "connections") + b.ReportMetric(float64(active), "active_connections") + b.ReportAllocs() + // Production reuses the manager-owned request snapshot after the + // first tick. Keep setup growth outside the steady-state budget. + rm.cancelDisconnectedRequests(now, clientDisconnectProbeGrace, connectionPeerClosed) + b.ResetTimer() + for i := 0; i < b.N; i++ { + rm.cancelDisconnectedRequests(now, clientDisconnectProbeGrace, connectionPeerClosed) + } + }) + } + } +} diff --git a/suites/scenarios/14_issue_regression/issue_27595_direct_disconnect_release.py b/suites/scenarios/14_issue_regression/issue_27595_direct_disconnect_release.py new file mode 100644 index 0000000000000..4a1278c68175a --- /dev/null +++ b/suites/scenarios/14_issue_regression/issue_27595_direct_disconnect_release.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 + +import os +import socket +import sys +import threading +import time + +import pymysql + + +HOST = os.getenv("MO_HOST", "127.0.0.1") +PORT = int(os.getenv("MO_PORT", "6001")) +USER = os.getenv("MO_USER", "root") +PASSWORD = os.getenv("MO_PASSWORD", "111") +DATABASE = os.getenv("MO_DATABASE", "issue_27595") + + +def connect(*, database=None, autocommit=True, read_timeout=10): + return pymysql.connect( + host=HOST, + port=PORT, + user=USER, + password=PASSWORD, + database=database, + autocommit=autocommit, + connect_timeout=10, + read_timeout=read_timeout, + write_timeout=10, + ) + + +def execute(conn, sql): + with conn.cursor() as cursor: + cursor.execute(sql) + return cursor.fetchall() + + +def wait_for_statement(observer, connection_id, sql_fragment, timeout=5): + deadline = time.monotonic() + timeout + normalized_fragment = " ".join(sql_fragment.lower().split()) + while time.monotonic() < deadline: + for row in execute(observer, "show processlist"): + normalized_cells = [" ".join(str(cell).lower().split()) for cell in row] + if str(connection_id) in normalized_cells and any( + normalized_fragment in cell for cell in normalized_cells + ): + return True + time.sleep(0.05) + return False + + +def disconnect(conn): + conn._sock.shutdown(socket.SHUT_RDWR) + conn._sock.close() + + +def close_quietly(conn): + if conn is not None: + try: + conn.close() + except Exception: + pass + + +def main(): + admin = setup = holder = observer = waiter = None + holder_finished = threading.Event() + holder_result = {} + try: + admin = connect() + execute(admin, f"drop database if exists `{DATABASE}`") + execute(admin, f"create database `{DATABASE}`") + + setup = connect(database=DATABASE) + execute( + setup, + "create table lock_probe (" + "user_id varchar(128) not null, " + "session_id varchar(64) not null, " + "status varchar(20) not null, " + "primary key (user_id, session_id))", + ) + execute(setup, "insert into lock_probe values ('u', 's', 'running')") + + holder = connect(database=DATABASE, autocommit=False, read_timeout=70) + observer = connect(database=DATABASE) + waiter = connect(database=DATABASE, autocommit=False, read_timeout=5) + connection_id = execute(holder, "select connection_id()")[0][0] + locked = execute( + holder, + "select status from lock_probe " + "where user_id = 'u' and session_id = 's' for update", + ) + if locked != (("running",),): + raise AssertionError(f"holder did not acquire the expected row: {locked}") + + def run_sleep(): + try: + execute(holder, "select sleep(60)") + holder_result["success"] = True + except Exception as exc: + holder_result["error"] = repr(exc) + finally: + holder_finished.set() + + worker = threading.Thread(target=run_sleep, daemon=True) + worker.start() + if not wait_for_statement(observer, connection_id, "select sleep(60)"): + raise AssertionError("holder did not enter SELECT SLEEP(60)") + + disconnect(holder) + started = time.monotonic() + rows = execute( + waiter, + "select status from lock_probe " + "where user_id = 'u' and session_id = 's' for update", + ) + elapsed = time.monotonic() - started + execute(waiter, "commit") + + if rows != (("running",),): + print(f"FAIL: waiter returned unexpected rows: {rows}") + return 1 + if elapsed >= 5: + print(f"FAIL: disconnected holder retained its lock for {elapsed:.3f}s") + return 1 + if not holder_finished.wait(5): + print("FAIL: disconnected holder statement did not terminate") + return 1 + if holder_result.get("success"): + print("FAIL: disconnected holder statement returned success") + return 1 + + worker.join() + print(f"PASS: disconnected holder released its lock in {elapsed:.3f}s") + return 0 + finally: + close_quietly(waiter) + close_quietly(observer) + close_quietly(holder) + close_quietly(setup) + if admin is not None: + try: + execute(admin, f"drop database if exists `{DATABASE}`") + except Exception: + pass + close_quietly(admin) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/suites/scenarios/14_issue_regression/issue_27595_direct_disconnect_release_test.py b/suites/scenarios/14_issue_regression/issue_27595_direct_disconnect_release_test.py new file mode 100644 index 0000000000000..ac0ef88df0ac6 --- /dev/null +++ b/suites/scenarios/14_issue_regression/issue_27595_direct_disconnect_release_test.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 + +import importlib.util +import socket +import threading +import unittest +from pathlib import Path +from unittest import mock + + +SCRIPT = Path(__file__).with_name("issue_27595_direct_disconnect_release.py") +SPEC = importlib.util.spec_from_file_location("issue_27595", SCRIPT) +ISSUE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(ISSUE) + + +class FakeSocket: + def __init__(self, disconnected): + self.disconnected = disconnected + self.shutdown_how = None + self.closed = False + + def shutdown(self, how): + self.shutdown_how = how + self.disconnected.set() + + def close(self): + self.closed = True + + +class FakeConnection: + def __init__(self, name, disconnected=None): + self.name = name + self.closed = False + self._sock = FakeSocket(disconnected) if disconnected else None + + def close(self): + self.closed = True + + +class DirectDisconnectReleaseTest(unittest.TestCase): + def make_connections(self): + disconnected = threading.Event() + connections = [ + FakeConnection("admin"), + FakeConnection("setup"), + FakeConnection("holder", disconnected), + FakeConnection("observer"), + FakeConnection("waiter"), + ] + return connections, disconnected + + def run_main(self, *, waiter_rows=(("running",),), holder_succeeds=False): + connections, disconnected = self.make_connections() + + def execute(conn, sql): + normalized = " ".join(sql.lower().split()) + if conn.name == "holder" and normalized == "select connection_id()": + return ((7,),) + if conn.name == "holder" and "for update" in normalized: + return (("running",),) + if conn.name == "holder" and normalized == "select sleep(60)": + disconnected.wait(2) + if holder_succeeds: + return ((0,),) + raise OSError("client socket closed") + if conn.name == "waiter" and "for update" in normalized: + return waiter_rows + return () + + with ( + mock.patch.object(ISSUE, "connect", side_effect=connections), + mock.patch.object(ISSUE, "execute", side_effect=execute), + mock.patch.object(ISSUE, "wait_for_statement", return_value=True), + mock.patch.object(ISSUE.time, "monotonic", side_effect=[10, 10.5]), + ): + result = ISSUE.main() + return result, connections + + def test_connect_passes_network_and_timeout_options(self): + with mock.patch.object(ISSUE.pymysql, "connect", return_value="connection") as connect: + result = ISSUE.connect(database="db", autocommit=False, read_timeout=5) + + self.assertEqual("connection", result) + connect.assert_called_once_with( + host=ISSUE.HOST, + port=ISSUE.PORT, + user=ISSUE.USER, + password=ISSUE.PASSWORD, + database="db", + autocommit=False, + connect_timeout=10, + read_timeout=5, + write_timeout=10, + ) + + def test_wait_for_statement_matches_connection_and_sql(self): + observer = FakeConnection("observer") + processlist = ( + ("node", 6, "other", "select sleep(60)"), + ("node", 7, "holder", "SELECT SLEEP(60)"), + ) + with ( + mock.patch.object(ISSUE, "execute", return_value=processlist), + mock.patch.object(ISSUE.time, "monotonic", side_effect=[0, 0.1]), + ): + self.assertTrue(ISSUE.wait_for_statement(observer, 7, "select sleep(60)")) + + def test_main_releases_lock_after_direct_disconnect(self): + result, connections = self.run_main() + + self.assertEqual(0, result) + self.assertEqual(socket.SHUT_RDWR, connections[2]._sock.shutdown_how) + self.assertTrue(connections[2]._sock.closed) + self.assertTrue(all(conn.closed for conn in connections)) + + def test_main_rejects_unexpected_waiter_rows(self): + result, _ = self.run_main(waiter_rows=()) + self.assertEqual(1, result) + + def test_main_rejects_successful_disconnected_statement(self): + result, _ = self.run_main(holder_succeeds=True) + self.assertEqual(1, result) + + def test_main_requires_observed_sleep_statement(self): + connections, _ = self.make_connections() + with ( + mock.patch.object(ISSUE, "connect", side_effect=connections), + mock.patch.object(ISSUE, "execute", return_value=(("running",),)), + mock.patch.object(ISSUE, "wait_for_statement", return_value=False), + ): + with self.assertRaisesRegex(AssertionError, "did not enter"): + ISSUE.main() + + +if __name__ == "__main__": + unittest.main()