diff options
| -rw-r--r-- | src/main.go | 53 |
1 files changed, 37 insertions, 16 deletions
diff --git a/src/main.go b/src/main.go index a95834a..43fd31c 100644 --- a/src/main.go +++ b/src/main.go @@ -3,10 +3,11 @@ package main import ( "fmt" "log" - "net/http" + "os" + "os/signal" + "syscall" "github.com/gliderlabs/ssh" - "github.com/gorilla/mux" ) var ( @@ -16,6 +17,7 @@ var ( ) func main() { + // create a map mapping a city to an amount of hits metricsCityNum = make(map[string]int) @@ -25,8 +27,14 @@ func main() { // parse flags building a config struct config := parseFlags() + // if a state exists, load it + err := loadState(config) + if err != nil { + log.Printf("Could not load the state: %s", err) + } + // start the ssh server - log.Println("Starting SSH listener") + log.Printf("Starting SSH listener on port %d", config.sshPort) go func() { sshPortString := fmt.Sprintf(":%d", config.sshPort) listenErr := ssh.ListenAndServe(sshPortString, nil, ssh.PasswordAuth(handlePass)) @@ -35,20 +43,33 @@ func main() { } }() - // start the http server logging the metrics - log.Println("Starting HTTP metrics listener") + // Start the exit handler handing SIGTERM syscalls + // The handler writes the state (the amount of hits processed until the + // syscall) to the statefile + exitHandler(config) - r := mux.NewRouter() - r.HandleFunc("/", indexHandler) - r.HandleFunc("/metrics", metricsHandler) - r.HandleFunc("/locations", locationHandlerEndpoint) + // start the http server exposing the metrics + // this is blocking + setupHTTPServer(config) +} - // start the http server exposing the metrics and the locations - httpPortString := fmt.Sprintf(":%d", config.httpPort) - listenErr := http.ListenAndServe(httpPortString, r) +func exitHandler(config config) { - // handle potential errors - if listenErr != nil { - log.Fatalln(listenErr.Error()) - } + // relay incoming signals to the signalChannel + signalChannel := make(chan os.Signal, 2) + signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM) + + // wait until a signal is given, then write the state to a file, THEN exit + go func() { + <-signalChannel + fmt.Println("Ending process... Saving the state") + + // write the state to a file + err := WriteStateToFile(config) + if err != nil { + log.Printf("Error writing the stateFile: %s", err) + } + + os.Exit(0) + }() } |
