mirror/merlin/merlin.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"text/tabwriter"
	"time"

	"git.csclub.uwaterloo.ca/public/merlin/common"
	"golang.org/x/sys/unix"
)

var (
	cfg            common.Config
	outLogger      *log.Logger
	errLogger      *log.Logger
	repoMap        map[string]*common.Repo
	repoIdx        int
	numJobsRunning int
)

func getAndRunCommand(conn net.Conn) {
	defer conn.Close()

	var buf bytes.Buffer
	_, err := io.Copy(&buf, conn)
	if err != nil {
		errLogger.Println(err.Error())
		return
	}

	command := buf.String()
	args := strings.Split(command, ":")
	respondAndLogErr := func(msg string) {
		outLogger.Println(msg)
		conn.Write([]byte(msg))
	}

	if args[0] == "status" {
		status := tabwriter.NewWriter(conn, 5, 5, 5, ' ', 0)
		fmt.Fprintf(status, "Repository\tLast Synced\tNext Expected Sync\n")

		// for time formating see https://pkg.go.dev/time#pkg-constants
		for name, repo := range repoMap {
			fmt.Fprintf(status, "%s\t%s\t%s\n",
				name,
				time.Unix(repo.State.LastAttemptStartTime, 0).Format(time.RFC1123),
				time.Unix(repo.State.LastAttemptRunTime+int64(repo.Frequency), 0).Format(time.RFC1123),
			)
		}

		status.Flush()
	} else if args[0] == "sync" {
		if len(args) != 2 {
			respondAndLogErr("Could not parse sync command, forced sync fails.")
			return
		}

		if repo, inMap := repoMap[args[1]]; inMap {
			outLogger.Println("Attempting to force sync of " + repo.Name)
			if repo.RunIfPossible() {
				conn.Write([]byte("Forced sync for " + repo.Name))
				numJobsRunning++
			} else {
				respondAndLogErr("Cannot force sync: " + repo.Name + ", already syncing.")
			}
		} else {
			respondAndLogErr(args[1] + " is not tracked so cannot sync")
		}
	} else {
		respondAndLogErr("Received unrecognized command: " + command)
	}
}

func unixSockListener(connChan chan net.Conn, stopLisChan chan struct{}) {
	// must remove cfg.SockPath otherwise get "bind: address already in use"
	if filepath.Ext(cfg.SockPath) != ".sock" {
		panic(fmt.Errorf("Socket file must end with .sock"))
	} else if pathInfo, _ := os.Stat(cfg.SockPath); pathInfo.IsDir() {
		panic(fmt.Errorf("Value specified for socket file is a directory"))
	} else if err := os.Remove(cfg.SockPath); err != nil {
		panic(err)
	}

	ear, err := net.Listen("unix", cfg.SockPath)
	if err != nil {
		panic(err)
	}

	go func() {
		for {
			// will exit when ear is closed
			conn, err := ear.Accept()
			if err != nil {
				errLogger.Println(err.Error())
				return
			}
			connChan <- conn
		}
	}()

	<-stopLisChan
	ear.Close()
}

// We use a round-robin strategy. It's not the most efficient, but it's simple
// (read: easy to understand) and guarantees each repo will eventually get a chance to run.
func runAsManyAsPossible() {
	repos := cfg.Repos
	startIdx := repoIdx

	for numJobsRunning < cfg.MaxJobs {
		repo := repos[repoIdx]
		if repo.RunIfPossible() {
			numJobsRunning++
		}
		repoIdx = (repoIdx + 1) % len(repos)
		if repoIdx == startIdx {
			// we've come full circle
			return
		}
	}
}

func main() {
	outLogger = common.OutLogger()
	errLogger = common.ErrLogger()

	// check that merlin is run as mirror user
	// check that mirror user has pid of 1001

	doneChan := make(chan common.Result)
	stopChan := make(chan struct{})
	connChan := make(chan net.Conn)
	stopLisChan := make(chan struct{})

	stopSig := make(chan os.Signal, 1)
	reloadSig := make(chan os.Signal, 1)
	signal.Notify(stopSig, syscall.SIGINT, syscall.SIGTERM)
	signal.Notify(reloadSig, syscall.SIGHUP)

	unix.Umask(002)

	numJobsRunning = 0

	loadConfig := func() {
		cfg = common.GetConfig(doneChan, stopChan)
		outLogger.Println("Loaded config:\n" + fmt.Sprintf("%+v\n", cfg))

		repoMap = make(map[string]*common.Repo)
		for _, repo := range cfg.Repos {
			repoMap[repo.Name] = repo
		}

		repoIdx = 0
		go unixSockListener(connChan, stopLisChan)
	}

	loadConfig()
	// IsRunning must be false otherwise repo will never sync
	for _, repo := range repoMap {
		repo.State.IsRunning = false
	}
	runAsManyAsPossible()

runLoop:
	for {
		select {
		case <-stopSig:
			close(stopChan)
			close(stopLisChan)
			break runLoop

		case <-reloadSig:
			stopLisChan <- struct{}{}
			loadConfig()

		case done := <-doneChan:
			repoMap[done.Name].SyncCompleted(done.Exit)
			numJobsRunning--

		case conn := <-connChan:
			getAndRunCommand(conn)

		case <-time.After(1 * time.Minute):
		}
		runAsManyAsPossible()
	}

	for {
		select {
		case done := <-doneChan:
			repoMap[done.Name].SyncCompleted(done.Exit)
			numJobsRunning--

		case <-time.After(1 * time.Second):
		}
		if numJobsRunning <= 0 {
			return
		}
	}
}