orca/internal/sshpush/fanout.go

package sshpush

import (
	"context"
	"fmt"
	"os"
	"sync"

	"golang.org/x/sync/errgroup"

	"git.cloudinit.dev/coreci/orca/internal/emitter"
)

// DefaultFanoutConcurrency is the default bounded-concurrency limit for
// fan-out operations (I-B-001). The Transport.ExecAll and WriteAll
// methods use this when the caller does not override it.
const DefaultFanoutConcurrency = 8

// ExecAll runs cmd on all peers in parallel with bounded concurrency
// (default 8, I-B-001). Returns per-peer output and per-peer errors. A
// nil entry in the errors map means that peer succeeded; the output map
// contains that peer's stdout. The returned error is non-nil only if
// the fan-out itself failed (e.g., context cancelled before any peer
// ran); per-peer failures are in the errors map.
func (t *Transport) ExecAll(ctx context.Context, peers []string, cmd string) (map[string][]byte, map[string]error) {
	return t.ExecAllWithConcurrency(ctx, peers, cmd, DefaultFanoutConcurrency)
}

// ExecAllWithConcurrency is ExecAll with an explicit concurrency limit.
// A limit <= 0 uses DefaultFanoutConcurrency.
func (t *Transport) ExecAllWithConcurrency(ctx context.Context, peers []string, cmd string, concurrency int) (map[string][]byte, map[string]error) {
	if concurrency <= 0 {
		concurrency = DefaultFanoutConcurrency
	}
	out := make(map[string][]byte, len(peers))
	errs := make(map[string]error, len(peers))
	var mu sync.Mutex
	g, gctx := errgroup.WithContext(ctx)
	g.SetLimit(concurrency)
	for _, p := range peers {
		peer := p
		g.Go(func() error {
			o, err := t.Exec(gctx, peer, cmd)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs[peer] = err
				return nil // per-peer error; do not cancel the group
			}
			out[peer] = o
			return nil
		})
	}
	_ = g.Wait()
	return out, errs
}

// WriteAll writes the given files to each peer in parallel with bounded
// concurrency (default 8, I-B-001). The files map is keyed by peer; each
// peer's files are written sequentially (to preserve order and avoid
// intra-peer races on shared paths). Returns per-peer errors; a peer
// missing from the map or with a nil entry succeeded. The returned
// error is non-nil only if the fan-out itself failed (context cancelled).
func (t *Transport) WriteAll(ctx context.Context, peers []string, files map[string][]emitter.File) map[string]error {
	return t.WriteAllWithConcurrency(ctx, peers, files, DefaultFanoutConcurrency)
}

// WriteAllWithConcurrency is WriteAll with an explicit concurrency limit.
// A limit <= 0 uses DefaultFanoutConcurrency.
func (t *Transport) WriteAllWithConcurrency(ctx context.Context, peers []string, files map[string][]emitter.File, concurrency int) map[string]error {
	if concurrency <= 0 {
		concurrency = DefaultFanoutConcurrency
	}
	errs := make(map[string]error, len(peers))
	var mu sync.Mutex
	g, gctx := errgroup.WithContext(ctx)
	g.SetLimit(concurrency)
	for _, p := range peers {
		peer := p
		peerFiles := files[peer]
		g.Go(func() error {
			for _, f := range peerFiles {
				mode := parseMode(f.Mode)
				if _, err := t.WriteFileIdempotent(gctx, peer, f.Path, []byte(f.Content), mode); err != nil {
					mu.Lock()
					errs[peer] = fmt.Errorf("sshpush: write %s on %s: %w", f.Path, peer, err)
					mu.Unlock()
					return nil // per-peer error; do not cancel the group
				}
			}
			return nil
		})
	}
	_ = g.Wait()
	return errs
}

// parseMode parses an octal mode string like "0644" into an os.FileMode.
// Returns 0644 on parse failure (a safe default for non-executable
// config files).
func parseMode(s string) os.FileMode {
	var m uint32
	for _, r := range s {
		if r < '0' || r > '7' {
			return 0o644
		}
		m = m<<3 | uint32(r-'0')
	}
	if m == 0 {
		return 0o644
	}
	return os.FileMode(m)
}