aiops/backend/pkg/remoteexec/persistent_client.go

// persistent_client.go - 任务级 SSH 长连接，多次 Exec 复用同一 client // Logan.Li
package remoteexec

import (
	"bytes"
	"context"
	"fmt"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

// PersistentClient 可复用的 SSH 长连接（任务生命周期内持有，Close 释放） // Logan.Li
type PersistentClient struct {
	cfg    SSHConfig
	client *ssh.Client
	mu     sync.Mutex
}

// NewPersistentClient 建立 SSH 长连接 // Logan.Li
func NewPersistentClient(cfg SSHConfig) (*PersistentClient, error) {
	client, err := connect(cfg)
	if err != nil {
		return nil, err
	}
	return &PersistentClient{cfg: cfg, client: client}, nil
}

// Exec 在长连接上执行命令；连接断开时自动重连一次 // Logan.Li
func (p *PersistentClient) Exec(command string, timeout time.Duration) (*ExecResult, error) {
	return p.ExecContext(context.Background(), command, timeout)
}

// ExecContext 在长连接上执行命令，并在调用方取消时关闭当前远端 session。 // Logan.Li
func (p *PersistentClient) ExecContext(ctx context.Context, command string, timeout time.Duration) (*ExecResult, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	result, err := p.execLockedContext(ctx, command, timeout)
	if err == nil {
		return result, nil
	}
	if ctx.Err() != nil {
		return result, ctx.Err()
	}
	// 连接级失败：重连后重试一次 // Logan.Li
	if p.client != nil {
		_ = p.client.Close()
		p.client = nil
	}
	client, connErr := connect(p.cfg)
	if connErr != nil {
		return result, err
	}
	p.client = client
	return p.execLockedContext(ctx, command, timeout)
}

// execLocked 假定已持锁，在当前 client 上执行 // Logan.Li
func (p *PersistentClient) execLocked(command string, timeout time.Duration) (*ExecResult, error) {
	return p.execLockedContext(context.Background(), command, timeout)
}

// execLockedContext 假定已持锁，并将调用方上下文传递给当前 SSH session。 // Logan.Li
func (p *PersistentClient) execLockedContext(ctx context.Context, command string, timeout time.Duration) (*ExecResult, error) {
	if p.client == nil {
		client, err := connect(p.cfg)
		if err != nil {
			return nil, err
		}
		p.client = client
	}
	return execWithClientContext(ctx, p.client, p.cfg, command, timeout)
}

// Close 关闭长连接 // Logan.Li
func (p *PersistentClient) Close() error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.client == nil {
		return nil
	}
	err := p.client.Close()
	p.client = nil
	return err
}

// execWithClient 在已有 SSH client 上执行命令（供 Exec / PersistentClient 共用） // Logan.Li
func execWithClient(client *ssh.Client, cfg SSHConfig, command string, timeout time.Duration) (*ExecResult, error) {
	return execWithClientContext(context.Background(), client, cfg, command, timeout)
}

// execWithClientContext 在已有 SSH client 上执行命令，并响应调用方取消。 // Logan.Li
func execWithClientContext(parent context.Context, client *ssh.Client, cfg SSHConfig, command string, timeout time.Duration) (*ExecResult, error) {
	start := time.Now()
	result := &ExecResult{Host: cfg.Host, Port: cfg.Port}

	session, err := client.NewSession()
	if err != nil {
		result.Duration = time.Since(start).Milliseconds()
		result.Error = fmt.Sprintf("创建SSH会话失败: %v", err)
		return result, err
	}
	defer session.Close()

	if timeout <= 0 {
		timeout = 60 * time.Second
	}

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

	type execOut struct {
		output []byte
		err    error
	}
	ch := make(chan execOut, 1)

	go func() {
		var stdout, stderr bytes.Buffer
		session.Stdout = &stdout
		session.Stderr = &stderr
		runErr := session.Run(command)

		output := stdout.Bytes()
		if stderr.Len() > 0 {
			if len(output) > 0 {
				output = append(output, '\n')
			}
			output = append(output, stderr.Bytes()...)
		}
		ch <- execOut{output: output, err: runErr}
	}()

	select {
	case <-ctx.Done():
		_ = session.Signal(ssh.SIGKILL)
		_ = session.Close()
		result.Duration = time.Since(start).Milliseconds()
		if parent.Err() != nil {
			result.Error = "执行已取消"
		} else {
			result.Error = fmt.Sprintf("执行超时（%v）", timeout)
		}
		result.ExitCode = -1
		return result, ctx.Err()
	case out := <-ch:
		result.Duration = time.Since(start).Milliseconds()
		result.Output = string(out.output)

		if out.err != nil {
			if exitErr, ok := out.err.(*ssh.ExitError); ok {
				result.ExitCode = exitErr.ExitStatus()
				result.Success = false
			} else {
				result.ExitCode = -1
				result.Error = out.err.Error()
				result.Success = false
				return result, out.err
			}
		} else {
			result.ExitCode = 0
			result.Success = true
		}
	}
	return result, nil
}