package tidbruntime

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

const (
	defaultTiDBPort              = 4000
	defaultTiDBStatusPort        = 10080
	defaultTiDBPDPort            = 2379
	defaultTiDBKVPort            = 20160
	sharedTiUPPlaygroundLockName = "tiup-playground.lock"
	sharedTiDBTemplateLockName   = "tidb-template.lock"
	sharedTiDBCloneLockPrefix    = "tidb-clone"
	sharedTiDBCloneQueueLockName = "tidb-clone-queue.lock"
	sharedTiDBShutdownDelay      = 5 * time.Second
	defaultTiUPInstallTimeout    = 20 * time.Minute
	defaultTiUPStartTimeout      = 12 * time.Minute
	defaultTiUPStopTimeout       = 30 * time.Second
	defaultTiUPCleanTimeout      = 30 * time.Second
	defaultTiUPPlaygroundVersion = "v8.5.5"
	defaultTiUPLockPollInterval  = 100 * time.Millisecond
	defaultTiDBCloneConcurrency  = 2
	defaultTiDBReusePingTimeout  = 10 * time.Second
)

type sharedTiDBInstance struct {
	adminDSN string
	stop     func() error
}

type tiupPlaygroundSpec struct {
	Version    string
	Tag        string
	Host       string
	PortOffset int
	DBPort     int
	PDPort     int
	KVPort     int
	StatusPort int
}

type tidbBootstrapAdminDB struct {
	db    *sql.DB
	close func() error
}

type sharedTiDBState struct {
	Tag            string `json:"tag"`
	OwnerPID       int    `json:"ownerPid"`
	RootPID        int    `json:"rootPid,omitempty"`
	SessionID      int    `json:"sessionId,omitempty"`
	ProcessGroupID int    `json:"processGroupId,omitempty"`
	AdminDSN       string `json:"adminDsn"`
}

var (
	sharedTiDBMu           sync.Mutex
	sharedTiDB             *sharedTiDBInstance
	sharedTiDBRefs         int
	sharedTiDBInitializing bool
	sharedTiDBInitWaitCh   chan struct{}
	sharedTiDBShutdown     *time.Timer
)

var (
	acquireSharedTiDBRuntimeFn       = acquireSharedTiDBRuntime
	startSharedTiDBInstanceFn        = startSharedTiDBInstance
	startTiUPPlaygroundServerFn      = startTiUPPlaygroundServer
	openTiDBBootstrapAdminDBFn       = openTiDBBootstrapAdminDB
	pingTiDBAdminFn                  = pingTiDBAdmin
	execCommandContextFn             = exec.CommandContext
	sharedTiDBAfterFuncFn            = time.AfterFunc
	sharedTiDBStateDirFn             = sharedTiDBStateDir
	processExistsFn                  = processExists
	stopTiUPPlaygroundStateFn        = stopTiUPPlaygroundState
	stopTiUPPlaygroundProcessGroupFn = stopTiUPPlaygroundProcessGroup
	listTiUPPlaygroundSessionPIDsFn  = listTiUPPlaygroundSessionPIDs
	startTiUPPlaygroundWatchdogFn    = startTiUPPlaygroundWatchdog
	cleanTiUPPlaygroundFn            = cleanTiUPPlayground
	tidbTiUPInstallTimeout           = defaultTiUPInstallTimeout
	tidbTiUPStartTimeout             = defaultTiUPStartTimeout
	tidbTiUPStopTimeout              = defaultTiUPStopTimeout
	tidbTiUPCleanTimeout             = defaultTiUPCleanTimeout
)

func resolveTiDBRootDSN(t *testing.T) string {
	t.Helper()

	if trimmed := strings.TrimSpace(os.Getenv("AGENTTEAM_TIDB_TEST_DSN")); trimmed != "" {
		return trimmed
	}

	return acquireSharedTiDBRuntimeFn(t).adminDSN
}

func resolveTiDBRootDSNWithCleanup(ctx context.Context) (string, func(), error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if trimmed := strings.TrimSpace(os.Getenv("AGENTTEAM_TIDB_TEST_DSN")); trimmed != "" {
		return trimmed, func() {}, nil
	}

	instance, err := acquireSharedTiDBRuntimeInstance(ctx)
	if err != nil {
		return "", nil, err
	}
	var releaseOnce sync.Once
	return instance.adminDSN, func() {
		releaseOnce.Do(func() {
			releaseSharedTiDBRuntime(instance)
		})
	}, nil
}

func acquireSharedTiDBRuntime(t *testing.T) *sharedTiDBInstance {
	t.Helper()

	instance, err := acquireSharedTiDBRuntimeInstance(context.Background())
	skipSharedTiDBIfNeeded(t, err)
	require.NoError(t, err)

	t.Cleanup(func() { releaseSharedTiDBRuntime(instance) })
	return instance
}

func acquireSharedTiDBRuntimeInstance(ctx context.Context) (*sharedTiDBInstance, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	for {
		if err := ctx.Err(); err != nil {
			return nil, err
		}

		sharedTiDBMu.Lock()
		for sharedTiDBInitializing {
			waitCh := sharedTiDBInitWaitCh
			sharedTiDBMu.Unlock()
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-waitCh:
			}
			sharedTiDBMu.Lock()
		}
		if sharedTiDBShutdown != nil {
			sharedTiDBShutdown.Stop()
			sharedTiDBShutdown = nil
		}
		if sharedTiDB != nil {
			instance := sharedTiDB
			if sharedTiDBRefs == 0 {
				sharedTiDBMu.Unlock()
				if !pingSharedTiDBAdminDSN(instance.adminDSN, defaultTiDBReusePingTimeout) {
					_ = instance.close()
					resetPreparedTiDBTemplateCache()

					sharedTiDBMu.Lock()
					if sharedTiDB == instance && sharedTiDBRefs == 0 {
						sharedTiDB = nil
					}
					sharedTiDBMu.Unlock()
					continue
				}

				sharedTiDBMu.Lock()
				if sharedTiDB != instance {
					sharedTiDBMu.Unlock()
					continue
				}
			}

			sharedTiDBRefs++
			sharedTiDBMu.Unlock()
			return instance, nil
		}

		sharedTiDBInitializing = true
		sharedTiDBInitWaitCh = make(chan struct{})
		sharedTiDBMu.Unlock()

		var (
			instance *sharedTiDBInstance
			err      error
		)
		instance, err = startSharedTiDBInstanceFn(ctx)

		sharedTiDBMu.Lock()
		if err == nil && instance != nil && sharedTiDB == nil {
			sharedTiDB = instance
			sharedTiDBRefs = 1
		}
		sharedTiDBInitializing = false
		close(sharedTiDBInitWaitCh)
		sharedTiDBInitWaitCh = nil
		sharedTiDBMu.Unlock()

		if err != nil {
			return nil, err
		}
		return instance, nil
	}
}

func releaseSharedTiDBRuntime(instance *sharedTiDBInstance) {
	sharedTiDBMu.Lock()
	if sharedTiDB != instance {
		sharedTiDBMu.Unlock()
		return
	}
	if sharedTiDBRefs <= 0 {
		sharedTiDBMu.Unlock()
		return
	}

	sharedTiDBRefs--
	if sharedTiDBRefs > 0 {
		sharedTiDBMu.Unlock()
		return
	}

	if sharedTiDBShutdown != nil {
		sharedTiDBShutdown.Stop()
	}
	sharedTiDBShutdown = sharedTiDBAfterFuncFn(sharedTiDBShutdownDelay, func() {
		sharedTiDBMu.Lock()
		if sharedTiDB != instance || sharedTiDBRefs != 0 {
			sharedTiDBMu.Unlock()
			return
		}
		sharedTiDB = nil
		sharedTiDBShutdown = nil
		sharedTiDBMu.Unlock()

		_ = instance.close()
	})
	sharedTiDBMu.Unlock()
}

func startSharedTiDBInstance(ctx context.Context) (*sharedTiDBInstance, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	var (
		adminDSN string
		stop     func() error
	)
	if err := withSharedTiUPPlaygroundLockContext(ctx, tidbTiUPStartTimeout, func() error {
		state, ok, err := reusableSharedTiDBState()
		if err != nil {
			return err
		}
		if ok {
			leasePath, err := acquireSharedTiDBLease(state.Tag)
			if err != nil {
				return err
			}
			if !processExistsFn(state.OwnerPID) {
				adoptedState, adoptedStop, err := adoptSharedTiDBState(state, leasePath)
				if err != nil {
					_ = removeSharedTiDBLease(leasePath)
					return err
				}
				state = adoptedState
				stop = adoptedStop
			} else {
				stop = func() error {
					return releaseSharedTiDBLease(leasePath)
				}
			}
			adminDSN = state.AdminDSN
			return nil
		}

		if err := cleanupOrphanedSharedTiDBPlaygrounds(); err != nil {
			return fmt.Errorf("cleanup orphaned shared tidb playgrounds: %w", err)
		}

		adminDSN, stop, err = startTiUPPlaygroundServerFn(ctx)
		return err
	}); err != nil {
		return nil, err
	}

	pingCtx, pingCancel := context.WithTimeout(ctx, 10*time.Second)
	defer pingCancel()

	adminDB, err := openTiDBBootstrapAdminDBFn(adminDSN)
	if err != nil {
		_ = stop()
		return nil, fmt.Errorf("connect shared tidb admin db: %w", err)
	}
	defer func() { _ = adminDB.Close() }()

	if err := pingTiDBAdminFn(pingCtx, adminDB.db); err != nil {
		_ = stop()
		return nil, fmt.Errorf("ping shared tidb admin db: %w", err)
	}

	return &sharedTiDBInstance{
		adminDSN: adminDSN,
		stop:     stop,
	}, nil
}

func adoptSharedTiDBState(state sharedTiDBState, leasePath string) (sharedTiDBState, func() error, error) {
	adoptedState := state
	adoptedState.OwnerPID = os.Getpid()

	statePath, err := writeSharedTiDBState(adoptedState)
	if err != nil {
		return sharedTiDBState{}, nil, fmt.Errorf("write adopted shared tidb state: %w", err)
	}
	restoreOriginalState := func() {
		_, _ = writeSharedTiDBState(state)
	}

	watchdogReader, watchdogWriter, err := os.Pipe()
	if err != nil {
		restoreOriginalState()
		return sharedTiDBState{}, nil, fmt.Errorf("create adopted shared tidb watchdog pipe: %w", err)
	}
	if _, _, err := startTiUPPlaygroundWatchdogFn(watchdogReader, adoptedState.RootPID, adoptedState.SessionID, adoptedState.Tag, statePath); err != nil {
		_ = watchdogReader.Close()
		_ = watchdogWriter.Close()
		restoreOriginalState()
		return sharedTiDBState{}, nil, fmt.Errorf("start adopted shared tidb watchdog: %w", err)
	}

	var stopOnce sync.Once
	stop := func() error {
		var stopErr error
		stopOnce.Do(func() {
			_ = releaseSharedTiDBLease(leasePath)
			unused, err := sharedTiDBRuntimeUnused(adoptedState.Tag)
			if err != nil {
				stopErr = err
				return
			}
			if !unused {
				return
			}
			_ = watchdogWriter.Close()
		})
		return stopErr
	}
	return adoptedState, stop, nil
}

func openTiDBBootstrapAdminDB(dsn string) (*tidbBootstrapAdminDB, error) {
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, err
	}
	return &tidbBootstrapAdminDB{
		db:    db,
		close: db.Close,
	}, nil
}

func (db *tidbBootstrapAdminDB) Close() error {
	if db == nil || db.close == nil {
		return nil
	}
	return db.close()
}

func configuredTiUPPlaygroundVersion() string {
	if version := strings.TrimSpace(os.Getenv("AGENTTEAM_TIDB_PLAYGROUND_VERSION")); version != "" {
		return version
	}
	return defaultTiUPPlaygroundVersion
}

func newTiUPPlaygroundSpec(version string, tag string, host string, portOffset int) tiupPlaygroundSpec {
	return tiupPlaygroundSpec{
		Version:    version,
		Tag:        tag,
		Host:       host,
		PortOffset: portOffset,
		DBPort:     defaultTiDBPort + portOffset,
		PDPort:     defaultTiDBPDPort + portOffset,
		KVPort:     defaultTiDBKVPort + portOffset,
		StatusPort: defaultTiDBStatusPort + portOffset,
	}
}

func buildStartTiUPPlaygroundArgs(spec tiupPlaygroundSpec) []string {
	args := []string{
		"playground",
		"--db", "1",
		"--pd", "1",
		"--kv", "1",
		"--tiflash", "0",
		"--without-monitor",
		"--host", spec.Host,
		"--tag", spec.Tag,
		"--port-offset", fmt.Sprintf("%d", spec.PortOffset),
	}
	if spec.Version == "" {
		return args
	}
	return append([]string{"playground", spec.Version}, args[1:]...)
}

func tiupPlaygroundRuntimeEnv(env []string) []string {
	if len(env) == 0 {
		return nil
	}

	filtered := make([]string, 0, len(env))
	for _, entry := range env {
		name, _, ok := strings.Cut(entry, "=")
		if !ok {
			filtered = append(filtered, entry)
			continue
		}
		switch {
		case strings.EqualFold(name, "http_proxy"),
			strings.EqualFold(name, "https_proxy"),
			strings.EqualFold(name, "all_proxy"):
			continue
		default:
			filtered = append(filtered, entry)
		}
	}
	return filtered
}

func tiupPlaygroundComponentSpecs(version string) []string {
	if version == "" {
		return nil
	}
	return []string{
		fmt.Sprintf("pd:%s", version),
		fmt.Sprintf("tikv:%s", version),
		fmt.Sprintf("tidb:%s", version),
	}
}

func tiupComponentBinaryPath(homeDir string, component string, version string, binary string) string {
	return filepath.Join(homeDir, ".tiup", "components", component, version, binary)
}

func missingTiUPPlaygroundComponents(version string) []string {
	if version == "" {
		return nil
	}
	homeDir, err := os.UserHomeDir()
	if err != nil {
		return tiupPlaygroundComponentSpecs(version)
	}

	required := []struct {
		component string
		binary    string
	}{
		{component: "pd", binary: "pd-server"},
		{component: "tikv", binary: "tikv-server"},
		{component: "tidb", binary: "tidb-server"},
	}

	missing := make([]string, 0, len(required))
	for _, item := range required {
		if _, err := os.Stat(tiupComponentBinaryPath(homeDir, item.component, version, item.binary)); err != nil {
			missing = append(missing, fmt.Sprintf("%s:%s", item.component, version))
		}
	}
	return missing
}

func ensureTiUPPlaygroundComponentsAvailable(ctx context.Context, version string) error {
	if version == "" {
		return nil
	}
	missing := missingTiUPPlaygroundComponents(version)
	if len(missing) == 0 {
		return nil
	}

	args := append([]string{"install", "--force"}, missing...)
	output, err := runTiUPPlaygroundCommandContext(ctx, tidbTiUPInstallTimeout, args...)
	if err != nil {
		return maybeSkipTiUPPlaygroundError("tiup install playground components", output, err)
	}
	return nil
}

func startTiUPPlaygroundServer(ctx context.Context) (string, func() error, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return "", nil, err
	}

	if err := ensureCommandAvailable("tiup"); err != nil {
		return "", nil, maybeSkipTiUPPlaygroundError("tiup playground start", nil, err)
	}

	version := configuredTiUPPlaygroundVersion()
	if err := ensureTiUPPlaygroundComponentsAvailable(ctx, version); err != nil {
		return "", nil, err
	}

	portOffset, err := reserveTiUPPlaygroundPortOffset()
	if err != nil {
		return "", nil, maybeSkipTiUPPlaygroundError("tiup playground start", nil, err)
	}
	spec := newTiUPPlaygroundSpec(
		version,
		fmt.Sprintf("agentteam-test-tidb-%d", time.Now().UnixNano()),
		"127.0.0.1",
		portOffset,
	)
	adminDSN := tiupPlaygroundAdminDSN(spec.Host, spec.DBPort)

	runCtx, cancel := context.WithCancel(context.Background())
	startupDone := make(chan struct{})
	defer close(startupDone)
	go func() {
		select {
		case <-ctx.Done():
			cancel()
		case <-startupDone:
		}
	}()

	cmd := execCommandContextFn(runCtx, "tiup", buildStartTiUPPlaygroundArgs(spec)...)
	var output bytes.Buffer
	cmd.Stdout = &output
	cmd.Stderr = &output
	cmd.Env = tiupPlaygroundRuntimeEnv(os.Environ())
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	if err := cmd.Start(); err != nil {
		cancel()
		return "", nil, maybeSkipTiUPPlaygroundError("tiup playground start", output.Bytes(), err)
	}

	waitCh := make(chan error, 1)
	go func() {
		waitCh <- cmd.Wait()
	}()

	deadline := time.Now().Add(tidbTiUPStartTimeout)
	if ctxDeadline, ok := ctx.Deadline(); ok && ctxDeadline.Before(deadline) {
		deadline = ctxDeadline
	}
	var lastErr error
	for time.Now().Before(deadline) {
		select {
		case waitErr := <-waitCh:
			cancel()
			if cmd.Process != nil {
				_ = stopTiUPPlaygroundProcessTree(cmd.Process.Pid, cmd.Process.Pid, tidbTiUPStopTimeout)
			}
			_ = cleanTiUPPlaygroundFn(spec.Tag)
			return "", nil, processExitedBeforeReadyError("tiup playground start", output.Bytes(), waitErr)
		case <-ctx.Done():
			cancel()
			if cmd.Process != nil {
				_ = stopTiUPPlaygroundProcessTree(cmd.Process.Pid, cmd.Process.Pid, tidbTiUPStopTimeout)
			}
			_ = cleanTiUPPlaygroundFn(spec.Tag)
			return "", nil, ctx.Err()
		default:
		}

		pingCtx, pingCancel := context.WithTimeout(ctx, 2*time.Second)
		db, openErr := sql.Open("mysql", adminDSN)
		if openErr == nil {
			lastErr = pingTiDBAdminFn(pingCtx, db)
			_ = db.Close()
		} else {
			lastErr = openErr
		}
		pingCancel()
		select {
		case waitErr := <-waitCh:
			cancel()
			if cmd.Process != nil {
				_ = stopTiUPPlaygroundProcessTree(cmd.Process.Pid, cmd.Process.Pid, tidbTiUPStopTimeout)
			}
			_ = cleanTiUPPlaygroundFn(spec.Tag)
			return "", nil, processExitedBeforeReadyError("tiup playground start", output.Bytes(), waitErr)
		default:
		}
		if lastErr == nil {
			runtimeState := sharedTiDBState{
				Tag:       spec.Tag,
				OwnerPID:  os.Getpid(),
				RootPID:   cmd.Process.Pid,
				SessionID: cmd.Process.Pid,
				AdminDSN:  adminDSN,
			}
			statePath, err := writeSharedTiDBState(runtimeState)
			if err != nil {
				cancel()
				_ = stopTiUPPlaygroundStateFn(runtimeState, tidbTiUPStopTimeout)
				_ = cleanTiUPPlaygroundFn(spec.Tag)
				return "", nil, fmt.Errorf("write shared tidb state: %w", err)
			}

			leasePath, err := acquireSharedTiDBLease(runtimeState.Tag)
			if err != nil {
				_ = os.Remove(statePath)
				cancel()
				_ = stopTiUPPlaygroundStateFn(runtimeState, tidbTiUPStopTimeout)
				_ = cleanTiUPPlaygroundFn(spec.Tag)
				return "", nil, fmt.Errorf("acquire shared tidb lease: %w", err)
			}

			watchdogReader, watchdogWriter, err := os.Pipe()
			if err != nil {
				_ = removeSharedTiDBLease(leasePath)
				_ = os.Remove(statePath)
				cancel()
				_ = stopTiUPPlaygroundStateFn(runtimeState, tidbTiUPStopTimeout)
				_ = cleanTiUPPlaygroundFn(spec.Tag)
				return "", nil, fmt.Errorf("create shared tidb watchdog pipe: %w", err)
			}

			if _, _, err := startTiUPPlaygroundWatchdogFn(watchdogReader, runtimeState.RootPID, runtimeState.SessionID, spec.Tag, statePath); err != nil {
				_ = watchdogReader.Close()
				_ = watchdogWriter.Close()
				_ = removeSharedTiDBLease(leasePath)
				_ = os.Remove(statePath)
				cancel()
				_ = stopTiUPPlaygroundStateFn(runtimeState, tidbTiUPStopTimeout)
				_ = cleanTiUPPlaygroundFn(spec.Tag)
				return "", nil, fmt.Errorf("start shared tidb watchdog: %w", err)
			}

			var stopOnce sync.Once
			stop := func() error {
				var stopErr error
				stopOnce.Do(func() {
					_ = releaseSharedTiDBLease(leasePath)
					unused, err := sharedTiDBRuntimeUnused(runtimeState.Tag)
					if err != nil {
						stopErr = err
						return
					}
					if !unused {
						return
					}
					_ = watchdogWriter.Close()
					cancel()
				})
				return stopErr
			}
			return adminDSN, stop, nil
		}
		if err := sleepContext(ctx, 250*time.Millisecond); err != nil {
			cancel()
			if cmd.Process != nil {
				_ = stopTiUPPlaygroundProcessTree(cmd.Process.Pid, cmd.Process.Pid, tidbTiUPStopTimeout)
			}
			_ = cleanTiUPPlaygroundFn(spec.Tag)
			return "", nil, err
		}
	}

	select {
	case waitErr := <-waitCh:
		cancel()
		if cmd.Process != nil {
			_ = stopTiUPPlaygroundProcessTree(cmd.Process.Pid, cmd.Process.Pid, tidbTiUPStopTimeout)
		}
		_ = cleanTiUPPlaygroundFn(spec.Tag)
		return "", nil, processExitedBeforeReadyError("tiup playground start", output.Bytes(), waitErr)
	default:
	}

	cancel()
	if cmd.Process != nil {
		_ = stopTiUPPlaygroundProcessTree(cmd.Process.Pid, cmd.Process.Pid, tidbTiUPStopTimeout)
	}
	waitErr := <-waitCh
	_ = cleanTiUPPlaygroundFn(spec.Tag)
	// A zero exit may have raced the final non-blocking waitCh check.
	if waitErr == nil {
		return "", nil, processExitedBeforeReadyError("tiup playground start", output.Bytes(), nil)
	}
	return "", nil, maybeSkipTiUPPlaygroundError("tiup playground start", output.Bytes(), fmt.Errorf("tiup playground did not become ready: %w", lastErr))
}

func sharedTiDBStateDir() string {
	return filepath.Join(os.TempDir(), "agentteam-shared-tidb-state")
}

// Serialize tiup start/clean operations across package test processes because
// tiup mutates shared component directories under ~/.tiup while preparing a
// playground runtime.
func withSharedTiUPPlaygroundLock(timeout time.Duration, fn func() error) (err error) {
	return withSharedTiUPPlaygroundLockContext(context.Background(), timeout, fn)
}

func withSharedTiUPPlaygroundLockContext(ctx context.Context, timeout time.Duration, fn func() error) (err error) {
	return withSharedTestkitFileLockContext(ctx, sharedTiUPPlaygroundLockName, timeout, fn)
}

func withSharedTiDBTemplateLockContext(ctx context.Context, timeout time.Duration, fn func() error) (err error) {
	return withSharedTestkitFileLockContext(ctx, sharedTiDBTemplateLockName, timeout, fn)
}

func withSharedTiDBCloneSlotContext(ctx context.Context, timeout time.Duration, fn func() error) (err error) {
	if fn == nil {
		return nil
	}

	lockFile, err := acquireSharedTestkitSlotLockContext(
		ctx,
		sharedTiDBCloneLockPrefix,
		configuredTiDBCloneConcurrency(),
		timeout,
	)
	if err != nil {
		return err
	}
	defer func() {
		if closeErr := lockFile.Close(); err == nil && closeErr != nil {
			err = fmt.Errorf("close shared tidb clone lock file: %w", closeErr)
		}
	}()
	defer func() {
		if unlockErr := releaseSharedTestkitFileLock(lockFile); err == nil && unlockErr != nil {
			err = unlockErr
		}
	}()

	return fn()
}

func withSharedTiDBAllCloneSlotsContext(ctx context.Context, timeout time.Duration, fn func() error) (err error) {
	if fn == nil {
		return nil
	}

	lockFiles, err := acquireSharedTestkitAllSlotLocksContext(
		ctx,
		sharedTiDBCloneLockPrefix,
		configuredTiDBCloneConcurrency(),
		timeout,
	)
	if err != nil {
		return err
	}
	defer func() {
		for _, lockFile := range lockFiles {
			if closeErr := lockFile.Close(); err == nil && closeErr != nil {
				err = fmt.Errorf("close shared tidb clone lock file: %w", closeErr)
			}
		}
	}()
	defer func() {
		for _, lockFile := range lockFiles {
			if unlockErr := releaseSharedTestkitFileLock(lockFile); err == nil && unlockErr != nil {
				err = unlockErr
			}
		}
	}()

	return fn()
}

func withSharedTestkitFileLock(lockName string, timeout time.Duration, fn func() error) (err error) {
	return withSharedTestkitFileLockContext(context.Background(), lockName, timeout, fn)
}

func withSharedTestkitFileLockContext(ctx context.Context, lockName string, timeout time.Duration, fn func() error) (err error) {
	if fn == nil {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return err
	}

	lockFile, err := openSharedTestkitLockFile(lockName)
	if err != nil {
		return err
	}
	defer func() {
		if closeErr := lockFile.Close(); err == nil && closeErr != nil {
			err = fmt.Errorf("close shared tiup playground lock file: %w", closeErr)
		}
	}()

	if err := acquireSharedTestkitFileLockContext(ctx, lockFile, timeout); err != nil {
		return err
	}
	defer func() {
		if unlockErr := releaseSharedTestkitFileLock(lockFile); err == nil && unlockErr != nil {
			err = unlockErr
		}
	}()

	return fn()
}

func acquireSharedTestkitSlotLock(prefix string, slots int, timeout time.Duration) (*os.File, error) {
	return acquireSharedTestkitSlotLockContext(context.Background(), prefix, slots, timeout)
}

func acquireSharedTestkitAllSlotLocksContext(ctx context.Context, prefix string, slots int, timeout time.Duration) (_ []*os.File, err error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if slots <= 0 {
		slots = defaultTiDBCloneConcurrency
	}
	if timeout <= 0 {
		timeout = tidbTestDatabaseSetupTimeout
	}

	deadline := time.Now().Add(timeout)
	queueLockFile, err := openSharedTestkitLockFile(sharedTiDBCloneQueueLockName)
	if err != nil {
		return nil, err
	}
	defer queueLockFile.Close()
	if err := acquireSharedTestkitFileLockUntilContext(ctx, queueLockFile, deadline); err != nil {
		return nil, fmt.Errorf("acquire shared tidb testkit slot queue lock: %w", err)
	}
	queueLocked := true
	defer func() {
		if queueLocked {
			if unlockErr := releaseSharedTestkitFileLock(queueLockFile); err == nil && unlockErr != nil {
				err = unlockErr
			}
		}
	}()

	lockFilesBySlot := make([]*os.File, slots)
	defer func() {
		if err == nil {
			return
		}
		for _, lockFile := range lockFilesBySlot {
			if lockFile == nil {
				continue
			}
			_ = releaseSharedTestkitFileLock(lockFile)
			_ = lockFile.Close()
		}
	}()

	acquired := 0
	for acquired < slots {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		for slot := 0; slot < slots; slot++ {
			if lockFilesBySlot[slot] != nil {
				continue
			}
			lockFile, err := openSharedTestkitLockFile(fmt.Sprintf("%s-%d.lock", prefix, slot))
			if err != nil {
				return nil, err
			}
			err = syscall.Flock(int(lockFile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
			if err == nil {
				lockFilesBySlot[slot] = lockFile
				acquired++
				continue
			}
			_ = lockFile.Close()
			if !errors.Is(err, syscall.EWOULDBLOCK) && !errors.Is(err, syscall.EAGAIN) {
				return nil, fmt.Errorf("acquire shared tidb testkit slot lock: %w", err)
			}
		}
		if acquired == slots {
			break
		}
		if time.Now().After(deadline) {
			return nil, fmt.Errorf("acquire all shared tidb testkit slot locks: timed out after %s", timeout)
		}
		if err := sleepContext(ctx, defaultTiUPLockPollInterval); err != nil {
			return nil, err
		}
	}

	if unlockErr := releaseSharedTestkitFileLock(queueLockFile); unlockErr != nil {
		return nil, unlockErr
	}
	queueLocked = false
	return lockFilesBySlot, nil
}

func acquireSharedTestkitSlotLockContext(ctx context.Context, prefix string, slots int, timeout time.Duration) (*os.File, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if slots <= 0 {
		slots = defaultTiDBCloneConcurrency
	}
	if timeout <= 0 {
		timeout = tidbTestDatabaseSetupTimeout
	}

	deadline := time.Now().Add(timeout)
	queueLockFile, err := openSharedTestkitLockFile(sharedTiDBCloneQueueLockName)
	if err != nil {
		return nil, err
	}
	defer queueLockFile.Close()
	if err := acquireSharedTestkitFileLockUntilContext(ctx, queueLockFile, deadline); err != nil {
		return nil, fmt.Errorf("acquire shared tidb testkit slot queue lock: %w", err)
	}

	for {
		if err := ctx.Err(); err != nil {
			_ = releaseSharedTestkitFileLock(queueLockFile)
			return nil, err
		}
		for slot := 0; slot < slots; slot++ {
			lockFile, err := openSharedTestkitLockFile(fmt.Sprintf("%s-%d.lock", prefix, slot))
			if err != nil {
				_ = releaseSharedTestkitFileLock(queueLockFile)
				return nil, err
			}
			err = syscall.Flock(int(lockFile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
			if err == nil {
				if unlockErr := releaseSharedTestkitFileLock(queueLockFile); unlockErr != nil {
					_ = lockFile.Close()
					return nil, unlockErr
				}
				return lockFile, nil
			}
			_ = lockFile.Close()
			if !errors.Is(err, syscall.EWOULDBLOCK) && !errors.Is(err, syscall.EAGAIN) {
				_ = releaseSharedTestkitFileLock(queueLockFile)
				return nil, fmt.Errorf("acquire shared tidb testkit slot lock: %w", err)
			}
		}
		if unlockErr := releaseSharedTestkitFileLock(queueLockFile); unlockErr != nil {
			return nil, unlockErr
		}
		if time.Now().After(deadline) {
			return nil, fmt.Errorf("acquire shared tidb testkit slot lock: timed out after %s", timeout)
		}
		if err := sleepContext(ctx, defaultTiUPLockPollInterval); err != nil {
			return nil, err
		}
		if err := acquireSharedTestkitFileLockUntilContext(ctx, queueLockFile, deadline); err != nil {
			return nil, fmt.Errorf("acquire shared tidb testkit slot queue lock: %w", err)
		}
	}
}

func configuredTiDBCloneConcurrency() int {
	value := strings.TrimSpace(os.Getenv("AGENTTEAM_TIDB_CLONE_CONCURRENCY"))
	if value == "" {
		return defaultTiDBCloneConcurrency
	}
	parsed, err := strconv.Atoi(value)
	if err != nil || parsed <= 0 {
		return defaultTiDBCloneConcurrency
	}
	return parsed
}

func openSharedTestkitLockFile(lockName string) (*os.File, error) {
	stateDir := sharedTiDBStateDirFn()
	if err := os.MkdirAll(stateDir, 0o755); err != nil {
		return nil, fmt.Errorf("create shared tidb testkit state dir: %w", err)
	}

	lockPath := filepath.Join(stateDir, lockName)
	lockFile, err := os.OpenFile(lockPath, os.O_CREATE|os.O_RDWR, 0o644)
	if err != nil {
		return nil, fmt.Errorf("open shared tidb testkit lock file: %w", err)
	}
	return lockFile, nil
}

func acquireSharedTestkitFileLock(lockFile *os.File, timeout time.Duration) error {
	return acquireSharedTestkitFileLockContext(context.Background(), lockFile, timeout)
}

func acquireSharedTestkitFileLockContext(ctx context.Context, lockFile *os.File, timeout time.Duration) error {
	if timeout <= 0 {
		timeout = tidbTiUPStartTimeout
	}
	return acquireSharedTestkitFileLockUntilContext(ctx, lockFile, time.Now().Add(timeout))
}

func acquireSharedTestkitFileLockUntil(lockFile *os.File, deadline time.Time) error {
	return acquireSharedTestkitFileLockUntilContext(context.Background(), lockFile, deadline)
}

func acquireSharedTestkitFileLockUntilContext(ctx context.Context, lockFile *os.File, deadline time.Time) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if lockFile == nil {
		return fmt.Errorf("shared tidb testkit lock file is nil")
	}

	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		err := syscall.Flock(int(lockFile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
		if err == nil {
			return nil
		}
		if !errors.Is(err, syscall.EWOULDBLOCK) && !errors.Is(err, syscall.EAGAIN) {
			return fmt.Errorf("acquire shared tidb testkit lock: %w", err)
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("acquire shared tidb testkit lock: timed out")
		}
		if err := sleepContext(ctx, defaultTiUPLockPollInterval); err != nil {
			return err
		}
	}
}

func releaseSharedTestkitFileLock(lockFile *os.File) error {
	if lockFile == nil {
		return nil
	}
	if err := syscall.Flock(int(lockFile.Fd()), syscall.LOCK_UN); err != nil {
		return fmt.Errorf("release shared tidb testkit lock: %w", err)
	}
	return nil
}

func writeSharedTiDBState(state sharedTiDBState) (string, error) {
	stateDir := sharedTiDBStateDirFn()
	if err := os.MkdirAll(stateDir, 0o755); err != nil {
		return "", fmt.Errorf("create shared tidb state dir: %w", err)
	}

	if strings.TrimSpace(state.Tag) == "" {
		return "", fmt.Errorf("shared tidb state tag is required")
	}

	payload, err := json.Marshal(state)
	if err != nil {
		return "", fmt.Errorf("marshal shared tidb state: %w", err)
	}

	statePath := filepath.Join(stateDir, strings.TrimSpace(state.Tag)+".json")
	tmpFile, err := os.CreateTemp(stateDir, "agentteam-shared-tidb-state-*.tmp")
	if err != nil {
		return "", fmt.Errorf("create shared tidb state temp file: %w", err)
	}
	tmpPath := tmpFile.Name()
	defer func() { _ = os.Remove(tmpPath) }()

	if _, err := tmpFile.Write(payload); err != nil {
		_ = tmpFile.Close()
		return "", fmt.Errorf("write shared tidb state: %w", err)
	}
	if err := tmpFile.Close(); err != nil {
		return "", fmt.Errorf("close shared tidb state temp file: %w", err)
	}
	if err := os.Rename(tmpPath, statePath); err != nil {
		return "", fmt.Errorf("move shared tidb state into place: %w", err)
	}
	return statePath, nil
}

func acquireSharedTiDBLease(tag string) (string, error) {
	leasePath, err := sharedTiDBLeasePath(tag, os.Getpid())
	if err != nil {
		return "", err
	}
	if err := os.WriteFile(leasePath, []byte(strconv.Itoa(os.Getpid())), 0o644); err != nil {
		return "", fmt.Errorf("write shared tidb lease: %w", err)
	}
	return leasePath, nil
}

func releaseSharedTiDBLease(path string) error {
	// A lease marks that this test process may still start more parallel TiDB
	// tests. Keep it until the process exits; liveness checks remove stale
	// lease files by PID.
	return nil
}

func removeSharedTiDBLease(path string) error {
	if strings.TrimSpace(path) == "" {
		return nil
	}
	if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) {
		return fmt.Errorf("remove shared tidb lease: %w", err)
	}
	return nil
}

func sharedTiDBLeasePath(tag string, pid int) (string, error) {
	trimmedTag := strings.TrimSpace(tag)
	if trimmedTag == "" {
		return "", fmt.Errorf("shared tidb lease tag is required")
	}
	if strings.ContainsAny(trimmedTag, `/\`) {
		return "", fmt.Errorf("shared tidb lease tag %q must not contain path separators", trimmedTag)
	}
	if pid <= 0 {
		return "", fmt.Errorf("shared tidb lease pid must be positive")
	}
	stateDir := sharedTiDBStateDirFn()
	if err := os.MkdirAll(stateDir, 0o755); err != nil {
		return "", fmt.Errorf("create shared tidb lease dir: %w", err)
	}
	return filepath.Join(stateDir, fmt.Sprintf("%s.lease.%d", trimmedTag, pid)), nil
}

func sharedTiDBRuntimeUnused(tag string) (bool, error) {
	pids, err := liveSharedTiDBLeasePIDs(tag)
	if err != nil {
		return false, err
	}
	return len(pids) == 0, nil
}

func liveSharedTiDBLeasePIDs(tag string) ([]int, error) {
	trimmedTag := strings.TrimSpace(tag)
	if trimmedTag == "" {
		return nil, nil
	}
	pattern := filepath.Join(sharedTiDBStateDirFn(), trimmedTag+".lease.*")
	paths, err := filepath.Glob(pattern)
	if err != nil {
		return nil, fmt.Errorf("list shared tidb leases: %w", err)
	}

	pids := make([]int, 0, len(paths))
	prefix := trimmedTag + ".lease."
	for _, path := range paths {
		pidText := strings.TrimPrefix(filepath.Base(path), prefix)
		pid, err := strconv.Atoi(pidText)
		if err != nil || pid <= 0 {
			_ = os.Remove(path)
			continue
		}
		if processExistsFn(pid) {
			pids = append(pids, pid)
			continue
		}
		_ = os.Remove(path)
	}
	sort.Ints(pids)
	return pids, nil
}

func removeSharedTiDBLeaseFiles(tag string) error {
	trimmedTag := strings.TrimSpace(tag)
	if trimmedTag == "" {
		return nil
	}
	pattern := filepath.Join(sharedTiDBStateDirFn(), trimmedTag+".lease.*")
	paths, err := filepath.Glob(pattern)
	if err != nil {
		return fmt.Errorf("list shared tidb leases: %w", err)
	}
	var firstErr error
	for _, path := range paths {
		if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) && firstErr == nil {
			firstErr = fmt.Errorf("remove shared tidb lease %s: %w", path, err)
		}
	}
	return firstErr
}

func reusableSharedTiDBState() (sharedTiDBState, bool, error) {
	states, err := readSharedTiDBStates()
	if err != nil {
		return sharedTiDBState{}, false, err
	}
	for _, state := range states {
		if strings.TrimSpace(state.AdminDSN) == "" || strings.TrimSpace(state.Tag) == "" {
			continue
		}

		if !pingSharedTiDBState(state, defaultTiDBReusePingTimeout) {
			continue
		}
		if processExistsFn(state.OwnerPID) {
			return state, true, nil
		}
		liveLeasePIDs, err := liveSharedTiDBLeasePIDs(state.Tag)
		if err != nil {
			return sharedTiDBState{}, false, err
		}
		if len(liveLeasePIDs) > 0 {
			return state, true, nil
		}
	}
	return sharedTiDBState{}, false, nil
}

func pingSharedTiDBState(state sharedTiDBState, timeout time.Duration) bool {
	return pingSharedTiDBAdminDSN(state.AdminDSN, timeout)
}

func pingSharedTiDBAdminDSN(adminDSN string, timeout time.Duration) bool {
	if strings.TrimSpace(adminDSN) == "" {
		return false
	}
	if timeout <= 0 {
		timeout = defaultTiDBReusePingTimeout
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	adminDB, err := openTiDBBootstrapAdminDBFn(adminDSN)
	if err != nil {
		return false
	}
	defer func() { _ = adminDB.Close() }()

	if err := pingTiDBAdminFn(ctx, adminDB.db); err != nil {
		return false
	}
	return true
}

func readSharedTiDBStates() ([]sharedTiDBState, error) {
	stateDir := sharedTiDBStateDirFn()
	entries, err := os.ReadDir(stateDir)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("read shared tidb state dir: %w", err)
	}

	states := make([]sharedTiDBState, 0)
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(strings.ToLower(entry.Name()), ".json") {
			continue
		}
		statePath := filepath.Join(stateDir, entry.Name())
		stateData, err := os.ReadFile(statePath)
		if err != nil {
			return nil, fmt.Errorf("read shared tidb state %s: %w", statePath, err)
		}

		var state sharedTiDBState
		if err := json.Unmarshal(stateData, &state); err != nil {
			_ = os.Remove(statePath)
			return nil, fmt.Errorf("decode shared tidb state %s: %w", statePath, err)
		}
		states = append(states, state)
	}
	return states, nil
}

func cleanupOrphanedSharedTiDBPlaygrounds() error {
	stateDir := sharedTiDBStateDirFn()
	entries, err := os.ReadDir(stateDir)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("read shared tidb state dir: %w", err)
	}

	var firstErr error
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(strings.ToLower(entry.Name()), ".json") {
			continue
		}

		statePath := filepath.Join(stateDir, entry.Name())
		stateData, err := os.ReadFile(statePath)
		if err != nil {
			if firstErr == nil {
				firstErr = fmt.Errorf("read shared tidb state %s: %w", statePath, err)
			}
			continue
		}

		var state sharedTiDBState
		if err := json.Unmarshal(stateData, &state); err != nil {
			_ = os.Remove(statePath)
			if firstErr == nil {
				firstErr = fmt.Errorf("decode shared tidb state %s: %w", statePath, err)
			}
			continue
		}
		if processExistsFn(state.OwnerPID) {
			continue
		}
		unused, err := sharedTiDBRuntimeUnused(state.Tag)
		if err != nil {
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		if !unused {
			continue
		}

		stopErr := stopTiUPPlaygroundStateFn(state, tidbTiUPStopTimeout)
		cleanErr := cleanTiUPPlaygroundFn(state.Tag)
		if stopErr == nil && cleanErr == nil {
			_ = os.Remove(statePath)
			_ = removeSharedTiDBLeaseFiles(state.Tag)
			continue
		}
		if firstErr == nil {
			if stopErr != nil && cleanErr != nil {
				firstErr = fmt.Errorf("stop orphaned shared tidb playground %s: %v; clean playground data: %w", state.Tag, stopErr, cleanErr)
			} else if stopErr != nil {
				firstErr = stopErr
			} else {
				firstErr = cleanErr
			}
		}
	}
	return firstErr
}

func stopTiUPPlaygroundState(state sharedTiDBState, timeout time.Duration) error {
	if state.SessionID > 0 || state.RootPID > 0 {
		return stopTiUPPlaygroundProcessTree(state.RootPID, state.SessionID, timeout)
	}
	return stopTiUPPlaygroundProcessGroupFn(state.ProcessGroupID, timeout)
}

func stopTiUPPlaygroundProcessTree(rootPID int, sessionID int, timeout time.Duration) error {
	if rootPID <= 0 && sessionID <= 0 {
		return nil
	}
	if timeout <= 0 {
		timeout = tidbTiUPStopTimeout
	}

	if err := signalTiUPPlaygroundProcesses(rootPID, sessionID, syscall.SIGTERM); err != nil {
		return err
	}
	if waitForTiUPPlaygroundProcessesExit(rootPID, sessionID, timeout) {
		return nil
	}
	if err := signalTiUPPlaygroundProcesses(rootPID, sessionID, syscall.SIGKILL); err != nil {
		return err
	}
	if waitForTiUPPlaygroundProcessesExit(rootPID, sessionID, timeout) {
		return nil
	}

	remaining, err := tiupPlaygroundProcessIDs(rootPID, sessionID)
	if err != nil {
		return err
	}
	return fmt.Errorf("tiup playground root pid %d session %d did not exit; remaining pids: %v", rootPID, sessionID, remaining)
}

func signalTiUPPlaygroundProcesses(rootPID int, sessionID int, signal syscall.Signal) error {
	var firstErr error
	if sessionID > 0 {
		if err := syscall.Kill(-sessionID, signal); err != nil && !errors.Is(err, syscall.ESRCH) {
			firstErr = fmt.Errorf("signal tiup playground session %d: %w", sessionID, err)
		}
	}

	pids, err := tiupPlaygroundProcessIDs(rootPID, sessionID)
	if len(pids) == 0 {
		if firstErr != nil {
			return firstErr
		}
		return err
	}

	for _, pid := range pids {
		process, findErr := os.FindProcess(pid)
		if findErr != nil {
			continue
		}
		if signalErr := process.Signal(signal); signalErr != nil && processExistsFn(pid) && firstErr == nil {
			firstErr = fmt.Errorf("signal tiup playground pid %d: %w", pid, signalErr)
		}
	}
	if firstErr != nil {
		return firstErr
	}
	return err
}

func waitForTiUPPlaygroundProcessesExit(rootPID int, sessionID int, timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if !tiupPlaygroundProcessesExist(rootPID, sessionID) {
			return true
		}
		time.Sleep(100 * time.Millisecond)
	}
	return !tiupPlaygroundProcessesExist(rootPID, sessionID)
}

func tiupPlaygroundProcessesExist(rootPID int, sessionID int) bool {
	if rootPID > 0 && processExistsFn(rootPID) {
		return true
	}

	pids, err := listTiUPPlaygroundSessionPIDsFn(sessionID)
	return err == nil && len(pids) > 0
}

func tiupPlaygroundProcessIDs(rootPID int, sessionID int) ([]int, error) {
	pids := []int{}
	var firstErr error
	if rootPID > 0 {
		pids = append(pids, rootPID)
		descendantPIDs, err := listTiUPPlaygroundDescendantPIDs(rootPID)
		pids = append(pids, descendantPIDs...)
		if err != nil {
			firstErr = err
		}
	}

	sessionPIDs, err := listTiUPPlaygroundSessionPIDsFn(sessionID)
	pids = append(pids, sessionPIDs...)
	if err != nil && rootPID <= 0 {
		firstErr = err
	}
	pids = uniqueTiUPPlaygroundPIDs(pids)
	if firstErr != nil {
		return pids, firstErr
	}
	return pids, nil
}

func listTiUPPlaygroundDescendantPIDs(rootPID int) ([]int, error) {
	if rootPID <= 0 {
		return nil, nil
	}

	output, err := exec.Command("ps", "-A", "-o", "pid=", "-o", "ppid=").CombinedOutput()
	trimmed := strings.TrimSpace(string(output))
	if err != nil {
		return nil, fmt.Errorf("list tiup playground descendants for root pid %d: %s: %w", rootPID, trimmed, err)
	}
	return descendantPIDsFromProcessTable(rootPID, string(output))
}

func descendantPIDsFromProcessTable(rootPID int, output string) ([]int, error) {
	if rootPID <= 0 {
		return nil, nil
	}

	childrenByParent := map[int][]int{}
	for _, line := range strings.Split(output, "\n") {
		fields := strings.Fields(line)
		if len(fields) == 0 {
			continue
		}
		if len(fields) < 2 {
			if isProcessTableHeaderLine(line) {
				continue
			}
			return nil, fmt.Errorf("parse process table line %q: expected pid and ppid", line)
		}
		pid, err := strconv.Atoi(strings.TrimSpace(fields[0]))
		if err != nil {
			if isProcessTableHeaderLine(line) {
				continue
			}
			return nil, fmt.Errorf("parse process table pid %q: %w", fields[0], err)
		}
		ppid, err := strconv.Atoi(strings.TrimSpace(fields[1]))
		if err != nil {
			if isProcessTableHeaderLine(line) {
				continue
			}
			return nil, fmt.Errorf("parse process table parent pid %q: %w", fields[1], err)
		}
		if pid <= 0 || ppid <= 0 {
			continue
		}
		childrenByParent[ppid] = append(childrenByParent[ppid], pid)
	}

	pids := []int{}
	queue := append([]int(nil), childrenByParent[rootPID]...)
	seen := map[int]bool{rootPID: true}
	for len(queue) > 0 {
		pid := queue[0]
		queue = queue[1:]
		if seen[pid] {
			continue
		}
		seen[pid] = true
		pids = append(pids, pid)
		queue = append(queue, childrenByParent[pid]...)
	}
	return uniqueTiUPPlaygroundPIDs(pids), nil
}

func isProcessTableHeaderLine(line string) bool {
	normalized := strings.ToLower(strings.TrimSpace(line))
	return strings.Contains(normalized, "pid")
}

func listTiUPPlaygroundSessionPIDs(sessionID int) ([]int, error) {
	if sessionID <= 0 {
		return nil, nil
	}

	output, err := exec.Command("ps", "-A", "-o", "pid=", "-o", "pgid=").CombinedOutput()
	trimmed := strings.TrimSpace(string(output))
	if err != nil {
		return nil, fmt.Errorf("list tiup playground process group %d: %s: %w", sessionID, trimmed, err)
	}
	if trimmed == "" {
		return nil, nil
	}

	pids := []int{}
	for _, line := range strings.Split(trimmed, "\n") {
		fields := strings.Fields(line)
		if len(fields) < 2 {
			if isProcessTableHeaderLine(line) {
				continue
			}
			return nil, fmt.Errorf("parse tiup playground process group line %q: expected pid and pgid", line)
		}
		pid, parseErr := strconv.Atoi(fields[0])
		if parseErr != nil {
			if isProcessTableHeaderLine(line) {
				continue
			}
			return nil, fmt.Errorf("parse tiup playground process group pid %q: %w", fields[0], parseErr)
		}
		processGroupID, parseErr := strconv.Atoi(fields[1])
		if parseErr != nil {
			if isProcessTableHeaderLine(line) {
				continue
			}
			return nil, fmt.Errorf("parse tiup playground process group id %q: %w", fields[1], parseErr)
		}
		if processGroupID == sessionID && pid > 0 && processExistsFn(pid) {
			pids = append(pids, pid)
		}
	}
	return uniqueTiUPPlaygroundPIDs(pids), nil
}

func uniqueTiUPPlaygroundPIDs(pids []int) []int {
	if len(pids) == 0 {
		return nil
	}

	sort.Ints(pids)
	unique := make([]int, 0, len(pids))
	last := 0
	for _, pid := range pids {
		if pid <= 0 {
			continue
		}
		if len(unique) > 0 && pid == last {
			continue
		}
		unique = append(unique, pid)
		last = pid
	}
	return unique
}

func stopTiUPPlaygroundProcessGroup(processGroupID int, timeout time.Duration) error {
	if processGroupID <= 0 {
		return nil
	}
	if timeout <= 0 {
		timeout = tidbTiUPStopTimeout
	}

	if err := syscall.Kill(-processGroupID, syscall.SIGTERM); err != nil && !errors.Is(err, syscall.ESRCH) {
		return fmt.Errorf("signal tiup playground process group %d: %w", processGroupID, err)
	}
	if waitForProcessGroupExit(processGroupID, timeout) {
		return nil
	}
	if err := syscall.Kill(-processGroupID, syscall.SIGKILL); err != nil && !errors.Is(err, syscall.ESRCH) {
		return fmt.Errorf("kill tiup playground process group %d: %w", processGroupID, err)
	}
	if waitForProcessGroupExit(processGroupID, timeout) {
		return nil
	}
	return fmt.Errorf("tiup playground process group %d did not exit", processGroupID)
}

func waitForProcessGroupExit(processGroupID int, timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if !processGroupExists(processGroupID) {
			return true
		}
		time.Sleep(100 * time.Millisecond)
	}
	return !processGroupExists(processGroupID)
}

func processGroupExists(processGroupID int) bool {
	if processGroupID <= 0 {
		return false
	}
	err := syscall.Kill(-processGroupID, syscall.Signal(0))
	if err == nil {
		return true
	}
	return !errors.Is(err, syscall.ESRCH)
}

func startTiUPPlaygroundWatchdog(reader *os.File, rootPID int, sessionID int, tag string, statePath string) (*os.Process, <-chan error, error) {
	if reader == nil {
		return nil, nil, fmt.Errorf("shared tidb watchdog reader is required")
	}
	processGroupID := sessionID
	if processGroupID <= 0 {
		processGroupID = rootPID
	}

	watchdog := exec.Command("sh", "-c", `
cat >/dev/null
if [ ! -f "$AGENTTEAM_TIDB_STATE_PATH" ]; then
  exit 0
fi

idle_deadline=$(( $(date +%s) + ${AGENTTEAM_TIDB_IDLE_WAIT_SECONDS:-600} ))
while :; do
  if [ ! -f "$AGENTTEAM_TIDB_STATE_PATH" ]; then
    exit 0
  fi

  live_lease=""
  for lease_path in "$AGENTTEAM_TIDB_STATE_DIR"/"$AGENTTEAM_TIDB_TAG".lease.*; do
    [ -e "$lease_path" ] || continue
    lease_pid="${lease_path##*.lease.}"
    case "$lease_pid" in
      ""|*[!0-9]*)
        rm -f "$lease_path"
        continue
        ;;
    esac
    if kill -0 "$lease_pid" >/dev/null 2>&1; then
      live_lease=1
      break
    fi
    rm -f "$lease_path"
  done

  if [ -z "$live_lease" ]; then
    if [ "$(date +%s)" -ge "$idle_deadline" ]; then
      break
    fi
  else
    idle_deadline=$(( $(date +%s) + ${AGENTTEAM_TIDB_IDLE_WAIT_SECONDS:-600} ))
  fi
  sleep 1
done

process_group_owns_tag() {
  [ "$AGENTTEAM_TIDB_PROCESS_GROUP_ID" -gt 0 ] || return 2
  [ -n "$AGENTTEAM_TIDB_TAG" ] || return 2
  process_table="$(ps -A -o pgid= -o command= 2>/dev/null)" || return 2
  printf '%s\n' "$process_table" |
    awk -v pgid="$AGENTTEAM_TIDB_PROCESS_GROUP_ID" -v tag="$AGENTTEAM_TIDB_TAG" '
      NF < 2 || $1 !~ /^[0-9]+$/ { invalid = 1; next }
      { rows++ }
      $1 == pgid && index($0, tag) > 0 { found = 1 }
      END {
        if (rows == 0 || invalid) exit 2
        exit found ? 0 : 1
      }
    '
}

for signal_name in TERM KILL; do
  process_group_owns_tag
  ownership_status=$?
  case "$ownership_status" in
    0)
      kill -"$signal_name" "-$AGENTTEAM_TIDB_PROCESS_GROUP_ID" >/dev/null 2>&1 || true
      sleep 1
      ;;
    1)
      break
      ;;
    *)
      exit 1
      ;;
  esac
done

process_group_owns_tag
ownership_status=$?
case "$ownership_status" in
  1)
    ;;
  *)
    exit 1
    ;;
esac

clean_output="$(tiup clean "$AGENTTEAM_TIDB_TAG" 2>&1)" || clean_status=$?
if [ "${clean_status:-0}" -ne 0 ]; then
  lower_output="$(printf '%s' "$clean_output" | tr '[:upper:]' '[:lower:]')"
  case "$lower_output" in
    *"not found"*|*"no such file"*)
      :
      ;;
    *)
      exit 1
      ;;
  esac
fi
rm -f "$AGENTTEAM_TIDB_STATE_PATH"
rm -f "$AGENTTEAM_TIDB_STATE_DIR"/"$AGENTTEAM_TIDB_TAG".lease.*
`)
	watchdog.Stdin = reader
	watchdog.Stdout = io.Discard
	watchdog.Stderr = io.Discard
	watchdog.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	watchdog.Env = append(
		os.Environ(),
		fmt.Sprintf("AGENTTEAM_TIDB_PROCESS_GROUP_ID=%d", processGroupID),
		fmt.Sprintf("AGENTTEAM_TIDB_TAG=%s", strings.TrimSpace(tag)),
		fmt.Sprintf("AGENTTEAM_TIDB_STATE_DIR=%s", filepath.Dir(strings.TrimSpace(statePath))),
		fmt.Sprintf("AGENTTEAM_TIDB_STATE_PATH=%s", strings.TrimSpace(statePath)),
	)
	if err := watchdog.Start(); err != nil {
		_ = reader.Close()
		return nil, nil, err
	}
	_ = reader.Close()
	done := make(chan error, 1)
	go func() {
		done <- watchdog.Wait()
		close(done)
	}()
	return watchdog.Process, done, nil
}

func reserveTiUPPlaygroundPortOffset() (int, error) {
	const (
		minPortOffset = 1000
		maxPortOffset = 30000
	)

	for portOffset := minPortOffset; portOffset <= maxPortOffset; portOffset++ {
		spec := newTiUPPlaygroundSpec("", "", "127.0.0.1", portOffset)
		if !tiupPlaygroundPortsValid(spec) {
			continue
		}
		if tiupPlaygroundPortsAvailable(spec) {
			return portOffset, nil
		}
	}
	return 0, fmt.Errorf("reserve tiup playground port offset: exhausted retries")
}

func tiupPlaygroundPortsValid(spec tiupPlaygroundSpec) bool {
	ports := []int{spec.DBPort, spec.PDPort, spec.KVPort, spec.StatusPort}
	for _, port := range ports {
		if port <= 0 || port > 65535 {
			return false
		}
	}
	return tiupPlaygroundPortsDistinct(ports)
}

func tiupPlaygroundPortsAvailable(spec tiupPlaygroundSpec) bool {
	for _, port := range []int{spec.DBPort, spec.PDPort, spec.KVPort, spec.StatusPort} {
		listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
		if err != nil {
			return false
		}
		_ = listener.Close()
	}
	if !tiupPlaygroundTempStorageLockAvailable(spec) {
		return false
	}
	return true
}

func tiupPlaygroundTempStorageLockAvailable(spec tiupPlaygroundSpec) bool {
	lockPath := tiupPlaygroundTempStorageLockPath(spec)
	if _, err := os.Stat(lockPath); errors.Is(err, os.ErrNotExist) {
		return true
	} else if err != nil {
		return false
	}

	lockFile, err := os.OpenFile(lockPath, os.O_RDWR, 0)
	if err != nil {
		return false
	}
	defer lockFile.Close()

	err = syscall.Flock(int(lockFile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
	if err == nil {
		_ = syscall.Flock(int(lockFile.Fd()), syscall.LOCK_UN)
		return true
	}
	if errors.Is(err, syscall.EWOULDBLOCK) || errors.Is(err, syscall.EAGAIN) {
		return false
	}
	return false
}

func tiupPlaygroundTempStorageLockPath(spec tiupPlaygroundSpec) string {
	tempStorageIdentity := fmt.Sprintf("%s:%d/0.0.0.0:%d", spec.Host, spec.DBPort, spec.StatusPort)
	encodedIdentity := base64.StdEncoding.EncodeToString([]byte(tempStorageIdentity))
	return filepath.Join(os.TempDir(), fmt.Sprintf("%d_tidb", os.Getuid()), encodedIdentity, "tmp-storage", "_dir.lock")
}

func tiupPlaygroundPortsDistinct(ports []int) bool {
	seen := make(map[int]struct{}, len(ports))
	for _, port := range ports {
		if _, ok := seen[port]; ok {
			return false
		}
		seen[port] = struct{}{}
	}
	return true
}

func cleanTiUPPlayground(tag string) error {
	if strings.TrimSpace(tag) == "" {
		return nil
	}
	output, err := runTiUPPlaygroundCommand(tidbTiUPCleanTimeout, "clean", tag)
	if err != nil {
		message := strings.ToLower(strings.TrimSpace(string(output)))
		if strings.Contains(message, "not found") || strings.Contains(message, "no such file") {
			return nil
		}
		return fmt.Errorf("tiup clean %s: %s: %w", tag, strings.TrimSpace(string(output)), err)
	}
	return nil
}

func runTiUPPlaygroundCommand(timeout time.Duration, args ...string) ([]byte, error) {
	return runTiUPPlaygroundCommandContext(context.Background(), timeout, args...)
}

func runTiUPPlaygroundCommandContext(ctx context.Context, timeout time.Duration, args ...string) ([]byte, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	cmdCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	output, err := execCommandContextFn(cmdCtx, "tiup", args...).CombinedOutput()
	if err != nil {
		if ctxErr := cmdCtx.Err(); ctxErr != nil {
			return output, fmt.Errorf("%w: %w", ctxErr, err)
		}
	}
	return output, err
}

func processExitedBeforeReadyError(action string, output []byte, err error) error {
	if err == nil {
		message := strings.TrimSpace(string(output))
		if message == "" {
			return fmt.Errorf("%s: process exited before becoming ready", action)
		}
		return fmt.Errorf("%s: %s: process exited before becoming ready", action, message)
	}
	return maybeSkipTiUPPlaygroundError(action, output, err)
}

func maybeSkipTiUPPlaygroundError(action string, output []byte, err error) error {
	if err == nil {
		return nil
	}

	if errors.Is(err, exec.ErrNotFound) {
		return &skipTestError{reason: fmt.Sprintf("tiup playground is required for self-hosted tidb integration tests (%s)", action)}
	}

	message := commandMessage(output, err)
	lower := strings.ToLower(message)
	if errors.Is(err, os.ErrPermission) ||
		strings.Contains(lower, "permission denied") ||
		strings.Contains(lower, "operation not permitted") ||
		strings.Contains(lower, "read-only file system") ||
		isTiDBTempStorageLockMessage(lower) {
		return &skipTestError{reason: fmt.Sprintf("self-hosted tidb integration tests require local socket access and writable tiup state (%s)", action)}
	}
	if errors.Is(err, context.DeadlineExceeded) ||
		strings.Contains(lower, "temporary failure in name resolution") ||
		strings.Contains(lower, "no route to host") ||
		strings.Contains(lower, "connection refused") ||
		strings.Contains(lower, "i/o timeout") ||
		strings.Contains(lower, "dial tcp") ||
		strings.Contains(lower, "lookup") ||
		strings.Contains(lower, "download") {
		return &skipTestError{reason: fmt.Sprintf("tiup playground components must be installed or downloadable for self-hosted tidb integration tests (%s)", action)}
	}
	return fmt.Errorf("%s: %s: %w", action, strings.TrimSpace(string(output)), err)
}

func isTiDBTempStorageLockMessage(lower string) bool {
	return strings.Contains(lower, "fslock: lock is held") ||
		strings.Contains(lower, "temporary storage dir has been occupied") ||
		strings.Contains(lower, "current temporary storage dir has been occupied")
}

func commandMessage(output []byte, err error) string {
	message := strings.TrimSpace(string(output))
	if message == "" && err != nil {
		message = strings.TrimSpace(err.Error())
	}
	return message
}

func sleepContext(ctx context.Context, delay time.Duration) error {
	if delay <= 0 {
		return nil
	}
	timer := time.NewTimer(delay)
	defer timer.Stop()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func tiupPlaygroundAdminDSN(host string, port int) string {
	return fmt.Sprintf("root@tcp(%s:%d)/?multiStatements=true&parseTime=true&loc=UTC&time_zone=%%27%%2B00%%3A00%%27", host, port)
}

func pingTiDBAdmin(ctx context.Context, db *sql.DB) error {
	if db == nil {
		return fmt.Errorf("tidb db is nil")
	}
	if err := db.PingContext(ctx); err != nil {
		return err
	}

	var version string
	if err := db.QueryRowContext(ctx, `SELECT VERSION()`).Scan(&version); err != nil {
		return err
	}
	if !strings.Contains(strings.ToLower(strings.TrimSpace(version)), "tidb") {
		return fmt.Errorf("connected database does not report a TiDB version: %q", version)
	}

	var txnMode string
	if err := db.QueryRowContext(ctx, `SELECT @@tidb_txn_mode`).Scan(&txnMode); err != nil {
		return err
	}
	if strings.TrimSpace(txnMode) == "" {
		return fmt.Errorf("tidb transaction mode is empty")
	}
	return nil
}

func (instance *sharedTiDBInstance) close() error {
	if instance.stop != nil {
		return withSharedTiUPPlaygroundLock(tidbTiUPStartTimeout, instance.stop)
	}
	return nil
}

func skipSharedTiDBIfNeeded(t *testing.T, err error) {
	t.Helper()

	if err == nil {
		return
	}

	var skipErr *skipTestError
	if errors.As(err, &skipErr) {
		t.Skip(skipErr.reason)
	}
}