mirror/merlin/merlin.go

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"golang.org/x/sys/unix"

	"git.csclub.uwaterloo.ca/public/merlin/arthur"
	"git.csclub.uwaterloo.ca/public/merlin/config"
	"git.csclub.uwaterloo.ca/public/merlin/logger"
	"git.csclub.uwaterloo.ca/public/merlin/sync"
)

var DEFAULT_CONFIG_PATH = "merlin-config.ini"

func main() {

	// custom help message
	flag.Usage = func() {
		w := flag.CommandLine.Output()

		fmt.Fprintf(w, "USAGE: merlin [-h | --help] [--config=<config-path>]\n")
		flag.PrintDefaults()
	}

	// parse command args
	// if more are added can prob use a library
	configPath := flag.String("config", DEFAULT_CONFIG_PATH, "alternate config file")
	flag.Parse()

	// check the user that program is running under (should be mirror)
	if os.Getuid() == 0 {
		panic("Merlin should not be run as root")
	}

	// receives a Result struct when a repo is done/fails syncing
	doneChan := make(chan config.SyncResult)
	// closed when merlin is told to stop running
	stopChan := make(chan struct{})

	// receives a Conn when a client makes a connection to unix socket
	connChan := make(chan net.Conn)
	// closed or receives a signal to stop listening to the unix socket
	stopLisChan := make(chan struct{})

	// gets unblocked when SIGINT or SIGTERM is sent, will begin process of stopping the program
	stopSig := make(chan os.Signal, 1)
	signal.Notify(stopSig, syscall.SIGINT, syscall.SIGTERM)

	// gets unblocked with SIGHUP is sent, will attempt to reload the config
	reloadSig := make(chan os.Signal, 1)
	signal.Notify(reloadSig, syscall.SIGHUP)

	unix.Umask(002)
	repoIdx := 0
	numJobsRunning := 0

	loadConfig := func() {
		config.LoadConfig(*configPath, doneChan, stopChan)
		logger.OutLog("Loaded config:\n" + fmt.Sprintf("%+v\n", config.Conf))

		// reset the round-robin index and recreate the socket listener
		repoIdx = 0
		go arthur.StartListener(connChan, stopLisChan)
	}

	// 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.
	runAsManyAsPossible := func() {
		startIdx := repoIdx

		for numJobsRunning < config.Conf.MaxJobs {
			if sync.SyncIfPossible(config.Repos[repoIdx], false) {
				numJobsRunning++
			}
			repoIdx = (repoIdx + 1) % len(config.Repos)
			if repoIdx == startIdx {
				// we've tried to run every repo and have come full circle
				return
			}
		}
	}

	loadConfig()
	// ensure that IsRunning is false otherwise repo will never sync.
	// only on startup can we assume that repos were not previously syncing
	// since reloading the config does not stop repos with a sync in progress
	for _, repo := range config.Repos {
		repo.State.IsRunning = false
	}
	runAsManyAsPossible()

runLoop:
	for {
		select {
		case <-stopSig: // caught a SIGINT or SIGTERM
			// kill all syncing repos and the socket listener
			close(stopChan)
			close(stopLisChan)
			break runLoop

		case <-reloadSig: // caught a SIGHUP
			// temporary stop the socket listener and load the config again
			stopLisChan <- struct{}{}
			loadConfig()

		case done := <-doneChan: // a sync is done and sends its exit status
			// repo could be removed from config while sync in progress
			if repo, check := config.RepoMap[done.Name]; check {
				sync.SyncCompleted(repo, done.Exit)
			}
			numJobsRunning--

		case conn := <-connChan: // a connection was accepted by the socket listener
			command, repoName := arthur.GetCommand(conn)
			switch command {
			case "status":
				arthur.SendStatus(conn)
			case "sync":
				if arthur.ForceSync(conn, repoName) {
					numJobsRunning++
				}
			default:
				arthur.SendAndLog(conn, "Received unrecognized command: "+command)
			}
			// close the received connection so that the message is sent
			conn.Close()

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

	// allow some time for jobs to terminate before force exiting the program
	go func() {
		<-time.After(time.Minute)
		logger.ErrLog("One minute has passed, forcefully exiting the program")
		os.Exit(1)
	}()

	// wait on repos to get terminated
	for {
		select {
		case done := <-doneChan:
			if repo, check := config.RepoMap[done.Name]; check {
				sync.SyncCompleted(repo, done.Exit)
			}
			numJobsRunning--

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

		if numJobsRunning <= 0 {
			os.Exit(0)
		}
	}
}