perf: SSH uploads fetch pieces ahead and downloads post them side by side

internal/agent/sshfs.go

@@ -22,6 +22,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unicode/utf8"
 
@@ -524,24 +525,53 @@ func (a *Agent) doSSHFSSend(ctx context.Context, env proto.TaskEnvelope) proto.T
 	if err := s.Start("cat -- " + shq(args.Path)); err != nil {
 		return proto.TaskResult{Error: err.Error()}
 	}
+	// Read in order from cat, posted to the server a few at a time.
 	sizeHdr := strconv.FormatInt(size, 10)
 	n := pieceCount(size, args.Chunk)
-	for i := 0; i < n; i++ {
+	pctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	var (
+		wg    sync.WaitGroup
+		once  sync.Once
+		first error
+	)
+	fail := func(err error) { once.Do(func() { first = err; cancel() }) }
+	sem := make(chan struct{}, xferWorkers)
+	for i := 0; i < n && pctx.Err() == nil; i++ {
 		buf := make([]byte, min(int64(args.Chunk), size-int64(i)*int64(args.Chunk)))
 		if _, err := io.ReadFull(out, buf); err != nil {
-			return proto.TaskResult{Error: "файл изменился во время чтения: " + err.Error()}
+			fail(errors.New("файл изменился во время чтения: " + err.Error()))
+			break
 		}
-		err := a.xferRequest(ctx, func(ctx context.Context) (*http.Request, error) {
-			req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.xferURL(args.Xfer, i), bytes.NewReader(buf))
-			if err == nil {
-				req.Header.Set("X-Size", sizeHdr)
-				req.Header.Set("Content-Type", "application/octet-stream")
+		select {
+		case sem <- struct{}{}:
+		case <-pctx.Done():
+		}
+		if pctx.Err() != nil {
+			break
+		}
+		wg.Add(1)
+		go func(i int, buf []byte) {
+			defer func() { <-sem; wg.Done() }()
+			err := a.xferRequest(pctx, func(ctx context.Context) (*http.Request, error) {
+				req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.xferURL(args.Xfer, i), bytes.NewReader(buf))
+				if err == nil {
+					req.Header.Set("X-Size", sizeHdr)
+					req.Header.Set("Content-Type", "application/octet-stream")
+				}
+				return req, err
+			}, func(*http.Response) error { return nil })
+			if err != nil {
+				fail(errors.New("передача прервана: " + err.Error()))
 			}
-			return req, err
-		}, func(*http.Response) error { return nil })
-		if err != nil {
-			return proto.TaskResult{Error: "передача прервана: " + err.Error()}
-		}
+		}(i, buf)
+	}
+	wg.Wait()
+	if first == nil && ctx.Err() != nil {
+		first = ctx.Err()
+	}
+	if first != nil {
+		return proto.TaskResult{Error: first.Error()}
 	}
 	// All of it is read; the session closes (deferred) without waiting on a
 	// cat that may have more to say.
@@ -572,12 +602,15 @@ func (a *Agent) doSSHFSPut(ctx context.Context, env proto.TaskEnvelope) proto.Ta
 	if err := s.Start(writeCommand(args.Path)); err != nil {
 		return proto.TaskResult{Error: err.Error()}
 	}
+	// Pieces come in a few ahead; stdin takes them in order. ack: how many
+	// are written, for the server to let go.
 	n := pieceCount(args.Size, args.Chunk)
-	for i := 0; i < n; i++ {
+	var written atomic.Int64
+	err = ordered(ctx, n, xferWorkers, func(ctx context.Context, i int) ([]byte, error) {
 		want := min(int64(args.Chunk), args.Size-int64(i)*int64(args.Chunk))
 		var data []byte
 		err := a.xferRequest(ctx, func(ctx context.Context) (*http.Request, error) {
-			return http.NewRequestWithContext(ctx, http.MethodGet, a.xferURL(args.Xfer, i)+"?ack="+strconv.Itoa(i), nil)
+			return http.NewRequestWithContext(ctx, http.MethodGet, a.xferURL(args.Xfer, i)+"?ack="+strconv.FormatInt(written.Load(), 10), nil)
 		}, func(resp *http.Response) error {
 			b, err := io.ReadAll(io.LimitReader(resp.Body, int64(args.Chunk)+1))
 			if err == nil && int64(len(b)) != want {
@@ -586,13 +619,15 @@ func (a *Agent) doSSHFSPut(ctx context.Context, env proto.TaskEnvelope) proto.Ta
 			data = b
 			return err
 		})
-		if err == nil {
-			_, err = in.Write(data)
-		}
-		if err != nil {
-			_ = s.Close()
-			return proto.TaskResult{Error: "загрузка прервана: " + err.Error()}
-		}
+		return data, err
+	}, func(i int, b []byte) error {
+		_, err := in.Write(b)
+		written.Store(int64(i + 1))
+		return err
+	})
+	if err != nil {
+		_ = s.Close()
+		return proto.TaskResult{Error: "загрузка прервана: " + err.Error()}
 	}
 	_ = in.Close()
 	if err := s.Wait(); err != nil {