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)))) }