lobby/daemon/main.go

package main

import (
	"context"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/by-cx/lobby/common"
	"github.com/by-cx/lobby/nats_driver"
	"github.com/by-cx/lobby/redis_driver"
	"github.com/by-cx/lobby/server"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

var discoveryStorage server.Discoveries = server.Discoveries{}
var driver common.Driver
var localHost server.LocalHost

var config Config

var shuttingDown bool

func init() {
	// Load config from environment variables
	config = *GetConfig()

	// Setup discovery storage
	discoveryStorage.LogChannel = make(chan string)
	discoveryStorage.TTL = config.TTL

	// localhost initization
	localHost = server.LocalHost{
		LabelsPath:            config.LabelsPath,
		HostnameOverride:      config.HostName,
		InitialLabels:         config.Labels,
		RuntimeLabelsFilename: config.RuntimeLabelsFilename,
	}

	// Setup driver
	if config.Driver == "NATS" {
		driver = &nats_driver.Driver{
			NATSUrl:              config.NATSURL,
			NATSDiscoveryChannel: config.NATSDiscoveryChannel,

			LogChannel: discoveryStorage.LogChannel,
		}
	} else if config.Driver == "Redis" {
		driver = &redis_driver.Driver{
			Host:     config.RedisHost,
			Port:     uint(config.RedisPort),
			Password: config.RedisPassword,
			Channel:  config.RedisChannel,
			DB:       uint(config.RedisDB),

			LogChannel: discoveryStorage.LogChannel,
		}
	} else {
		log.Fatalf("unsupported driver %s", config.Driver)
	}

}

// cleanDiscoveryPool clears the local server map and keeps only the alive servers
func cleanDiscoveryPool() {
	for {
		discoveryStorage.Clean()
		time.Sleep(time.Duration(config.CleanEvery) * time.Second)
	}

}

// sendGoodbyePacket is almost same as sendDiscoveryPacket but it's not running in loop
// and it adds goodbye message so other nodes know this node is gonna die.
func sendGoodbyePacket() {
	discovery, err := localHost.GetIdentification()
	if err != nil {
		log.Printf("sending discovery identification error: %v\n", err)
	}

	err = driver.SendGoodbyePacket(discovery)
	if err != nil {
		log.Println(err)
	}
}

// sendDisoveryPacket sends discovery packet regularly so the network know we exist
func sendDiscoveryPacket() {
	for {
		discovery, err := localHost.GetIdentification()
		if err != nil {
			log.Printf("sending discovery identification error: %v\n", err)
		}

		err = driver.SendDiscoveryPacket(discovery)
		if err != nil {
			log.Println(err)
		}
		time.Sleep(time.Duration(config.KeepAlive) * time.Second)

		if shuttingDown {
			break
		}
	}
}

// Print logs acquired from disovery storage
func printDiscoveryLogs() {
	for {
		logMessage := <-discoveryStorage.LogChannel
		log.Println(logMessage)
	}
}

func main() {
	var err error

	// Closing the logging channel
	defer close(discoveryStorage.LogChannel)
	defer driver.Close()

	// ------------------------
	// Server discovering stuff
	// ------------------------

	// Connect to the NATS service
	driver.RegisterSubscribeFunction(func(d server.Discovery) {
		discoveryStorage.Add(d)
	})
	driver.RegisterUnsubscribeFunction(func(d server.Discovery) {
		discoveryStorage.Delete(d.Hostname)
	})

	err = driver.Init()
	if err != nil {
		log.Fatalln(err)
	}

	go printDiscoveryLogs()
	go cleanDiscoveryPool()

	// If config.Register is false this instance won't be registered with other nodes
	if config.Register {
		go sendDiscoveryPacket()
	} else {
		log.Println("standalone mode, I won't register myself")
	}

	// --------
	// REST API
	// --------
	e := echo.New()

	// Middleware
	if len(config.Token) > 0 {
		e.Use(TokenMiddleware)
	}
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())

	// Routes
	if config.DisableAPI {
		e.GET("/", listHandler)
		e.GET("/v1/discovery", getIdentificationHandler)
		e.GET("/v1/discoveries", listHandler)
		e.POST("/v1/labels", addLabelsHandler)
		e.DELETE("/v1/labels", deleteLabelsHandler)
		e.GET("/v1/prometheus/:name", prometheusHandler)
	}

	// ------------------------------
	// Termination signals processing
	// ------------------------------

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

	go func(e *echo.Echo, config Config) {
		sig := <-signals
		shuttingDown = true
		if config.Register {
			log.Printf("%s signal received, sending goodbye packet\n", sig.String())
			sendGoodbyePacket()
			time.Sleep(1 * time.Second) // we wait for a few seconds to let background jobs to finish their job
		} else {
			log.Printf("%s signal received", sig.String())
		}
		e.Shutdown(context.TODO())
	}(e, config)

	// Start server
	// In most cases this will end expectedly so it doesn't make sense to use the echo's approach to treat this message as an error.
	e.Logger.Info(e.Start(config.Host + ":" + strconv.Itoa(int(config.Port))))
}