opsxsql/backend/internal/ghost/ghost.go

package ghost

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"

	"sql-platform/internal/safe"
)

// GhOstTask defines a gh-ost execution task
type GhOstTask struct {
	Host            string `json:"host"`
	Port            int    `json:"port,omitempty"`
	User            string `json:"user"`
	Password        string `json:"password"`
	Database        string `json:"database"`
	Table           string `json:"table"`
	Alter           string `json:"alter"` // e.g. "ADD COLUMN status INT NOT NULL DEFAULT 0"
	ChunkSize       int    `json:"chunk_size,omitempty"`
	MaxLoad         string `json:"max_load,omitempty"`
	CriticalLoad    string `json:"critical_load,omitempty"`
	MaxLagMillis    int    `json:"max_lag_millis,omitempty"`
	AllowOnMaster   bool   `json:"allow_on_master"`
	DryRun          bool   `json:"dry_run"`
	Execute         bool   `json:"execute"`
	ExactRowCount   bool   `json:"exact_row_count"`
	PostponeCutOver bool   `json:"postpone_cut_over"`
	AliyunRDS       bool   `json:"aliyun_rds,omitempty"` // --aliyun-rds
	AWSAurora       bool   `json:"aws_aurora,omitempty"` // --aws
}

// GhOstResult holds the execution result
type GhOstResult struct {
	Success  bool   `json:"success"`
	Output   string `json:"output"`
	Error    string `json:"error"`
	Command  string `json:"command"`
	Duration string `json:"duration"`
	TaskID   string `json:"task_id"`
}

// ProgressCallback is called with stdout lines for real-time progress
type ProgressCallback func(taskID string, line string)

// TaskManager manages async gh-ost tasks
type TaskManager struct {
	mu     sync.RWMutex
	tasks  map[string]*GhOstResult
	status map[string]string // running, completed, failed
}

var Manager = &TaskManager{
	tasks:  make(map[string]*GhOstResult),
	status: make(map[string]string),
}

// RunGhOst executes gh-ost synchronously
func RunGhOst(ctx context.Context, task GhOstTask, ghOstPath string) (*GhOstResult, error) {
	start := time.Now()

	args := []string{
		"--host=" + task.Host,
		fmt.Sprintf("--port=%d", task.Port),
		"--user=" + task.User,
		"--password=" + task.Password,
		"--database=" + task.Database,
		"--table=" + task.Table,
		"--alter=" + task.Alter,
		"--verbose",
		"--skip-port-validation",
		"--allow-on-master",
		"--initially-drop-ghost-table",
		"--initially-drop-old-table",
		"--initially-drop-socket-file",
	}

	if task.ChunkSize > 0 {
		args = append(args, fmt.Sprintf("--chunk-size=%d", task.ChunkSize))
	}
	if task.MaxLoad != "" {
		args = append(args, "--max-load="+task.MaxLoad)
	}
	if task.CriticalLoad != "" {
		args = append(args, "--critical-load="+task.CriticalLoad)
	}
	if task.MaxLagMillis > 0 {
		args = append(args, fmt.Sprintf("--max-lag-millis=%d", task.MaxLagMillis))
	}
	if task.AllowOnMaster {
		args = append(args, "--allow-on-master")
	}
	if task.ExactRowCount {
		args = append(args, "--exact-rowcount")
	}
	if task.PostponeCutOver {
		args = append(args, "--postpone-cut-over-flag-file=/tmp/gh-ost-cutover.flag")
	}
	// Cloud-specific flags
	if task.AliyunRDS {
		args = append(args, "--aliyun-rds")
	}
	if task.AWSAurora {
		args = append(args, "--aws")
	}

	// Dry-Run vs Execute
	// gh-ost defaults to noop (dry-run), only --execute actually applies changes
	if task.Execute {
		args = append(args, "--execute")
	}
	// always verbose for output

	cmd := exec.CommandContext(ctx, ghOstPath, args...)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()

	result := &GhOstResult{
		Command:  "gh-ost " + maskSecrets(strings.Join(args, " ")),
		Duration: time.Since(start).String(),
		Output:   stdout.String(),
	}

	if err != nil {
		result.Success = false
		result.Error = stderr.String() + "\n" + err.Error()
		return result, err
	}

	result.Success = true
	return result, nil
}

// RunGhOstAsync runs gh-ost in background and tracks status
func RunGhOstAsync(task GhOstTask, ghOstPath string, cb ProgressCallback) string {
	taskID := fmt.Sprintf("ghost-%d", time.Now().UnixNano())

	Manager.mu.Lock()
	Manager.status[taskID] = "running"
	Manager.tasks[taskID] = &GhOstResult{TaskID: taskID}
	Manager.mu.Unlock()

	run := func() {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Hour)
		defer cancel()

		start := time.Now()

		args := buildArgs(task, ghOstPath)

		cmd := exec.CommandContext(ctx, ghOstPath, args...)
		var stdout, stderr bytes.Buffer
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr

		err := cmd.Run()

		// 注意：这里不能用 defer Unlock —— 回调要放到锁外执行，
		// 所以显式在写完 result 后 Unlock（上面已写明原因）。
		Manager.mu.Lock()

		result := &GhOstResult{
			TaskID:   taskID,
			Command:  "gh-ost " + maskSecrets(strings.Join(args, " ")),
			Duration: time.Since(start).String(),
			Output:   stdout.String(),
			Success:  err == nil,
		}

		if err != nil {
			result.Error = stderr.String() + "\n" + err.Error()
			Manager.status[taskID] = "failed"
		} else {
			Manager.status[taskID] = "completed"
		}

		Manager.tasks[taskID] = result
		Manager.mu.Unlock()

		// 回调必须放在锁**外面**：调用方如果在回调里查任务状态
		// （GetTaskStatus 需要 RLock），持有写锁时调用会永久死锁。
		// 当前所有调用方传的都是 nil，属于防御性修改。
		if cb != nil {
			cb(taskID, result.Output)
		}
	}
	safe.Go("ghost:"+taskID, run)

	return taskID
}

// maskSecrets 将命令行中的密码参数值替换为 ******，防止经任务状态接口回传前端
func maskSecrets(cmdline string) string {
	return regexp.MustCompile(`(--password=)([^ ]+)`).ReplaceAllString(cmdline, `$1******`)
}

// GetTaskStatus returns the status of an async task
func (m *TaskManager) GetTaskStatus(taskID string) (string, *GhOstResult) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.status[taskID], m.tasks[taskID]
}

// ListTasks returns all tracked tasks
func (m *TaskManager) ListTasks() map[string]string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	result := make(map[string]string)
	for k, v := range m.status {
		result[k] = v
	}
	return result
}

func buildArgs(task GhOstTask, ghOstPath string) []string {
	args := []string{
		"--host=" + task.Host,
		fmt.Sprintf("--port=%d", task.Port),
		"--user=" + task.User,
		"--password=" + task.Password,
		"--database=" + task.Database,
		"--table=" + task.Table,
		"--alter=" + task.Alter,
		"--verbose",
		"--skip-port-validation",
		"--allow-on-master",
		"--initially-drop-ghost-table",
		"--initially-drop-old-table",
		"--initially-drop-socket-file",
	}

	if task.ChunkSize > 0 {
		args = append(args, fmt.Sprintf("--chunk-size=%d", task.ChunkSize))
	}
	if task.MaxLoad != "" {
		args = append(args, "--max-load="+task.MaxLoad)
	}
	if task.CriticalLoad != "" {
		args = append(args, "--critical-load="+task.CriticalLoad)
	}
	if task.MaxLagMillis > 0 {
		args = append(args, fmt.Sprintf("--max-lag-millis=%d", task.MaxLagMillis))
	}
	if task.AllowOnMaster {
		args = append(args, "--allow-on-master")
	}
	if task.ExactRowCount {
		args = append(args, "--exact-rowcount")
	}
	if task.PostponeCutOver {
		args = append(args, "--postpone-cut-over-flag-file=/tmp/gh-ost-cutover.flag")
	}
	// Cloud-specific flags
	if task.AliyunRDS {
		args = append(args, "--aliyun-rds")
	}
	if task.AWSAurora {
		args = append(args, "--aws")
	}

	// gh-ost defaults to noop, only --execute actually applies changes
	if task.Execute {
		args = append(args, "--execute")
	}

	return args
}