Shell任务超时 - 技术栈

背景：Shell功能为了实现实时转发消息需求，采用了两个线程，一个线程负责从io里面读取信息，当读取超过1行就转发，另一个线程监听Shell连接的会话（Session）是否结束，如果结束则通知第一个线程关闭。

问题：在执行一次输出多行的任务时，缺失第一行之后所有行数的数据

补充：在第一个负责从io里面读取信息的代码里面，为了实现受到关闭通知可以停止，所以用了Go语言的For和Select配合，select里面第一个是监听是否ctx.Done，第二个是默认的处理io里面一行Shell输出，因为外面for循环，可以一直处理处理io里面一行输出

分析：问领导讨论，第二个监听Sell连接的会话（Session）结束后，通知第一个线程关闭。第一个线程在关闭时，io里面还有积压的消息没有处理，这时候我们最好想办法处理当前io里面积压的消息。领导让我找一找有没有io接口里面有没有直接读取全部积压消息的方式。一开始问GPT，给出来的方式并没有解决问题，于是我问DeepSeek给出来的方式解决了问题，现在对DeepSeek的解决思路进行一个分析。

DeepSeek分析更具体，当前实现中，一旦ctx.Done()被触发，第一个处理消息的errgroup会立即返回，导致还有未读的数据留在stdout里面，这时候我希望将全部的数据读出来，添加到result里面，并记录日志。

注意，当context被取消时，可能stdout管到没有关闭，这时候需要一个非阻塞的方式读取数据，或者设置一个截止时间避免无限等待。当然当会话被关闭时，stdout也会被关闭，这是ReadString会返回io.EOF。

这时候知道使用ioutil里面的ReadAll

补充：问题代码

javascript 复制代码

func (s *ShellTask) Run(ctx context.Context, timeout time.Duration) (int, string, error) {
	zap.S().Debugf("shell script is running : %s", s.Script)
	// 创建SSH配置
	config := &ssh.ClientConfig{
		User: s.Node.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(s.Node.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         timeout,
	}

	// 连接SSH服务器
	client, err := DialContext(ctx, "tcp", fmt.Sprintf("%s:%d", s.Node.Ipv4Address, s.Node.Port), config)
	if err != nil {
		return 1, "", fmt.Errorf("Failed to dial SSH server: %s", err.Error())
	}
	defer client.Close()

	// 创建新会话
	session, err := client.NewSession()
	if err != nil {
		return 1, "", fmt.Errorf("Failed to create SSH session: %s", err.Error())
	}
	defer session.Close()
	//stdin, err := session.StdinPipe()
	stdout, err := session.StdoutPipe()
	if err != nil {
		return 1, "", fmt.Errorf("Error pipr stdout: %s", err.Error())
	}
	var b bytes.Buffer
	session.Stderr = &b
	err = session.Start(s.Script)
	if err != nil {
		return 1, "", fmt.Errorf("Error run command over SSH: %s", err.Error())
	}
	result := ""
	eg, ctx := errgroup.WithContext(ctx)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	eg.Go(func() error {
		for {
			select {
			case <-ctx.Done():
				if ctx.Err() == context.DeadlineExceeded {
					zap.S().Infof("Shell task execution timeout.")
				}
				_ = session.Close()
				return ctx.Err()
			}
		}
	})
	eg.Go(func() error {
		reader := bufio.NewReader(stdout)
		for {
			select {
			case <-ctx.Done():
				// 这里先将reader现在有的所有信息直接吐给result，可以先打印一个日志看一下
				
				// 如果 ctx 被取消了，退出循环
				if ctx.Err() == context.DeadlineExceeded {
					zap.S().Infof("Shell task execution timeout.")
				}
				return ctx.Err()
			default:
				line, err2 := reader.ReadString('\n')
				if err2 == io.EOF {
					continue
				}
				if err2 != nil && err2 != io.EOF {
					break
				}
				line = strings.Replace(line, "\n", "", -1)
				zap.S().Debugf("--> %s", line)
				err := instance.notifyTaskResult(s.ClientId, s.TaskId, line)
				if err != nil {
					zap.S().Errorf("Notify error. %s", err.Error())
				}
				result += line + "\n"
			}
		}
		return nil
	})
	eg.Go(func() error {
		err := session.Wait()
		zap.S().Debugf("Shell session finished. %v", err)
		cancel()
		return err
	})
	err = eg.Wait()
	zap.S().Infof("任务运行完成")
	if err != nil {
		var exitError *ssh.ExitError
		if errors.As(err, &exitError) {
			//session 运行错误
			return exitError.ExitStatus(), result, err
		} else if errors.Is(err, context.DeadlineExceeded) {
			//session超时错误
			return 124, result, err
		} else {
			//正常退出session
			//通常是context canceled
			return 0, result, err
		}
	}

	if b.Len() > 0 {
		return 128, result, errors.New(b.String())
	}
	return 0, result, nil

}

解决方法

javascript 复制代码

eg.Go(func() error {
    reader := bufio.NewReader(stdout)
    for {
        select {
        case <-ctx.Done():
            // 上下文取消时读取剩余数据
            if ctx.Err() == context.DeadlineExceeded {
                zap.S().Infof("Shell task execution timeout. Collecting remaining output...")
            } else {
                zap.S().Infof("Shell task cancelled. Collecting remaining output...")
            }
            
            // 读取所有剩余数据
            remaining, err := ioutil.ReadAll(reader)
            if err != nil && err != io.EOF {
                zap.S().Errorf("Error reading remaining data: %v", err)
            }
            
            // 如果有剩余数据则追加到结果
            if len(remaining) > 0 {
                cleaned := strings.ReplaceAll(string(remaining), "\n", "\\n")
                zap.S().Debugf("Captured remaining output: %s", cleaned)
                result += string(remaining)
            }
            
            return ctx.Err()
            
        default:
            line, err2 := reader.ReadString('\n')
            if err2 != nil {
                if err2 == io.EOF {
                    // 遇到EOF继续检查上下文状态
                    continue
                }
                // 其他错误直接返回
                return fmt.Errorf("stdout read error: %w", err2)
            }
            
            // 处理正常输出行
            line = strings.TrimSuffix(line, "\n")
            zap.S().Debugf("--> %s", line)
            
            if notifyErr := instance.notifyTaskResult(s.ClientId, s.TaskId, line); notifyErr != nil {
                zap.S().Errorf("Notification failed: %v", notifyErr)
            }
            
            result += line + "\n"
        }
    }
})