fix: Fix file timeout and retry processing for long transmission tasks (#13746)

core/app/task/task.go

@@ -2,6 +2,7 @@ package task
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"path"
@@ -10,7 +11,6 @@ import (
 
 	"github.com/1Panel-dev/1Panel/core/app/model"
 	"github.com/1Panel-dev/1Panel/core/app/repo"
-	"github.com/1Panel-dev/1Panel/core/buserr"
 	"github.com/1Panel-dev/1Panel/core/constant"
 	"github.com/1Panel-dev/1Panel/core/global"
 	"github.com/1Panel-dev/1Panel/core/i18n"
@@ -19,9 +19,13 @@ import (
 )
 
 type ActionFunc func(*Task) error
+type ContextActionFunc func(context.Context, *Task) error
 type RollbackFunc func(*Task)
 
+var ErrExecutionUnconfirmed = errors.New("task execution termination could not be confirmed")
+
 type Task struct {
+	TaskCtx           context.Context
 	Name              string
 	TaskID            string
 	Logger            *logrus.Logger
@@ -41,6 +45,9 @@ type SubTask struct {
 	Retry             int
 	Timeout           time.Duration
 	Action            ActionFunc
+	ContextAction     ContextActionFunc
+	ShouldRetry       func(error) bool
+	RetryBackoff      time.Duration
 	Rollback          RollbackFunc
 	Error             error
 	IgnoreErr         bool
@@ -113,7 +120,7 @@ func NewTask(name, operate, taskScope, taskID string, resourceID uint) (*Task, e
 		Operate:    operate,
 	}
 	taskRepo := repo.NewITaskRepo()
-	task := &Task{Name: name, logFile: logFile, Logger: logger, taskRepo: taskRepo, Task: taskModel}
+	task := &Task{TaskID: taskID, Name: name, logFile: logFile, Logger: logger, taskRepo: taskRepo, Task: taskModel}
 	return task, nil
 }
 
@@ -132,61 +139,139 @@ func (t *Task) AddSubTaskWithOps(name string, action ActionFunc, rollback Rollba
 	t.SubTasks = append(t.SubTasks, subTask)
 }
 
+func (t *Task) AddSubTaskWithContext(name string, action ContextActionFunc, rollback RollbackFunc, retry int, timeout time.Duration) *SubTask {
+	subTask := &SubTask{RootTask: t, Name: name, Retry: retry, Timeout: timeout, ContextAction: action, Rollback: rollback}
+	t.SubTasks = append(t.SubTasks, subTask)
+	return subTask
+}
+
+func (t *Task) Context() context.Context {
+	if t.TaskCtx != nil {
+		return t.TaskCtx
+	}
+	return context.Background()
+}
+
 func (t *Task) AddSubTaskWithIgnoreErr(name string, action ActionFunc) {
 	subTask := &SubTask{RootTask: t, Name: name, Retry: 0, Timeout: 10 * time.Minute, Action: action, Rollback: nil, IgnoreErr: true}
 	t.SubTasks = append(t.SubTasks, subTask)
 }
 
 func (s *SubTask) Execute() error {
+	if s.Timeout < 0 || s.Retry < 0 || (s.Action == nil && s.ContextAction == nil) {
+		return fmt.Errorf("invalid subtask execution options")
+	}
 	subTaskName := s.Name
 	if s.Name == "" {
 		subTaskName = i18n.GetMsgByKey("SubTask")
 	}
 	s.RootTask.LogStart(subTaskName)
 	var err error
+	attempted := false
 	for i := 0; i < s.Retry+1; i++ {
+		if err = s.RootTask.Context().Err(); err != nil {
+			break
+		}
 		if i > 0 {
 			s.RootTask.Log(i18n.GetWithName("TaskRetry", strconv.Itoa(i)))
 		}
-		ctx, cancel := context.WithTimeout(context.Background(), s.Timeout)
-		defer cancel()
-
-		done := make(chan error)
-		go func() {
-			done <- s.Action(s.RootTask)
-		}()
-
-		select {
-		case <-ctx.Done():
-			s.RootTask.Log(i18n.GetWithName("TaskTimeout", subTaskName))
-			if s.CancelWhenTimeout {
-				return buserr.New(i18n.GetWithName("TaskTimeout", subTaskName))
-			}
-		case err = <-done:
-			if err != nil {
-				s.RootTask.Log(i18n.GetWithNameAndErr("SubTaskFailed", subTaskName, err))
-			} else {
+		var started bool
+		err, started = s.executeAttempt(subTaskName)
+		attempted = attempted || started
+		if err == nil {
 			s.RootTask.Log(i18n.GetWithName("SubTaskSuccess", subTaskName))
 			return nil
 		}
+		s.RootTask.Log(i18n.GetWithNameAndErr("SubTaskFailed", subTaskName, err))
+		if i == s.Retry || !s.canRetry(err) {
+			break
 		}
-
-		if i == s.Retry {
-			if s.Rollback != nil {
+		timer := time.NewTimer(s.retryDelay(i))
+		select {
+		case <-s.RootTask.Context().Done():
+			timer.Stop()
+			err = s.RootTask.Context().Err()
+		case <-timer.C:
+		}
+	}
+	if parentErr := s.RootTask.Context().Err(); parentErr != nil {
+		err = errors.Join(parentErr, err)
+	}
+	if attempted && s.Rollback != nil && !errors.Is(err, ErrExecutionUnconfirmed) {
 		s.Rollback(s.RootTask)
 	}
-		}
-		time.Sleep(1 * time.Second)
-	}
 	return err
 }
 
+func (s *SubTask) executeAttempt(name string) (error, bool) {
+	parent := s.RootTask.Context()
+	var ctx context.Context
+	var cancel context.CancelFunc
+	if s.Timeout == 0 {
+		ctx, cancel = context.WithCancel(parent)
+	} else {
+		ctx, cancel = context.WithTimeout(parent, s.Timeout)
+	}
+	defer cancel()
+	if err := ctx.Err(); err != nil {
+		return err, false
+	}
+	done := make(chan error, 1)
+	go func() {
+		if s.ContextAction != nil {
+			done <- s.ContextAction(ctx, s.RootTask)
+		} else {
+			done <- s.Action(s.RootTask)
+		}
+	}()
+	var err error
+	select {
+	case err = <-done:
+	case <-ctx.Done():
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			s.RootTask.Log(i18n.GetWithName("TaskTimeout", name))
+		}
+		err = <-done
+	}
+	if ctx.Err() != nil {
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			return errors.Join(fmt.Errorf("%s: %w", i18n.GetWithName("TaskTimeout", name), ctx.Err()), err), true
+		}
+		return errors.Join(ctx.Err(), err), true
+	}
+	return err, true
+}
+
+func (s *SubTask) canRetry(err error) bool {
+	if s.RootTask.Context().Err() != nil || errors.Is(err, context.Canceled) || errors.Is(err, ErrExecutionUnconfirmed) {
+		return false
+	}
+	if errors.Is(err, context.DeadlineExceeded) && (s.CancelWhenTimeout || s.ShouldRetry == nil) {
+		return false
+	}
+	return s.ShouldRetry == nil || s.ShouldRetry(err)
+}
+
+func (s *SubTask) retryDelay(attempt int) time.Duration {
+	if s.RetryBackoff <= 0 {
+		return time.Second
+	}
+	delay := min(s.RetryBackoff, 30*time.Second)
+	for i := 0; i < attempt && delay < 30*time.Second; i++ {
+		delay = min(delay*2, 30*time.Second)
+	}
+	return delay
+}
+
 func (t *Task) updateTask(task *model.Task) {
 	_ = t.taskRepo.Update(context.Background(), task)
 }
 
 func (t *Task) Execute() error {
 	if err := t.taskRepo.Save(context.Background(), t.Task); err != nil {
+		if t.logFile != nil {
+			_ = t.logFile.Close()
+		}
 		return err
 	}
 	var err error
@@ -200,15 +285,20 @@ func (t *Task) Execute() error {
 				t.Rollbacks = append(t.Rollbacks, subTask.Rollback)
 			}
 		} else {
-			if subTask.IgnoreErr {
+			if subTask.IgnoreErr && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, ErrExecutionUnconfirmed) {
 				err = nil
 				continue
 			}
 			t.Task.ErrorMsg = err.Error()
 			t.Task.Status = constant.StatusFailed
+			if errors.Is(err, context.Canceled) && !errors.Is(err, ErrExecutionUnconfirmed) {
+				t.Task.Status = constant.StatusCanceled
+			}
+			if !errors.Is(err, ErrExecutionUnconfirmed) {
 				for _, rollback := range t.Rollbacks {
 					rollback(t)
 				}
+			}
 			t.updateTask(t.Task)
 			break
 		}
@@ -222,7 +312,9 @@ func (t *Task) Execute() error {
 	t.Log("[TASK-END]")
 	t.Task.EndAt = time.Now()
 	t.updateTask(t.Task)
+	if t.logFile != nil {
 		_ = t.logFile.Close()
+	}
 	return err
 }
 

core/utils/cmd/cmdx.go

@@ -30,6 +30,7 @@ type CommandHelper struct {
 	taskItem           *task.Task
 	logger             *log.Logger
 	IgnoreExist1       bool
+	preserveErrorCause bool
 }
 
 type Option func(*CommandHelper)
@@ -356,8 +357,18 @@ func (c *CommandHelper) run(name string, arg ...string) (string, error) {
 	}()
 	select {
 	case err := <-done:
+		if c.preserveErrorCause && newContext != nil && newContext.Err() != nil {
+			if cmd.Process != nil && cmd.Process.Pid > 0 {
+				_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+			}
+			return "", newContext.Err()
+		}
 		if err != nil {
-			return handleErr(&stdout, &stderr, c.IgnoreExist1, err)
+			out, resultErr := handleErr(&stdout, &stderr, c.IgnoreExist1, err)
+			if c.preserveErrorCause && resultErr != nil {
+				resultErr = &commandError{message: resultErr.Error(), cause: err}
+			}
+			return out, resultErr
 		}
 		return stdout.String(), nil
 	case <-contextDone(newContext):
@@ -374,6 +385,9 @@ func (c *CommandHelper) run(name string, arg ...string) (string, error) {
 			err = newContext.Err()
 		}
 		<-done
+		if c.preserveErrorCause {
+			err = &commandError{message: err.Error(), cause: newContext.Err()}
+		}
 		return "", err
 	}
 }
@@ -414,6 +428,18 @@ func WithContext(ctx context.Context) Option {
 	}
 }
 
+func WithErrorCause() Option {
+	return func(s *CommandHelper) { s.preserveErrorCause = true }
+}
+
+type commandError struct {
+	message string
+	cause   error
+}
+
+func (e *commandError) Error() string { return e.message }
+func (e *commandError) Unwrap() error { return e.cause }
+
 func WithTimeout(timeout time.Duration) Option {
 	return func(s *CommandHelper) {
 		s.timeout = timeout

core/utils/ssh/ssh.go

@@ -1,11 +1,13 @@
 package ssh
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
 	"path"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/1Panel-dev/1Panel/core/app/repo"
@@ -245,6 +247,105 @@ func (c *SSHClient) RunWithStreamOutput(command string, outputCallback func(stri
 	return err
 }
 
+var ErrCommandTerminationUnconfirmed = errors.New("remote command termination could not be confirmed; check the source node before retrying")
+
+func (c *SSHClient) RunWithStreamOutputContext(ctx context.Context, command string, outputCallback func(string)) error {
+	return c.runStreamContext(ctx, command, outputCallback, 15*time.Second, 45*time.Second)
+}
+
+func (c *SSHClient) runStreamContext(ctx context.Context, command string, outputCallback func(string), heartbeatInterval, heartbeatTimeout time.Duration) (err error) {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	ctx, cancel := context.WithCancelCause(ctx)
+	closed := make(chan struct{})
+	stop := context.AfterFunc(ctx, func() {
+		c.Close()
+		close(closed)
+	})
+	heartbeatDone := make(chan struct{})
+	go func() {
+		defer close(heartbeatDone)
+		c.watchStreamConnection(ctx, cancel, heartbeatInterval, heartbeatTimeout)
+	}()
+	defer func() {
+		cause := context.Cause(ctx)
+		cancel(nil)
+		c.Close()
+		<-heartbeatDone
+		if !stop() {
+			<-closed
+		}
+		if cause != nil {
+			err = errors.Join(cause, ErrCommandTerminationUnconfirmed, err)
+		}
+	}()
+	session, err := c.Client.NewSession()
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+	writer := &streamCallbackWriter{callback: outputCallback}
+	session.Stdout = writer
+	session.Stderr = writer
+	if err := session.Run(command); err != nil {
+		var exitErr *gossh.ExitError
+		if !errors.As(err, &exitErr) {
+			return errors.Join(ErrCommandTerminationUnconfirmed, err)
+		}
+		return err
+	}
+	return nil
+}
+
+func (c *SSHClient) watchStreamConnection(ctx context.Context, cancel context.CancelCauseFunc, interval, timeout time.Duration) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+		}
+		result := make(chan error, 1)
+		go func() {
+			_, _, err := c.Client.SendRequest("keepalive@openssh.com", true, nil)
+			result <- err
+		}()
+		timer := time.NewTimer(timeout)
+		select {
+		case err := <-result:
+			timer.Stop()
+			if err != nil {
+				cancel(err)
+				return
+			}
+		case <-ctx.Done():
+			timer.Stop()
+			<-result // Closing the owned transport releases SendRequest.
+			return
+		case <-timer.C:
+			cancel(errors.New("SSH keepalive response timed out"))
+			<-result
+			return
+		}
+	}
+}
+
+type streamCallbackWriter struct {
+	mu       sync.Mutex
+	callback func(string)
+}
+
+func (w *streamCallbackWriter) Write(p []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.callback != nil {
+		w.callback(string(p))
+	}
+	return len(p), nil
+}
+
 func DialWithTimeout(network, addr string, useProxy bool, config *gossh.ClientConfig) (*gossh.Client, error) {
 	var conn net.Conn
 	var err error