perf: several dragged items cross to another place as one tar stream; older agents get them one by one

2026-10-15 22:28:46 +03:00
1 parent 6916af4236
commit 7d1996dd7d
@@ -99,7 +99,24 @@ func treeSize(ctx context.Context, p string) int64 {
 // packTar writes src as a tar stream whose entries start at name. Owners are
 // left out: the numbers mean other people on another machine.
 func packTar(ctx context.Context, w io.Writer, src, name string, m *shipMeter) error {
+	return packTars(ctx, w, []shipItem{{src, name}}, m)
+}
+
+// shipItem is one thing sent: where it is and the name it goes under.
+type shipItem struct{ src, name string }
+
+// packTars writes the items as one tar stream, each under its name.
+func packTars(ctx context.Context, w io.Writer, items []shipItem, m *shipMeter) error {
 	tw := tar.NewWriter(w)
+	for _, it := range items {
+		if err := packInto(ctx, tw, it.src, it.name, m); err != nil {
+			return err
+		}
+	}
+	return tw.Close()
+}
+
+func packInto(ctx context.Context, tw *tar.Writer, src, name string, m *shipMeter) error {
 	err := filepath.WalkDir(src, func(p string, d fs.DirEntry, err error) error {
 		if err != nil {
 			return err
@@ -152,10 +169,7 @@ func packTar(ctx context.Context, w io.Writer, src, name string, m *shipMeter) e
 		m.files.Add(1)
 		return nil
 	})
-	if err != nil {
-		return err
-	}
-	return tw.Close()
+	return err
 }

 type countReader struct {
@@ -175,10 +189,24 @@ func (a *Agent) doShipOut(parent context.Context, env proto.TaskEnvelope) proto.
 	if err := json.Unmarshal(env.Args, &args); err != nil {
 		return proto.TaskResult{Error: err.Error()}
 	}
-	if err := shipName(args.Name); err != nil {
-		return proto.TaskResult{Error: err.Error()}
+	names := args.ShipNames()
+	for _, n := range names {
+		if err := shipName(n); err != nil {
+			return proto.TaskResult{Error: err.Error()}
+		}
+	}
+	if len(args.More) > 0 && args.SSH != nil {
+		return proto.TaskResult{Error: "по SSH-адресу - по одному"}
 	}
 	src := filepath.Clean(args.Src)
+	srcs := []string{args.Src}
+	for _, p := range args.More {
+		// The others: of the same folder only.
+		if path.Dir(path.Clean(p)) != path.Dir(path.Clean(args.Src)) {
+			return proto.TaskResult{Error: "несколько - только из одной папки"}
+		}
+		srcs = append(srcs, p)
+	}
 	var conn *sshfsConn
 	if args.Conn != "" {
 		// The files are on an SSH host this agent holds a connection to.
@@ -195,14 +223,20 @@ func (a *Agent) doShipOut(parent context.Context, env proto.TaskEnvelope) proto.
 	m := &shipMeter{}
 	m.note.Store("подсчёт")
 	go a.reportShip(ctx, args.Job, m, cancel)
-	if conn != nil {
-		m.total.Store(a.remoteSize(conn, src))
-	} else {
-		m.total.Store(treeSize(ctx, src))
+	var total int64
+	items := make([]shipItem, len(srcs))
+	for i, p := range srcs {
+		if conn != nil {
+			total += a.remoteSize(conn, path.Clean(p))
+		} else {
+			total += treeSize(ctx, filepath.Clean(p))
+		}
+		items[i] = shipItem{filepath.Clean(p), names[i]}
 	}
-	pack := func(w io.Writer) error { return packTar(ctx, w, src, args.Name, m) }
+	m.total.Store(total)
+	pack := func(w io.Writer) error { return packTars(ctx, w, items, m) }
 	if conn != nil {
-		pack = func(w io.Writer) error { return a.packRemote(ctx, w, conn, src, args.Name, m) }
+		pack = func(w io.Writer) error { return a.packRemote(ctx, w, conn, srcs, names, m) }
 	}

 	var err error
@@ -318,8 +352,11 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 	if err := json.Unmarshal(env.Args, &args); err != nil {
 		return proto.TaskResult{Error: err.Error()}
 	}
-	if err := shipName(args.Name); err != nil {
-		return proto.TaskResult{Error: err.Error()}
+	names := args.ShipNames()
+	for _, n := range names {
+		if err := shipName(n); err != nil {
+			return proto.TaskResult{Error: err.Error()}
+		}
 	}
 	dir := filepath.Clean(args.Dir)
 	var conn *sshfsConn
@@ -331,8 +368,10 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 		}
 		conn, dir = c, path.Clean(args.Dir)
 	} else {
-		if _, err := os.Lstat(filepath.Join(dir, args.Name)); err == nil && args.Conflict != "replace" && args.Conflict != "skip" {
-			return proto.TaskResult{Error: filepath.Join(dir, args.Name) + " уже существует"}
+		for _, n := range names {
+			if _, err := os.Lstat(filepath.Join(dir, n)); err == nil && args.Conflict != "replace" && args.Conflict != "skip" {
+				return proto.TaskResult{Error: filepath.Join(dir, n) + " уже существует"}
+			}
 		}
 		if err := os.MkdirAll(dir, 0o755); err != nil {
 			return proto.TaskResult{Error: err.Error()}
@@ -366,7 +405,7 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 		}
 	}()
 	if conn != nil {
-		skipped, err := a.unpackRemote(conn, pr, dir, args.Name, args.Conflict)
+		skipped, err := a.unpackRemote(conn, pr, dir, names, args.Conflict)
 		if err == nil {
 			_, _ = io.Copy(io.Discard, pr)
 		}
@@ -375,13 +414,17 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 			return proto.TaskResult{Error: err.Error()}
 		}
 		res := map[string]int{}
-		if skipped {
-			res["skipped"] = 1
+		if skipped > 0 {
+			res["skipped"] = skipped
 		}
 		data, _ := json.Marshal(res)
 		return proto.TaskResult{OK: true, Data: data}
 	}
-	u := &unpacker{dir: dir, name: args.Name, conflict: args.Conflict}
+	allowed := map[string]bool{}
+	for _, n := range names {
+		allowed[n] = true
+	}
+	u := &unpacker{dir: dir, names: allowed, conflict: args.Conflict}
 	if os.Geteuid() == 0 {
 		// New things belong to whoever owns the folder they land in.
 		if info, err := os.Stat(dir); err == nil {
@@ -412,7 +455,8 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 // the conflict rule for what is there already: folders merge, and files are
 // replaced or kept.
 type unpacker struct {
-	dir, name, conflict string
+	dir, conflict string
+	names         map[string]bool // the top entries the stream may carry
 	uid, gid            int
 	chown               bool
 	done, skipped       int
@@ -429,7 +473,8 @@ func (u *unpacker) unpack(tr *tar.Reader) error {
 			return err
 		}
 		clean := path.Clean(hdr.Name)
-		if clean != u.name && !strings.HasPrefix(clean, u.name+"/") {
+		top, _, _ := strings.Cut(clean, "/")
+		if !u.names[top] || clean == ".." || strings.HasPrefix(clean, "../") {
 			return fmt.Errorf("в архиве чужой путь %q", hdr.Name)
 		}
 		target := filepath.Join(u.dir, filepath.FromSlash(clean))