diff --git a/internal/agent/fsstream.go b/internal/agent/fsstream.go index bdaeccc..d59f1ea 100644 --- a/internal/agent/fsstream.go +++ b/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 +} diff --git a/internal/agent/ordered_test.go b/internal/agent/ordered_test.go new file mode 100644 index 0000000..e807393 --- /dev/null +++ b/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) + } +} diff --git a/internal/agent/sshfs.go b/internal/agent/sshfs.go index be496ea..032e853 100644 --- a/internal/agent/sshfs.go +++ b/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 {