repo-sync/server_cmd.go

package main

import (
	"context"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	cfg "github.com/grmrgecko/repo-sync/config"
	"github.com/grmrgecko/repo-sync/fetch"
	"github.com/grmrgecko/repo-sync/server"
	"github.com/grmrgecko/repo-sync/state"
	log "github.com/sirupsen/logrus"
)

// shutdownTimeout bounds the graceful HTTP shutdown at exit.
const shutdownTimeout = 30 * time.Second

// ServerCmd runs the caching mirror server.
type ServerCmd struct {
	ConfigPath string `help:"The path to the config file." optional:"" type:"existingfile"`
	Bind       string `name:"http.bind" help:"Override the configured bind address." optional:""`
	Port       uint   `name:"http.port" help:"Override the configured port." optional:""`
}

// updateCFG applies command line overrides onto a freshly loaded
// configuration, before it is installed and visible to serving goroutines.
func (c *ServerCmd) updateCFG(conf *cfg.Config) {
	if c.Bind != "" {
		conf.HTTP.BindAddr = c.Bind
	}
	if c.Port != 0 {
		conf.HTTP.Port = c.Port
	}
}

// newHTTPServer builds the listener configuration from the active config.
func newHTTPServer() *http.Server {
	conf := cfg.C.Load()
	return &http.Server{
		Addr:              conf.HTTP.ListenAddr(),
		Handler:           server.Handler(),
		ReadHeaderTimeout: conf.HTTP.ReadHeaderTimeout,
		ReadTimeout:       conf.HTTP.ReadTimeout,
		WriteTimeout:      conf.HTTP.WriteTimeout,
		IdleTimeout:       conf.HTTP.IdleTimeout,
	}
}

// startHTTP binds the listener and serves in the background, returning the
// server for a later shutdown.
func startHTTP() (*http.Server, error) {
	server := newHTTPServer()
	ln, err := net.Listen("tcp", server.Addr)
	if err != nil {
		return nil, err
	}
	log.WithField("addr", server.Addr).Info("Mirror server listening.")
	go func() {
		if err := server.Serve(ln); err != nil && err != http.ErrServerClosed {
			log.WithError(err).Error("HTTP server failed.")
		}
	}()
	return server, nil
}

// stopHTTP gracefully shuts a server down within the shutdown timeout.
func stopHTTP(server *http.Server) {
	ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
	defer cancel()
	_ = server.Shutdown(ctx)
}

// Run starts the mirror server: configuration, state, crawl loops, and the
// HTTP listener, then handles signals until shutdown.
func (c *ServerCmd) Run() error {
	// Load configuration and apply command line overrides before install.
	conf, err := cfg.Load(c.ConfigPath)
	if err != nil {
		return err
	}
	c.updateCFG(conf)
	cfg.C.Store(conf)
	cfg.SetupLogging(conf.Log)
	if conf.Crawler.UserAgent != "" {
		fetch.SetUserAgent(conf.Crawler.UserAgent)
	}
	fetch.SetRequestTimeout(conf.Crawler.RequestTimeout)

	// Load persisted crawl state and build the shared HTTP client.
	if err := state.Load(); err != nil {
		return err
	}
	fetch.Reload()

	// Start the background loops and the listener.
	loopCtx, loopCancel := context.WithCancel(context.Background())
	defer loopCancel()
	go state.S.FlushLoop(loopCtx)
	crawlDone := make(chan struct{})
	go func() {
		defer close(crawlDone)
		server.CrawlLoop(loopCtx)
	}()

	// stopCrawls cancels the loops and waits for in-flight crawls so their
	// outcomes are recorded before the final state flush. The crawl loop
	// must return first: a scheduler pass in flight can still dispatch
	// crawls, and a WaitGroup Add racing Wait is a misuse panic.
	stopCrawls := func() {
		loopCancel()
		<-crawlDone
		server.WaitCrawls()
	}

	httpServer, err := startHTTP()
	if err != nil {
		return err
	}

	// Handle signals: reload on SIGHUP, shut down on SIGINT/SIGTERM.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
	for sig := range signals {
		if sig != syscall.SIGHUP {
			break
		}
		log.Info("Reloading configuration.")
		oldAddr := cfg.C.Load().HTTP.ListenAddr()
		conf, err := cfg.Load(c.ConfigPath)
		if err != nil {
			log.WithError(err).Error("Unable to reload configuration; keeping previous values.")
			continue
		}
		c.updateCFG(conf)
		cfg.C.Store(conf)
		cfg.SetupLogging(conf.Log)
		if conf.Crawler.UserAgent != "" {
			fetch.SetUserAgent(conf.Crawler.UserAgent)
		}
		fetch.SetRequestTimeout(conf.Crawler.RequestTimeout)
		fetch.Reload()
		if conf.HTTP.ListenAddr() != oldAddr {
			stopHTTP(httpServer)
			if httpServer, err = startHTTP(); err != nil {
				// Exiting on a failed rebind must still flush crawl state.
				stopCrawls()
				if saveErr := state.S.Save(); saveErr != nil {
					log.WithError(saveErr).Error("Failed to save state while shutting down.")
				}
				return err
			}
		}
	}

	// Shut down: stop the listener, stop the loops, then flush state.
	log.Info("Shutting down.")
	stopHTTP(httpServer)
	stopCrawls()
	return state.S.Save()
}