move-data-nas/internal/syncengine/engine.go

package syncengine

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/syncserver/internal/config"
	"github.com/syncserver/internal/models"
	"github.com/syncserver/internal/wol"
)

type Engine struct {
	db           *sql.DB
	cfg          *config.Config
	queue        *Queue
	eventBus     *EventBus
	mu           sync.RWMutex
	stopped      bool
	lastProbeAt  atomic.Int64
}

type Event struct {
	Type      string
	JobID     int64
	MachineID int64
	Key       string
	Value     string
	Line      string
	Stream    string
}

func New(database interface{ SQLDB() *sql.DB }, cfg *config.Config) *Engine {
	e := &Engine{
		db:       database.SQLDB(),
		cfg:      cfg,
		queue:    NewQueue(),
		eventBus: NewEventBus(200),
	}
	return e
}

func (e *Engine) Start() {}
func (e *Engine) Stop()  {}

func (e *Engine) SubscribeJob(jobID int64) (chan Event, func()) {
	return e.eventBus.Subscribe(jobID)
}

func (e *Engine) SubscribeGlobal() (chan Event, func()) {
	return e.eventBus.SubscribeGlobal()
}

func (e *Engine) wakeMachine(ctx context.Context, m *models.Machine) {
	if m == nil || !m.WoLEnabled || m.MACAddress == nil {
		return
	}
	if wol.IsReachable(ctx, m.Host, m.Port, 3*time.Second) {
		return
	}
	mac, err := wol.ParseMAC(*m.MACAddress)
	if err != nil {
		return
	}
	bcast := ""
	if m.BroadcastAddr != nil {
		bcast = *m.BroadcastAddr
	}
	if err := wol.Send(mac, bcast); err != nil {
		slog.Warn("WoL send failed", "host", m.Host, "error", err)
		return
	}
	slog.Info("WoL magic packet sent (remote-to-remote)", "host", m.Host, "mac", *m.MACAddress)
}

func (e *Engine) Run(ctx context.Context, jobID int64, pairID int64) error {
	if e.queue.IsRunning(pairID) {
		existingJobID, _ := e.queue.GetJobID(pairID)
		return fmt.Errorf("%w: job %d is already running", ErrAlreadyRunning, existingJobID)
	}

	jobCtx, cancel := context.WithCancel(ctx)
	cancelledByUser := atomic.Bool{}
	wrappedCancel := func() {
		cancelledByUser.Store(true)
		cancel()
	}
	enqueueErr := e.queue.Enqueue(pairID, jobID, wrappedCancel)
	if enqueueErr != nil {
		return enqueueErr
	}
	defer e.queue.Dequeue(pairID)

	pairRepo := models.NewSyncPairRepository(e.db)
	pair, err := pairRepo.GetByID(pairID)
	if err != nil {
		return fmt.Errorf("fetching sync pair: %w", err)
	}

	machineRepo := models.NewMachineRepository(e.db)
	var srcMachine, dstMachine *models.Machine

	if pair.SourceMachineID != nil {
		m, _ := machineRepo.GetByID(*pair.SourceMachineID)
		srcMachine = m
	}
	if pair.DestMachineID != nil {
		m, _ := machineRepo.GetByID(*pair.DestMachineID)
		dstMachine = m
	}

	src := buildPath(pair.SourcePath, srcMachine)
	dst := buildPath(pair.DestPath, dstMachine)

	cfg := &SyncPairConfig{
		ID:              pair.ID,
		Name:            pair.Name,
		SourceMachineID: pair.SourceMachineID,
		SourcePath:      pair.SourcePath,
		DestMachineID:   pair.DestMachineID,
		DestPath:        pair.DestPath,
		Direction:       pair.Direction,
		RsyncFlags:      pair.RsyncFlags,
		ExcludePatterns: pair.ExcludePatternsList(),
	}
	if isRemoteToRemote(srcMachine, dstMachine) {
		cfg.Source = pair.SourcePath
		cfg.Dest = buildPath(pair.DestPath, dstMachine)
	} else {
		cfg.Source = src
		cfg.Dest = dst
	}

	logPath := filepath.Join(e.cfg.LogsDir(), fmt.Sprintf("%d.log", jobID))
	f, err := os.Create(logPath)
	if err != nil {
		slog.Warn("failed to create log file", "error", err)
	} else {
		f.Close()
	}

	var targetMachine *models.Machine
	var remotePort int
	if isRemoteToRemote(srcMachine, dstMachine) {
		targetMachine = srcMachine
		remotePort = srcMachine.Port
	} else if pair.Direction == "pull" && srcMachine != nil {
		targetMachine = srcMachine
		remotePort = srcMachine.Port
	} else if dstMachine != nil {
		targetMachine = dstMachine
		remotePort = dstMachine.Port
	}

	wolShouldRun := targetMachine != nil &&
		targetMachine.WoLEnabled &&
		targetMachine.MACAddress != nil

	if wolShouldRun {
		if wol.IsReachable(jobCtx, targetMachine.Host, remotePort, 3*time.Second) {
			slog.Info("machine already reachable, skipping WoL", "host", targetMachine.Host)
			if targetMachine.ID > 0 {
				e.setMachineStatus(targetMachine.ID, "online")
			}
		} else {
			e.setJobStatus(jobID, "waking_up")
			e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "waking_up"})

			mac, err := wol.ParseMAC(*targetMachine.MACAddress)
			if err == nil {
				bcast := ""
				if targetMachine.BroadcastAddr != nil {
					bcast = *targetMachine.BroadcastAddr
				}
				wolErr := wol.Send(mac, bcast)
				if wolErr != nil {
					slog.Warn("WoL send failed", "host", targetMachine.Host, "error", wolErr)
				} else {
					slog.Info("WoL magic packet sent", "host", targetMachine.Host, "mac", *targetMachine.MACAddress)
				}

				timeout := time.Duration(targetMachine.WakeTimeoutSeconds) * time.Second
				interval := time.Duration(targetMachine.WakeCheckIntervalSeconds) * time.Second
				if err := wol.WaitUntilReady(jobCtx, targetMachine.Host, remotePort, timeout, interval); err != nil {
					e.setJobStatus(jobID, "failed")
					if wolErr != nil {
						e.setJobError(jobID, "wol_send_failed", wolErr.Error())
						e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "failed", Line: fmt.Sprintf("WoL send failed: %v", wolErr)})
						return fmt.Errorf("WoL send failed: %w", wolErr)
					}
					if targetMachine.ID > 0 {
						e.setMachineStatus(targetMachine.ID, "offline")
					}
					e.setJobError(jobID, "wol_timeout", err.Error())
					e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "failed", Line: err.Error()})
					return fmt.Errorf("machine not ready: %w", err)
				}
				if targetMachine.ID > 0 {
					e.setMachineStatus(targetMachine.ID, "online")
				}
			}
		}
	}

	if isRemoteToRemote(srcMachine, dstMachine) && dstMachine != nil && dstMachine.WoLEnabled && dstMachine.MACAddress != nil {
		e.wakeMachine(jobCtx, dstMachine)
	}
	e.setJobStatus(jobID, "running")
	e.setJobLogFile(jobID, logPath)
	e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "running"})
	slog.Info("job started", "job_id", jobID, "pair", pair.Name)

	privKey, err := e.resolveSSHKey(targetMachine)
	if err != nil {
		slog.Warn("failed to resolve SSH key, using server key", "error", err)
		privKey = filepath.Join(e.cfg.SSHDir(), "id_ed25519")
	}

	runner := NewRsyncRunner(e.cfg.SSHDir(), privKey)
	logRepo := models.NewJobLogRepository(e.db)
	var outBuf, errBuf []string
	flush := func() {
		if len(outBuf) > 0 {
			logRepo.InsertBatch(jobID, "stdout", outBuf)
			for _, l := range outBuf {
				e.emit(Event{Type: "log", JobID: jobID, Stream: "stdout", Line: l})
			}
			outBuf = nil
		}
		if len(errBuf) > 0 {
			logRepo.InsertBatch(jobID, "stderr", errBuf)
			for _, l := range errBuf {
				e.emit(Event{Type: "log", JobID: jobID, Stream: "stderr", Line: l})
			}
			errBuf = nil
		}
	}

	onLine := func(stream, line string) {
		f, _ := os.OpenFile(logPath, os.O_APPEND|os.O_WRONLY, 0644)
		if f != nil {
			fmt.Fprintln(f, line)
			f.Close()
		}
		if stream == "stdout" {
			outBuf = append(outBuf, line)
		} else {
			errBuf = append(errBuf, line)
		}
		if len(outBuf) >= 50 || len(errBuf) >= 50 {
			flush()
		}
	}

	var result *RsyncResult
	if isRemoteToRemote(srcMachine, dstMachine) {
		result, err = runner.RunRemote(jobCtx, cfg, RemoteMachine{
			Host:    srcMachine.Host,
			Port:    srcMachine.Port,
			SSHUser: srcMachine.SSHUser,
			PrivKey: privKey,
		}, onLine)
	} else {
		result, err = runner.Run(jobCtx, cfg, onLine)
	}
	flush()

	if err != nil {
		if jobCtx.Err() != nil {
			code := "cancelled_shutdown"
			msg := "Job was cancelled due to server shutdown"
			if cancelledByUser.Load() {
				code = "cancelled_user"
				msg = "Job was cancelled by user"
			}
			e.setJobError(jobID, code, msg)
			e.setJobStatus(jobID, "cancelled")
			e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "cancelled", Line: msg})
			return jobCtx.Err()
		}
		e.setJobStatus(jobID, "failed")
		e.setJobError(jobID, "rsync_error", err.Error())
		e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "failed", Line: err.Error()})
		return fmt.Errorf("rsync error: %w", err)
	}

	if result.ExitCode != 0 {
		errMsg := result.Stderr
		errCode := "exit_code"
		if strings.Contains(errMsg, "cannot both be remote") {
			errCode = "remote_to_remote_unsupported"
			errMsg = "sync pair has both source and destination as remote machines; rsync requires one side to be the server (edit the sync pair to set one MachineID to null)"
		}
		e.setJobStatus(jobID, "failed")
		e.setJobError(jobID, errCode, errMsg)
		e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "failed", Line: errMsg})
		return fmt.Errorf("rsync exited with code %d: %s", result.ExitCode, result.Stderr)
	}

	e.setJobStatus(jobID, "success")
	e.emit(Event{Type: "status", JobID: jobID, Key: "status", Value: "success"})
	slog.Info("job completed", "job_id", jobID, "pair", pair.Name)
	e.persistAndClose(jobID)
	return nil
}

func (e *Engine) persistAndClose(jobID int64) {
	e.eventBus.CloseJobChannels(jobID)
}

func (e *Engine) resolveSSHKey(machine *models.Machine) (string, error) {
	if machine == nil || machine.SSHKeyID == nil {
		return filepath.Join(e.cfg.SSHDir(), "id_ed25519"), nil
	}
	sshKeyRepo := models.NewSSHKeyRepository(e.db)
	sshKey, err := sshKeyRepo.GetByID(*machine.SSHKeyID)
	if err != nil {
		return "", fmt.Errorf("fetching ssh key: %w", err)
	}
	return sshKey.PrivateKeyPath, nil
}

func (e *Engine) Cancel(jobID int64, syncPairID int64, byUser bool) bool {
	if e.queue.IsRunning(syncPairID) {
		e.queue.Cancel(syncPairID, byUser)
		return true
	}
	return false
}

func (e *Engine) setJobStatus(jobID int64, status string) {
	jobRepo := models.NewJobRepository(e.db)
	jobRepo.UpdateStatus(jobID, status)
}

func (e *Engine) setJobLogFile(jobID int64, path string) {
	jobRepo := models.NewJobRepository(e.db)
	jobRepo.SetLogFile(jobID, path)
}

func (e *Engine) setJobError(jobID int64, code, message string) {
	jobRepo := models.NewJobRepository(e.db)
	jobRepo.SetError(jobID, code, message)
}

func (e *Engine) setMachineStatus(machineID int64, status string) {
	machineRepo := models.NewMachineRepository(e.db)
	if err := machineRepo.UpdateStatus(machineID, status); err != nil {
		slog.Warn("failed to update machine status",
			"machine_id", machineID, "status", status, "error", err)
		return
	}
	e.emit(Event{
		Type:      "machine_status",
		MachineID: machineID,
		Key:       "status",
		Value:     status,
	})
}

func (e *Engine) emit(evt Event) {
	e.eventBus.Publish(evt)
}

func buildPath(path string, machine *models.Machine) string {
	if machine == nil {
		return path
	}
	return fmt.Sprintf("%s@%s:%s", machine.SSHUser, machine.Host, path)
}

func isRemoteToRemote(srcMachine, dstMachine *models.Machine) bool {
	return srcMachine != nil && dstMachine != nil
}

func (e *Engine) CreateJob(syncPairID int64, triggerType string) (int64, error) {
	jobRepo := models.NewJobRepository(e.db)
	id, err := jobRepo.Create(syncPairID, triggerType, "queued")
	if err != nil {
		return 0, err
	}
	return id, nil
}

const (
	probeThrottleSeconds = 10
	probeTimeout         = 1500 * time.Millisecond
	probeMaxConcurrent   = 20
)

func (e *Engine) ProbeAllMachines() {
	now := time.Now().UnixNano()
	last := e.lastProbeAt.Load()

	if now-last < int64(probeThrottleSeconds*time.Second) {
		slog.Debug("ProbeAllMachines: skipped (throttled)")
		return
	}
	if !e.lastProbeAt.CompareAndSwap(last, now) {
		return
	}

	machineRepo := models.NewMachineRepository(e.db)
	ms, err := machineRepo.GetAll()
	if err != nil {
		slog.Warn("ProbeAllMachines: list failed", "error", err)
		return
	}

	sem := make(chan struct{}, probeMaxConcurrent)
	var wg sync.WaitGroup

	for i := range ms {
		wg.Add(1)
		sem <- struct{}{}
		go func(m *models.Machine) {
			defer wg.Done()
			defer func() {
				if r := recover(); r != nil {
					slog.Error("probe goroutine panicked", "machine_id", m.ID, "panic", r)
				}
				<-sem
			}()

			ctx, cancel := context.WithTimeout(context.Background(), probeTimeout)
			defer cancel()

			status := "offline"
			if wol.IsReachable(ctx, m.Host, m.Port, probeTimeout) {
				status = "online"
			}
			e.setMachineStatus(m.ID, status)
		}(&ms[i])
	}
	wg.Wait()
}