优化FTP传输超时控制及失败重试机制

utils/ftpclient/ftpclient.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net"
 	"os"
 	"path/filepath"
 	"sync"
@@ -24,6 +25,7 @@ const (
 	defaultLogInterval     = 2 * time.Second
 	DefaultUploadTimeout   = 5 * time.Minute
 	DefaultDownloadTimeout = 5 * time.Minute
+	useFTPConnDeadline     = true
 )
 
 var (
@@ -118,6 +120,40 @@ type stopError struct{ err error }
 
 func (e *stopError) Error() string { return e.err.Error() }
 
+func dialFTP(ctx context.Context, serverAddr string) (*ftp.ServerConn, error) {
+	if !useFTPConnDeadline {
+		return ftp.Dial(serverAddr, ftp.DialWithContext(ctx))
+	}
+
+	return ftp.Dial(serverAddr, ftp.DialWithDialFunc(func(network, address string) (net.Conn, error) {
+		conn, err := (&net.Dialer{}).DialContext(ctx, network, address)
+		if err != nil {
+			return nil, err
+		}
+
+		if deadline, ok := ctx.Deadline(); ok {
+			if err := conn.SetDeadline(deadline); err != nil {
+				_ = conn.Close()
+				return nil, err
+			}
+		}
+
+		return conn, nil
+	}))
+}
+
+func waitRetry(ctx context.Context, interval time.Duration) error {
+	timer := time.NewTimer(interval)
+	defer timer.Stop()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}
+
 func UploadFileToFtp(ctx context.Context, localFile, serverAddr, loginUser, loginPass string, timeout time.Duration) (string, error) {
 	unlock, ok := tryLockFile(localFile)
 	if !ok {
@@ -158,7 +194,7 @@ func UploadFileToFtp(ctx context.Context, localFile, serverAddr, loginUser, logi
 		}
 
 		err := func() error {
-			c, err := ftp.Dial(serverAddr, ftp.DialWithContext(timeoutCtx))
+			c, err := dialFTP(timeoutCtx, serverAddr)
 			if err != nil {
 				return &stopError{err}
 			}
@@ -192,7 +228,11 @@ func UploadFileToFtp(ctx context.Context, localFile, serverAddr, loginUser, logi
 			if lfe, ok := err.(*stopError); ok {
 				return "", lfe.err
 			}
-			time.Sleep(defaultRtyInterval)
+
+			if err := waitRetry(timeoutCtx, defaultRtyInterval); err != nil {
+				return "", err
+			}
+
 			continue
 		}
 
@@ -236,7 +276,7 @@ func DownloadFileFromFtp(ctx context.Context, serverAddr, loginUser, loginPass,
 		}
 
 		err := func() error {
-			c, err := ftp.Dial(serverAddr, ftp.DialWithContext(timeoutCtx))
+			c, err := dialFTP(timeoutCtx, serverAddr)
 			if err != nil {
 				return &stopError{err}
 			}
@@ -279,7 +319,11 @@ func DownloadFileFromFtp(ctx context.Context, serverAddr, loginUser, loginPass,
 			if lfe, ok := err.(*stopError); ok {
 				return "", lfe.err
 			}
-			time.Sleep(defaultRtyInterval)
+
+			if err := waitRetry(timeoutCtx, defaultRtyInterval); err != nil {
+				return "", err
+			}
+
 			continue
 		}