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

internal/agent/fsstream.go

@@ -289,3 +289,49 @@ func (a *Agent) doFSSend(ctx context.Context, env proto.TaskEnvelope) proto.Task
 	data, _ := json.Marshal(map[string]any{"path": args.Path, "size": size})
 	return proto.TaskResult{OK: true, Data: data}
 }
+
+// ordered fetches pieces 0..n-1 with up to depth in flight and hands them to
+// use one by one in their order: a stream that must be written in sequence
+// (stdin over SSH) still gets its pieces fetched ahead.
+func ordered(ctx context.Context, n, depth int, fetch func(ctx context.Context, i int) ([]byte, error), use func(i int, b []byte) error) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	type got struct {
+		b   []byte
+		err error
+	}
+	slots := make([]chan got, n)
+	for i := range slots {
+		slots[i] = make(chan got, 1)
+	}
+	sem := make(chan struct{}, depth)
+	go func() {
+		for i := 0; i < n; i++ {
+			select {
+			case sem <- struct{}{}:
+			case <-ctx.Done():
+				return
+			}
+			go func(i int) {
+				b, err := fetch(ctx, i)
+				slots[i] <- got{b, err}
+			}(i)
+		}
+	}()
+	for i := 0; i < n; i++ {
+		var g got
+		select {
+		case g = <-slots[i]:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+		<-sem
+		if g.err != nil {
+			return g.err
+		}
+		if err := use(i, g.b); err != nil {
+			return err
+		}
+	}
+	return nil
+}

internal/agent/ordered_test.go

@@ -0,0 +1,45 @@
+package agent
+
+import (
+	"context"
+	"math/rand"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+// Pieces fetched side by side are used in their order, never more than
+// depth at once in flight.
+func TestOrderedKeepsOrder(t *testing.T) {
+	var inFlight, most int32
+	var got []int
+	err := ordered(context.Background(), 20, 3, func(ctx context.Context, i int) ([]byte, error) {
+		n := atomic.AddInt32(&inFlight, 1)
+		for {
+			m := atomic.LoadInt32(&most)
+			if n <= m || atomic.CompareAndSwapInt32(&most, m, n) {
+				break
+			}
+		}
+		time.Sleep(time.Duration(rand.Intn(5)) * time.Millisecond)
+		atomic.AddInt32(&inFlight, -1)
+		return []byte{byte(i)}, nil
+	}, func(i int, b []byte) error {
+		if int(b[0]) != i {
+			t.Errorf("piece %d carries %d", i, b[0])
+		}
+		got = append(got, i)
+		return nil
+	})
+	if err != nil || len(got) != 20 {
+		t.Fatalf("got %v, %v", got, err)
+	}
+	for i, v := range got {
+		if v != i {
+			t.Fatalf("order: %v", got)
+		}
+	}
+	if most > 3 || most < 2 {
+		t.Fatalf("at most %d in flight", most)
+	}
+}

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) {
+		select {
-			req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.xferURL(args.Xfer, i), bytes.NewReader(buf))
+		case sem <- struct{}{}:
-			if err == nil {
+		case <-pctx.Done():
-				req.Header.Set("X-Size", sizeHdr)
+		}
-				req.Header.Set("Content-Type", "application/octet-stream")
+		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
+		}(i, buf)
-		}, func(*http.Response) error { return nil })
+	}
-		if err != nil {
+	wg.Wait()
-			return proto.TaskResult{Error: "передача прервана: " + err.Error()}
+	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 {
+		return data, err
-			_, err = in.Write(data)
+	}, func(i int, b []byte) error {
-		}
+		_, err := in.Write(b)
-		if err != nil {
+		written.Store(int64(i + 1))
-			_ = s.Close()
+		return err
-			return proto.TaskResult{Error: "загрузка прервана: " + err.Error()}
+	})
-		}
+	if err != nil {
+		_ = s.Close()
+		return proto.TaskResult{Error: "загрузка прервана: " + err.Error()}
 	}
 	_ = in.Close()
 	if err := s.Wait(); err != nil {