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