clonetool/manager.go

package main

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"
)

const defaultBlockSize = 4 * 1024 * 1024
const defaultConnectTimeoutSec = 8

// SyncConfig holds everything the `sync` command needs. It is also handed
// (the transport-relevant fields of it) to control agents as part of
// connect_push/connect_pull requests, so they know how to reach the peer.
type SyncConfig struct {
	Job               string
	Source            string
	Dest              string
	BlockSize         int64
	Yes               bool
	Sudo              string // auto | always | never
	Deploy            bool   // copy this binary to remote hosts that lack it
	ConnectTimeoutSec int
	SSHBin            string
	SSHOpts           []string
	RemoteBin         string
	ManagerHost       string
}

func runSync(cfg SyncConfig) error {
	srcSpec, err := parseSpec(cfg.Source)
	if err != nil {
		return fmt.Errorf("--source: %w", err)
	}
	dstSpec, err := parseSpec(cfg.Dest)
	if err != nil {
		return fmt.Errorf("--dest: %w", err)
	}
	if err := checkNotSame(srcSpec, dstSpec); err != nil {
		return err
	}

	// Make sure each remote endpoint has a runnable clonetool, copying this
	// binary over if not (unless --deploy=false). A host that appears on
	// both sides is only probed once.
	srcRemoteBin, dstRemoteBin := cfg.RemoteBin, cfg.RemoteBin
	if !srcSpec.IsLocal() {
		if srcRemoteBin, err = resolveRemoteBin(&cfg, srcSpec, "source"); err != nil {
			return err
		}
	}
	if !dstSpec.IsLocal() {
		if !srcSpec.IsLocal() && sameHost(srcSpec, dstSpec) {
			dstRemoteBin = srcRemoteBin
		} else if dstRemoteBin, err = resolveRemoteBin(&cfg, dstSpec, "dest"); err != nil {
			return err
		}
	}

	srcCtrl, srcInfo, srcSudo, err := bringUpController(srcSpec, "source", &cfg, srcRemoteBin, srcSpec.Path)
	if err != nil {
		return fmt.Errorf("source: %w", err)
	}
	defer srcCtrl.Close()
	if !srcInfo.Exists {
		return fmt.Errorf("source %s does not exist", srcSpec)
	}

	dstCtrl, dstInfo, dstSudo, err := bringUpController(dstSpec, "dest", &cfg, dstRemoteBin, dstSpec.Path)
	if err != nil {
		return fmt.Errorf("dest: %w", err)
	}
	defer dstCtrl.Close()

	targetSize, err := computeTargetSize(srcSpec, srcInfo, dstSpec, dstInfo, cfg.Yes)
	if err != nil {
		return err
	}

	if err := dstCtrl.Prepare(dstSpec.Path, targetSize); err != nil {
		return fmt.Errorf("prepare destination: %w", err)
	}

	blockCount := (targetSize + cfg.BlockSize - 1) / cfg.BlockSize

	pp := newProgressPrinter(cfg.BlockSize)
	cb := transferCallbacks{onProgress: pp.print}

	bothLocal := srcSpec.IsLocal() && dstSpec.IsLocal()
	srcHost, srcUser := resolveConnectHost(srcSpec, &cfg)
	dstHost, dstUser := resolveConnectHost(dstSpec, &cfg)

	pushReq := CtrlMsg{
		Path: srcSpec.Path, Size: targetSize, BlockSize: cfg.BlockSize,
		PeerHost: dstHost, PeerUser: dstUser, PeerPath: dstSpec.Path, PeerLocal: bothLocal,
		RemoteBin: dstRemoteBin, SSHBin: cfg.SSHBin, SSHOpts: cfg.SSHOpts, ConnectTimeoutSec: cfg.ConnectTimeoutSec,
		Sudo: dstSudo,
	}
	if bothLocal {
		fmt.Fprintf(os.Stderr, "both source and dest are local; syncing directly (no ssh) ...\n")
	} else {
		fmt.Fprintf(os.Stderr, "attempting push %s -> %s ...\n", srcSpec, dstSpec)
	}
	ok, reason, err := srcCtrl.ConnectPush(pushReq, cb)
	if err != nil {
		return err
	}
	if !ok {
		pp.finish()
		fmt.Fprintf(os.Stderr, "push not possible (%s); trying pull %s <- %s ...\n", reason, dstSpec, srcSpec)
		pullReq := CtrlMsg{
			Path: dstSpec.Path, Size: targetSize, BlockSize: cfg.BlockSize,
			PeerHost: srcHost, PeerUser: srcUser, PeerPath: srcSpec.Path, PeerLocal: bothLocal,
			RemoteBin: srcRemoteBin, SSHBin: cfg.SSHBin, SSHOpts: cfg.SSHOpts, ConnectTimeoutSec: cfg.ConnectTimeoutSec,
			Sudo: srcSudo,
		}
		ok2, reason2, err2 := dstCtrl.ConnectPull(pullReq, cb)
		if err2 != nil {
			return err2
		}
		if !ok2 {
			return fmt.Errorf(
				"could not establish a direct connection in either direction (push: %s; pull: %s); "+
					"run the manager on the source or destination host, or set up SSH connectivity in at least one direction",
				reason, reason2)
		}
	}

	pp.finish()
	label := cfg.Job
	if label == "" {
		label = fmt.Sprintf("%s -> %s", srcSpec, dstSpec)
	}
	fmt.Fprintf(os.Stderr, "done: %s, %d blocks\n", label, blockCount)
	return nil
}

// bringUpController starts a control agent for spec and does its initial
// stat. With --sudo=always the agent is elevated from the start; with
// --sudo=auto a permission error on the stat triggers one transparent
// restart under sudo. The returned bool reports whether the agent (and any
// peer helper it later spawns for this side) is running elevated.
func bringUpController(spec Spec, tag string, cfg *SyncConfig, remoteBin, probePath string) (*Controller, PathInfo, bool, error) {
	sudo := cfg.Sudo == "always" && canElevate()
	c, err := startController(spec, tag, cfg, remoteBin, sudo)
	if err != nil {
		return nil, PathInfo{}, sudo, err
	}
	info, err := c.Stat(probePath)
	if err == nil {
		return c, info, sudo, nil
	}
	if errors.Is(err, errNeedPriv) && cfg.Sudo == "auto" && !sudo {
		fmt.Fprintf(os.Stderr, "%s: permission denied on %s; retrying via sudo ...\n", tag, probePath)
		c.Close()
		sudo = true
		if c, err = startController(spec, tag, cfg, remoteBin, true); err != nil {
			return nil, PathInfo{}, sudo, err
		}
		if info, err = c.Stat(probePath); err == nil {
			return c, info, sudo, nil
		}
	}
	c.Close()
	if sudo {
		return nil, PathInfo{}, sudo, fmt.Errorf(
			"%w; if this is a sudo failure, configure passwordless sudo for clonetool on %s or run the manager as root",
			err, hostLabel(spec))
	}
	return nil, PathInfo{}, sudo, err
}

func sameHost(a, b Spec) bool {
	return strings.EqualFold(a.Host, b.Host) && a.User == b.User
}

func hostLabel(spec Spec) string {
	if spec.IsLocal() {
		return "this machine"
	}
	return spec.Host
}

func resolveConnectHost(spec Spec, cfg *SyncConfig) (host, user string) {
	if !spec.IsLocal() {
		return spec.Host, spec.User
	}
	host = cfg.ManagerHost
	if host == "" {
		if h, err := os.Hostname(); err == nil {
			host = h
		}
	}
	return host, ""
}

// computeTargetSize applies the sizing/truncation rules: a block-device
// destination can't grow, so the source must fit inside it (sync exactly
// min(src,dst), leaving the remainder of the device untouched); a
// regular-file destination is truncated to the source's size, shrinking
// with confirmation if it currently holds more data than that.
func computeTargetSize(srcSpec Spec, srcInfo PathInfo, dstSpec Spec, dstInfo PathInfo, yes bool) (int64, error) {
	if dstInfo.Exists && dstInfo.IsDevice {
		if srcInfo.Size > dstInfo.Size {
			return 0, fmt.Errorf("source %s (%s) is larger than destination device %s (%s); a device can't be grown",
				srcSpec, humanBytes(srcInfo.Size), dstSpec, humanBytes(dstInfo.Size))
		}
		return srcInfo.Size, nil
	}
	if dstInfo.Exists && dstInfo.Size > srcInfo.Size {
		if !yes {
			if !confirmShrink(dstSpec, dstInfo.Size, srcInfo.Size) {
				return 0, fmt.Errorf("aborted: destination %s would be truncated from %s to %s", dstSpec, humanBytes(dstInfo.Size), humanBytes(srcInfo.Size))
			}
		}
	}
	return srcInfo.Size, nil
}

func confirmShrink(dstSpec Spec, oldSize, newSize int64) bool {
	fmt.Fprintf(os.Stderr, "destination %s is %s; it will be truncated to %s to match the source. Continue? [y/N] ",
		dstSpec, humanBytes(oldSize), humanBytes(newSize))
	reader := bufio.NewReader(os.Stdin)
	line, _ := reader.ReadString('\n')
	line = strings.TrimSpace(strings.ToLower(line))
	return line == "y" || line == "yes"
}

// rateTracker smooths a monotonically increasing block counter into a
// blocks/sec rate, sampling no more than a few times a second so a burst of
// same-tick progress messages doesn't produce a noisy instantaneous rate.
type rateTracker struct {
	lastT   time.Time
	lastVal int64
	rate    float64
}

func (r *rateTracker) update(val int64) {
	now := time.Now()
	if r.lastT.IsZero() {
		r.lastT, r.lastVal = now, val
		return
	}
	dt := now.Sub(r.lastT).Seconds()
	if dt < 0.4 {
		return
	}
	inst := float64(val-r.lastVal) / dt
	if r.rate == 0 {
		r.rate = inst
	} else {
		r.rate = 0.6*r.rate + 0.4*inst
	}
	r.lastT, r.lastVal = now, val
}

func (r *rateTracker) mib(blockSize int64) float64 {
	return r.rate * float64(blockSize) / (1024 * 1024)
}

// progressPrinter renders a single rewriting status line on stderr. Scan,
// source read and destination write all run concurrently, each tracked by
// its own phase-tagged progress message (see CtrlMsg.Phase), so it keeps the
// latest of each and shows them on one line with their own rolling
// throughput; the "scan H/M" segment and its read rate disappear once the
// destination fingerprint is complete.
type progressPrinter struct {
	blockSize int64
	active    bool

	hashed  int64
	copied  int64
	skipped int64
	srcRead int64
	written int64
	total   int64

	dstReadRate  rateTracker // destination scan/hash throughput
	srcReadRate  rateTracker // source read/compare throughput
	dstWriteRate rateTracker // destination write throughput
}

func newProgressPrinter(blockSize int64) *progressPrinter {
	return &progressPrinter{blockSize: blockSize}
}

func (p *progressPrinter) print(m CtrlMsg) {
	p.active = true
	if m.TotalBlocks > 0 {
		p.total = m.TotalBlocks
	}
	switch m.Phase {
	case "scan":
		if m.Hashed > p.hashed {
			p.hashed = m.Hashed
		}
		if p.total > 0 && p.hashed >= p.total {
			p.dstReadRate.rate = 0 // scan finished; stop showing a frozen rate
		} else {
			p.dstReadRate.update(p.hashed)
		}
	case "xfer":
		p.copied, p.skipped, p.srcRead = m.Copied, m.Skipped, m.SrcRead
		p.srcReadRate.update(p.srcRead)
	case "write":
		p.written = m.Written
		p.dstWriteRate.update(p.written)
	}

	done := p.copied + p.skipped
	total := p.total
	if total <= 0 {
		total = done
	}

	scan := ""
	if p.total > 0 && p.hashed < p.total {
		scan = fmt.Sprintf("  scan %s/%s", formatCount(p.hashed), formatCount(p.total))
	}
	fmt.Fprintf(os.Stderr, "\r  syncing %s %s/%s  copied=%s skipped=%s%s  rd(src) %5.1f  rd(dst) %5.1f  wr(dst) %5.1f MiB/s        ",
		progressBar(done, total), formatCount(done), formatCount(total),
		formatCount(p.copied), formatCount(p.skipped), scan,
		p.srcReadRate.mib(p.blockSize), p.dstReadRate.mib(p.blockSize), p.dstWriteRate.mib(p.blockSize))
}

// finish ends the current status line with a newline so following output
// (and the shell prompt) starts clean.
func (p *progressPrinter) finish() {
	if p.active {
		fmt.Fprintln(os.Stderr)
		p.active = false
	}
}

func progressBar(done, total int64) string {
	const w = 22
	pct := 0
	if total > 0 {
		pct = int(done * 100 / total)
	}
	if pct > 100 {
		pct = 100
	}
	n := pct * w / 100
	return fmt.Sprintf("[%s%s] %3d%%", strings.Repeat("=", n), strings.Repeat(" ", w-n), pct)
}

func formatCount(n int64) string {
	switch {
	case n >= 1_000_000:
		return fmt.Sprintf("%.1fM", float64(n)/1e6)
	case n >= 10_000:
		return fmt.Sprintf("%.0fk", float64(n)/1e3)
	default:
		return fmt.Sprintf("%d", n)
	}
}