b612/starssh
1
0
Fork 0
Code Issues Pull Requests Projects Releases 6 Wiki Activity

feat(sftp): 增加可选的并发传输策略

Browse Source 
- 增加 SFTP client 级配置，支持 packet size、单文件并发请求数、并发读和并发写
- 将吞吐优化限制在 StarSSH 托管的 SFTP client 路径中
- 上传和下载在显式启用时使用并发快路径，同时保留原子传输生命周期
- 避免快路径失败前提前上报 100% 进度
- 补充安全校验、托管 client 配置、进度、取消和下载对齐的回归测试
This commit is contained in:
兔子
2026-06-22 03:26:47 +08:00
parent 0c23e7d4bf
commit 672a111ec1
2 changed files with 1036 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sftp.go
+370 -62
View File
@@ -13,6 +13,7 @@ import (
"path"
"path/filepath"
"strings"
"sync"
"time"
"github.com/pkg/sftp"
@@ -36,6 +37,19 @@ type SFTPTransferOptions struct {
VerifySize *bool
VerifyChecksum *bool
TempSuffix string
Client SFTPClientOptions
}
// SFTPClientOptions controls the underlying SFTP protocol client.
//
// These options only apply when StarSSH creates the SFTP client internally.
// They are intentionally separate from BufferSize, which only controls the
// local copy buffer used by transfer progress reporting.
type SFTPClientOptions struct {
MaxPacketSize int
MaxConcurrentRequestsPerFile int
ConcurrentReads *bool
ConcurrentWrites *bool
}
type resolvedSFTPTransferOptions struct {
@@ -48,6 +62,38 @@ type resolvedSFTPTransferOptions struct {
VerifySize bool
VerifyChecksum bool
TempSuffix string
Client resolvedSFTPClientOptions
}
type resolvedSFTPClientOptions struct {
MaxPacketSize int
MaxConcurrentRequestsPerFile int
ConcurrentReads *bool
ConcurrentWrites *bool
}
type sftpConcurrentReaderFrom interface {
ReadFromWithConcurrency(io.Reader, int) (int64, error)
}
type sftpUploadProgressReader struct {
ctx context.Context
reader io.Reader
total int64
progress func(float64)
mu sync.Mutex
copied int64
}
type sftpDownloadProgressWriter struct {
ctx context.Context
writer io.Writer
total int64
progress func(float64)
mu sync.Mutex
copied int64
}
type SFTPErrorCategory string
@@ -106,6 +152,7 @@ var (
sftpVerifyLocalSizeFunc = verifyLocalSize
sftpLocalFileSHA256Func = localFileSHA256
sftpRemoteFileSHA256Func = remoteFileSHA256
sftpNewClientFunc = sftp.NewClient
)
func DefaultSFTPTransferOptions() SFTPTransferOptions {
@@ -121,6 +168,16 @@ func DefaultSFTPTransferOptions() SFTPTransferOptions {
}
}
func ThroughputSFTPTransferOptions() SFTPTransferOptions {
opts := DefaultSFTPTransferOptions()
opts.Client = SFTPClientOptions{
ConcurrentReads: SFTPBool(true),
ConcurrentWrites: SFTPBool(true),
MaxConcurrentRequestsPerFile: 32,
}
return opts
}
func SFTPBool(value bool) *bool {
return &value
}
@@ -316,51 +373,40 @@ func (fs *SFTPFileSystem) Rename(ctx context.Context, oldPath string, newPath st
})
}
func normalizeSFTPTransferOptions(options *SFTPTransferOptions) resolvedSFTPTransferOptions {
func normalizeSFTPTransferOptions(options *SFTPTransferOptions) (resolvedSFTPTransferOptions, error) {
opts := DefaultSFTPTransferOptions()
if options == nil {
return resolvedSFTPTransferOptions{
BufferSize: opts.BufferSize,
Progress: opts.Progress,
RetryCount: normalizeSFTPRetryCount(derefSFTPInt(opts.RetryCount, defaultSFTPRetryCount)),
RetryInitialBackoff: derefSFTPDuration(opts.RetryInitialBackoff, defaultSFTPRetryInitialBackoff),
AtomicUpload: derefSFTPBool(opts.AtomicUpload, true),
AtomicDownload: derefSFTPBool(opts.AtomicDownload, true),
VerifySize: derefSFTPBool(opts.VerifySize, true),
VerifyChecksum: derefSFTPBool(opts.VerifyChecksum, false),
TempSuffix: normalizeSFTPTempSuffix(opts.TempSuffix),
if options != nil {
if options.BufferSize > 0 {
opts.BufferSize = options.BufferSize
}
if options.Progress != nil {
opts.Progress = options.Progress
}
if options.RetryCount != nil {
opts.RetryCount = options.RetryCount
}
if options.RetryInitialBackoff != nil {
opts.RetryInitialBackoff = options.RetryInitialBackoff
}
if options.AtomicUpload != nil {
opts.AtomicUpload = options.AtomicUpload
}
if options.AtomicDownload != nil {
opts.AtomicDownload = options.AtomicDownload
}
if options.VerifySize != nil {
opts.VerifySize = options.VerifySize
}
if options.VerifyChecksum != nil {
opts.VerifyChecksum = options.VerifyChecksum
}
if strings.TrimSpace(options.TempSuffix) != "" {
opts.TempSuffix = options.TempSuffix
}
opts.Client = options.Client
}
if options.BufferSize > 0 {
opts.BufferSize = options.BufferSize
}
if options.Progress != nil {
opts.Progress = options.Progress
}
if options.RetryCount != nil {
opts.RetryCount = options.RetryCount
}
if options.RetryInitialBackoff != nil {
opts.RetryInitialBackoff = options.RetryInitialBackoff
}
if options.AtomicUpload != nil {
opts.AtomicUpload = options.AtomicUpload
}
if options.AtomicDownload != nil {
opts.AtomicDownload = options.AtomicDownload
}
if options.VerifySize != nil {
opts.VerifySize = options.VerifySize
}
if options.VerifyChecksum != nil {
opts.VerifyChecksum = options.VerifyChecksum
}
if strings.TrimSpace(options.TempSuffix) != "" {
opts.TempSuffix = options.TempSuffix
}
return resolvedSFTPTransferOptions{
resolved := resolvedSFTPTransferOptions{
BufferSize: opts.BufferSize,
Progress: opts.Progress,
RetryCount: normalizeSFTPRetryCount(derefSFTPInt(opts.RetryCount, defaultSFTPRetryCount)),
@@ -371,6 +417,11 @@ func normalizeSFTPTransferOptions(options *SFTPTransferOptions) resolvedSFTPTran
VerifyChecksum: derefSFTPBool(opts.VerifyChecksum, false),
TempSuffix: normalizeSFTPTempSuffix(opts.TempSuffix),
}
resolved.Client = normalizeSFTPClientOptions(opts.Client)
if err := validateResolvedSFTPTransferOptions(resolved); err != nil {
return resolvedSFTPTransferOptions{}, err
}
return resolved, nil
}
func derefSFTPBool(value *bool, fallback bool) bool {
@@ -409,13 +460,56 @@ func normalizeSFTPRetryCount(value int) int {
return value
}
func normalizeSFTPClientOptions(options SFTPClientOptions) resolvedSFTPClientOptions {
return resolvedSFTPClientOptions{
MaxPacketSize: options.MaxPacketSize,
MaxConcurrentRequestsPerFile: options.MaxConcurrentRequestsPerFile,
ConcurrentReads: options.ConcurrentReads,
ConcurrentWrites: options.ConcurrentWrites,
}
}
func validateResolvedSFTPTransferOptions(opts resolvedSFTPTransferOptions) error {
return validateResolvedSFTPClientOptions(opts.Client)
}
func validateResolvedSFTPClientOptions(opts resolvedSFTPClientOptions) error {
if opts.MaxPacketSize < 0 {
return errors.New("sftp max packet size must not be negative")
}
if opts.MaxConcurrentRequestsPerFile < 0 {
return errors.New("sftp max concurrent requests per file must not be negative")
}
return nil
}
func rejectExternalSFTPClientOptions(opts resolvedSFTPClientOptions) error {
if opts.MaxPacketSize != 0 ||
opts.MaxConcurrentRequestsPerFile != 0 ||
opts.ConcurrentReads != nil ||
opts.ConcurrentWrites != nil {
return errors.New("sftp client options require StarSSH-managed SFTP client")
}
return nil
}
func validateSFTPUploadOptions(opts resolvedSFTPTransferOptions) error {
if derefSFTPBool(opts.Client.ConcurrentWrites, false) && !opts.AtomicUpload {
return errors.New("sftp concurrent writes require atomic upload")
}
return nil
}
func (star *StarSSH) runSFTPClientOperation(ctx context.Context, operation string, remotePath string, fn func(*sftp.Client) error) error {
if err := ensureContext(ctx); err != nil {
return err
}
opts := normalizeSFTPTransferOptions(nil)
opts, err := normalizeSFTPTransferOptions(nil)
if err != nil {
return err
}
return executeSFTPRetry(ctx, operation, "", remotePath, opts, func(attempt int) error {
return star.withIsolatedSFTPClient(ctx, fn)
return star.withIsolatedSFTPClient(ctx, opts.Client, fn)
})
}
@@ -423,23 +517,27 @@ func (star *StarSSH) runSFTPClientOperationNoRetry(ctx context.Context, fn func(
if err := ensureContext(ctx); err != nil {
return err
}
return star.withIsolatedSFTPClient(ctx, fn)
return star.withIsolatedSFTPClient(ctx, resolvedSFTPClientOptions{}, fn)
}
func (star *StarSSH) CreateSftpClient() (*sftp.Client, error) {
return star.createSFTPClientWithOptions(resolvedSFTPClientOptions{})
}
func (star *StarSSH) createSFTPClientWithOptions(options resolvedSFTPClientOptions) (*sftp.Client, error) {
client, err := star.requireSSHClient()
if err != nil {
return nil, err
}
return sftp.NewClient(client)
return sftpNewClientFunc(client, buildSFTPClientOptions(options)...)
}
func (star *StarSSH) withIsolatedSFTPClient(ctx context.Context, fn func(*sftp.Client) error) error {
func (star *StarSSH) withIsolatedSFTPClient(ctx context.Context, options resolvedSFTPClientOptions, fn func(*sftp.Client) error) error {
if err := ensureContext(ctx); err != nil {
return err
}
client, err := star.CreateSftpClient()
client, err := star.createSFTPClientWithOptions(options)
if err != nil {
return err
}
@@ -448,6 +546,23 @@ func (star *StarSSH) withIsolatedSFTPClient(ctx context.Context, fn func(*sftp.C
return fn(client)
}
func buildSFTPClientOptions(options resolvedSFTPClientOptions) []sftp.ClientOption {
clientOptions := make([]sftp.ClientOption, 0, 4)
if options.MaxPacketSize > 0 {
clientOptions = append(clientOptions, sftp.MaxPacketChecked(options.MaxPacketSize))
}
if options.MaxConcurrentRequestsPerFile > 0 {
clientOptions = append(clientOptions, sftp.MaxConcurrentRequestsPerFile(options.MaxConcurrentRequestsPerFile))
}
if options.ConcurrentReads != nil {
clientOptions = append(clientOptions, sftp.UseConcurrentReads(*options.ConcurrentReads))
}
if options.ConcurrentWrites != nil {
clientOptions = append(clientOptions, sftp.UseConcurrentWrites(*options.ConcurrentWrites))
}
return clientOptions
}
func (star *StarSSH) getReusableSFTPClient() (*sftp.Client, error) {
if star == nil {
return nil, errors.New("ssh client is nil")
@@ -526,7 +641,7 @@ func (star *StarSSH) runSFTPWithRetry(
fn func(context.Context, *sftp.Client, resolvedSFTPTransferOptions) error,
) error {
return executeSFTPRetry(ctx, operation, localPath, remotePath, opts, func(attempt int) error {
return star.withIsolatedSFTPClient(ctx, func(client *sftp.Client) error {
return star.withIsolatedSFTPClient(ctx, opts.Client, func(client *sftp.Client) error {
return fn(ctx, client, opts)
})
})
@@ -537,7 +652,13 @@ func (star *StarSSH) SftpTransferOut(localFilePath, remotePath string) error {
}
func (star *StarSSH) SftpTransferOutContext(ctx context.Context, localFilePath, remotePath string, options *SFTPTransferOptions) error {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return err
}
if err := validateSFTPUploadOptions(opts); err != nil {
return err
}
return star.runSFTPWithRetry(ctx, "sftp_put_file", localFilePath, remotePath, opts, func(ctx context.Context, client *sftp.Client, opts resolvedSFTPTransferOptions) error {
return transferOutContext(ctx, client, localFilePath, remotePath, opts)
})
@@ -548,7 +669,16 @@ func SftpTransferOut(localFilePath, remotePath string, sftpClient *sftp.Client)
}
func SftpTransferOutWithContext(ctx context.Context, localFilePath, remotePath string, sftpClient *sftp.Client, options *SFTPTransferOptions) error {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return err
}
if err := validateSFTPUploadOptions(opts); err != nil {
return err
}
if err := rejectExternalSFTPClientOptions(opts.Client); err != nil {
return err
}
return executeSFTPRetry(ctx, "sftp_put_file", localFilePath, remotePath, opts, func(attempt int) error {
return transferOutContext(ctx, sftpClient, localFilePath, remotePath, opts)
})
@@ -559,7 +689,13 @@ func (star *StarSSH) SftpTransferOutByte(localData []byte, remotePath string) er
}
func (star *StarSSH) SftpTransferOutByteContext(ctx context.Context, localData []byte, remotePath string, options *SFTPTransferOptions) error {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return err
}
if err := validateSFTPUploadOptions(opts); err != nil {
return err
}
return star.runSFTPWithRetry(ctx, "sftp_put_bytes", "", remotePath, opts, func(ctx context.Context, client *sftp.Client, opts resolvedSFTPTransferOptions) error {
return transferOutByteContext(ctx, client, localData, remotePath, opts)
})
@@ -570,7 +706,16 @@ func SftpTransferOutByte(localData []byte, remotePath string, sftpClient *sftp.C
}
func SftpTransferOutByteWithContext(ctx context.Context, localData []byte, remotePath string, sftpClient *sftp.Client, options *SFTPTransferOptions) error {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return err
}
if err := validateSFTPUploadOptions(opts); err != nil {
return err
}
if err := rejectExternalSFTPClientOptions(opts.Client); err != nil {
return err
}
return executeSFTPRetry(ctx, "sftp_put_bytes", "", remotePath, opts, func(attempt int) error {
return transferOutByteContext(ctx, sftpClient, localData, remotePath, opts)
})
@@ -595,10 +740,13 @@ func (star *StarSSH) SftpTransferInByte(remotePath string) ([]byte, error) {
}
func (star *StarSSH) SftpTransferInByteContext(ctx context.Context, remotePath string, options *SFTPTransferOptions) ([]byte, error) {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return nil, err
}
var data []byte
err := star.runSFTPWithRetry(ctx, "sftp_get_bytes", "", remotePath, opts, func(ctx context.Context, client *sftp.Client, opts resolvedSFTPTransferOptions) error {
err = star.runSFTPWithRetry(ctx, "sftp_get_bytes", "", remotePath, opts, func(ctx context.Context, client *sftp.Client, opts resolvedSFTPTransferOptions) error {
out, runErr := transferInByteContext(ctx, client, remotePath, opts)
if runErr != nil {
return runErr
@@ -617,10 +765,16 @@ func SftpTransferInByte(remotePath string, sftpClient *sftp.Client) ([]byte, err
}
func SftpTransferInByteWithContext(ctx context.Context, remotePath string, sftpClient *sftp.Client, options *SFTPTransferOptions) ([]byte, error) {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return nil, err
}
if err := rejectExternalSFTPClientOptions(opts.Client); err != nil {
return nil, err
}
var data []byte
err := executeSFTPRetry(ctx, "sftp_get_bytes", "", remotePath, opts, func(attempt int) error {
err = executeSFTPRetry(ctx, "sftp_get_bytes", "", remotePath, opts, func(attempt int) error {
out, runErr := transferInByteContext(ctx, sftpClient, remotePath, opts)
if runErr != nil {
return runErr
@@ -639,7 +793,10 @@ func (star *StarSSH) SftpTransferIn(src, dst string) error {
}
func (star *StarSSH) SftpTransferInContext(ctx context.Context, src, dst string, options *SFTPTransferOptions) error {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return err
}
return star.runSFTPWithRetry(ctx, "sftp_get_file", dst, src, opts, func(ctx context.Context, client *sftp.Client, opts resolvedSFTPTransferOptions) error {
return transferInContext(ctx, client, src, dst, opts)
})
@@ -650,7 +807,13 @@ func SftpTransferIn(src, dst string, sftpClient *sftp.Client) error {
}
func SftpTransferInWithContext(ctx context.Context, src, dst string, sftpClient *sftp.Client, options *SFTPTransferOptions) error {
opts := normalizeSFTPTransferOptions(options)
opts, err := normalizeSFTPTransferOptions(options)
if err != nil {
return err
}
if err := rejectExternalSFTPClientOptions(opts.Client); err != nil {
return err
}
return executeSFTPRetry(ctx, "sftp_get_file", dst, src, opts, func(attempt int) error {
return transferInContext(ctx, sftpClient, src, dst, opts)
})
@@ -712,7 +875,7 @@ func transferOutContext(ctx context.Context, sftpClient *sftp.Client, localFileP
return err
}
if _, err := sftpCopyWithProgressFunc(ctx, dstFile, srcFile, opts.BufferSize, stat.Size(), opts.Progress); err != nil {
if _, err := copyUploadWithProgressContext(ctx, dstFile, srcFile, opts.BufferSize, stat.Size(), opts.Progress, opts); err != nil {
_ = dstFile.Close()
return err
}
@@ -792,7 +955,7 @@ func transferOutByteContext(ctx context.Context, sftpClient *sftp.Client, localD
}
reader := bytes.NewReader(localData)
if _, err := sftpCopyWithProgressFunc(ctx, dstFile, reader, opts.BufferSize, int64(len(localData)), opts.Progress); err != nil {
if _, err := copyUploadWithProgressContext(ctx, dstFile, reader, opts.BufferSize, int64(len(localData)), opts.Progress, opts); err != nil {
_ = dstFile.Close()
return err
}
@@ -877,7 +1040,7 @@ func transferInContext(ctx context.Context, sftpClient *sftp.Client, src, dst st
}()
}
if _, err := sftpCopyWithProgressFunc(ctx, dstFile, srcFile, opts.BufferSize, stat.Size(), opts.Progress); err != nil {
if _, err := copyDownloadWithProgressContext(ctx, dstFile, srcFile, opts.BufferSize, stat.Size(), opts.Progress, opts); err != nil {
_ = dstFile.Close()
return err
}
@@ -948,7 +1111,7 @@ func transferInByteContext(ctx context.Context, sftpClient *sftp.Client, remoteP
}
var out bytes.Buffer
if _, err := sftpCopyWithProgressFunc(ctx, &out, srcFile, opts.BufferSize, stat.Size(), opts.Progress); err != nil {
if _, err := copyDownloadWithProgressContext(ctx, &out, srcFile, opts.BufferSize, stat.Size(), opts.Progress, opts); err != nil {
return nil, err
}
@@ -1149,6 +1312,151 @@ func copyWithProgressContext(ctx context.Context, dst io.Writer, src io.Reader,
return copied, nil
}
func copyUploadWithProgressContext(
ctx context.Context,
dst io.Writer,
src io.Reader,
bufSize int,
total int64,
progress func(float64),
opts resolvedSFTPTransferOptions,
) (int64, error) {
if !derefSFTPBool(opts.Client.ConcurrentWrites, false) {
return sftpCopyWithProgressFunc(ctx, dst, src, bufSize, total, progress)
}
readerFrom, ok := dst.(sftpConcurrentReaderFrom)
if !ok {
return sftpCopyWithProgressFunc(ctx, dst, src, bufSize, total, progress)
}
if err := ensureContext(ctx); err != nil {
return 0, err
}
if progress != nil && total > 0 {
progress(0)
}
wrappedSrc := &sftpUploadProgressReader{
ctx: ctx,
reader: src,
total: total,
progress: progress,
}
written, err := readerFrom.ReadFromWithConcurrency(wrappedSrc, opts.Client.MaxConcurrentRequestsPerFile)
if err != nil {
return written, err
}
if err := ensureContext(ctx); err != nil {
return written, err
}
reportProgress(progress, written, total)
return written, nil
}
func copyDownloadWithProgressContext(
ctx context.Context,
dst io.Writer,
src io.Reader,
bufSize int,
total int64,
progress func(float64),
opts resolvedSFTPTransferOptions,
) (int64, error) {
if !derefSFTPBool(opts.Client.ConcurrentReads, false) {
return sftpCopyWithProgressFunc(ctx, dst, src, bufSize, total, progress)
}
writerTo, ok := src.(io.WriterTo)
if !ok {
return sftpCopyWithProgressFunc(ctx, dst, src, bufSize, total, progress)
}
if err := ensureContext(ctx); err != nil {
return 0, err
}
if progress != nil && total > 0 {
progress(0)
}
wrappedDst := &sftpDownloadProgressWriter{
ctx: ctx,
writer: dst,
total: total,
progress: progress,
}
written, err := writerTo.WriteTo(wrappedDst)
if err != nil {
return written, err
}
if err := ensureContext(ctx); err != nil {
return written, err
}
reportProgress(progress, written, total)
return written, nil
}
func (r *sftpUploadProgressReader) Read(p []byte) (int, error) {
if err := ensureContext(r.ctx); err != nil {
return 0, err
}
r.mu.Lock()
defer r.mu.Unlock()
if err := ensureContext(r.ctx); err != nil {
return 0, err
}
n, err := r.reader.Read(p)
if n > 0 {
r.copied += int64(n)
reportQueuedTransferProgress(r.progress, r.copied, r.total)
}
return n, err
}
func (w *sftpDownloadProgressWriter) Write(p []byte) (int, error) {
if err := ensureContext(w.ctx); err != nil {
return 0, err
}
w.mu.Lock()
defer w.mu.Unlock()
if err := ensureContext(w.ctx); err != nil {
return 0, err
}
n, err := w.writer.Write(p)
if n > 0 {
w.copied += int64(n)
reportQueuedTransferProgress(w.progress, w.copied, w.total)
}
if err != nil {
return n, err
}
if n != len(p) {
return n, io.ErrShortWrite
}
if err := ensureContext(w.ctx); err != nil {
return n, err
}
return n, nil
}
func reportQueuedTransferProgress(progress func(float64), copied int64, total int64) {
if progress == nil || total <= 0 {
return
}
percent := float64(copied) / float64(total) * 100
if percent >= 100 {
percent = 99
}
progress(percent)
}
func reportProgress(progress func(float64), copied int64, total int64) {
if progress == nil {
return