clonetool/manager.go

package main

import (
	"bufio"
	"encoding/hex"
	"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
	StateDir          string
	Yes               bool
	Force             bool
	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
	}

	srcCtrl, err := startController(srcSpec, "source", &cfg)
	if err != nil {
		return fmt.Errorf("start source control agent: %w", err)
	}
	defer srcCtrl.Close()

	dstCtrl, err := startController(dstSpec, "dest", &cfg)
	if err != nil {
		return fmt.Errorf("start dest control agent: %w", err)
	}
	defer dstCtrl.Close()

	srcInfo, err := srcCtrl.Stat(srcSpec.Path)
	if err != nil {
		return fmt.Errorf("stat source: %w", err)
	}
	if !srcInfo.Exists {
		return fmt.Errorf("source %s does not exist", srcSpec)
	}

	dstInfo, err := dstCtrl.Stat(dstSpec.Path)
	if err != nil {
		return fmt.Errorf("stat destination: %w", err)
	}

	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
	statePath := statePathFor(cfg.StateDir, cfg.Job)
	state, err := LoadJobState(statePath, cfg.Source, cfg.Dest, cfg.BlockSize, cfg.Force)
	if err != nil {
		return err
	}
	state.Resize(uint64(blockCount))

	lastSave := time.Now()
	sinceSave := 0
	checkpoint := func(force bool) {
		sinceSave++
		if !force && sinceSave < 2000 && time.Since(lastSave) < 5*time.Second {
			return
		}
		if err := state.Save(statePath); err != nil {
			fmt.Fprintf(os.Stderr, "warning: could not checkpoint job state: %v\n", err)
			return
		}
		lastSave = time.Now()
		sinceSave = 0
	}

	cb := transferCallbacks{
		onProgress: func(m CtrlMsg) { printProgress(m) },
		onBlockDone: func(r BlockResult) {
			h, err := hex.DecodeString(r.Hash)
			if err != nil || len(h) != 32 {
				return
			}
			if r.Index < uint64(len(state.Hashes)) {
				copy(state.Hashes[r.Index][:], h)
			}
			checkpoint(false)
		},
	}

	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: cfg.RemoteBin, SSHBin: cfg.SSHBin, SSHOpts: cfg.SSHOpts, ConnectTimeoutSec: cfg.ConnectTimeoutSec,
	}
	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, state.Hashes, cb)
	if err != nil {
		return err
	}
	if !ok {
		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: cfg.RemoteBin, SSHBin: cfg.SSHBin, SSHOpts: cfg.SSHOpts, ConnectTimeoutSec: cfg.ConnectTimeoutSec,
		}
		ok2, reason2, err2 := dstCtrl.ConnectPull(pullReq, state.Hashes, 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)
		}
	}

	fmt.Fprintln(os.Stderr)
	state.Size = targetSize
	checkpoint(true)
	fmt.Fprintf(os.Stderr, "done: job %q, %d blocks, source=%s dest=%s\n", cfg.Job, blockCount, srcSpec, dstSpec)
	return nil
}

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"
}

func printProgress(m CtrlMsg) {
	total := m.TotalBlocks
	if total <= 0 {
		total = m.Copied + m.Skipped
	}
	fmt.Fprintf(os.Stderr, "\rblocks: %d/%d  copied=%d skipped=%d          ", m.Copied+m.Skipped, total, m.Copied, m.Skipped)
}