golang shell

package shell

import (
	"context"
	"fmt"
	"os/exec"
	"time"
)

// 自定义输出结构体
type customOutput struct {
	outPut   chan string
	resetCtx chan struct{}
}

// Write 将输出写入到 customOutput 结构体中，并通知重置超时。
func (c customOutput) Write(p []byte) (int, error) {
	output := string(p)
	c.outPut <- output
	c.resetCtx <- struct{}{} // 每次有输出时通知续期
	return len(p), nil
}

type CustomShellCommand struct {
	cmd      *exec.Cmd
	outPipe  chan string
	resetCtx chan struct{}
	ctx      context.Context
	cancel   context.CancelFunc
	timer    *time.Timer
	timeout  time.Duration
}

// NewShell 创建新的 Shell 命令实例
func NewShell(command []string, timeout time.Duration) *CustomShellCommand {
	ctx, cancel := context.WithCancel(context.Background())
	outPipe := make(chan string, 100)
	resetCtx := make(chan struct{}, 1) // 用于通知续期

	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	cmd.Stdout = customOutput{outPipe, resetCtx}

	// 创建并返回 CustomShellCommand
	return &CustomShellCommand{
		cmd:      cmd,
		outPipe:  outPipe,
		resetCtx: resetCtx,
		ctx:      ctx,
		timeout:  timeout,
		cancel:   cancel,
		timer:    time.NewTimer(timeout), // 设置初始超时
	}
}

// Run 执行命令并处理续期超时
func (c *CustomShellCommand) Run() error {
	errChan := make(chan error, 1) // 创建一个缓冲区大小为1的错误通道

	go func() {
		err := c.cmd.Run() // 在 goroutine 中执行命令
		//if err != nil {
		//	fmt.Println("Command exited with error:", err)
		//}
		errChan <- err // 将错误发送到错误通道
		c.cancel()     // 当命令执行完成时取消上下文
	}()

	for {
		select {
		case <-c.ctx.Done(): // 当命令被取消或完成时退出
			//fmt.Println("Command context done.")
			return c.ctx.Err()
		case <-c.timer.C: // 当超时没有输出时，停止命令
			//fmt.Println("Command timed out.")
			c.cancel()
			return fmt.Errorf("command timed out")
		case <-c.resetCtx: // 当有输出时重置超时计时器
			if !c.timer.Stop() {
				<-c.timer.C // 确保通道被清空
			}
			c.timer.Reset(c.timeout) // 重置计时器为
		case err := <-errChan: // 捕获命令执行的错误
			return err // 返回错误
		}
	}
}

// Killed 取消任务
func (c *CustomShellCommand) Killed() {
	c.cancel()
}

// GetOutputPipe 返回命令的输出通道
func (c *CustomShellCommand) GetOutputPipe() <-chan string {
	return c.outPipe
}

// Close 关闭输出通道并取消上下文
func (c *CustomShellCommand) Close() {
	close(c.outPipe)
	c.cancel()
}