Gracefully shutdown server

Listen for SIGINT and SIGTERM signals and write the current state to the
key-value store before shutting down.

1 files changed, 33 insertions, 0 deletions

diff --git a/monzo_ynab/tokens.go b/monzo_ynab/tokens.go
index 4617595bf55a..d969ce6e43b4 100644
--- a/monzo_ynab/tokens.go
+++ b/monzo_ynab/tokens.go
@@ -17,6 +17,8 @@ import (
 	"os"
 	"time"
 	"kv"
+	"os/signal"
+	"syscall"
 )
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -118,6 +120,34 @@ func refreshTokens(refreshToken string) (string, string) {
 	return payload.AccessToken, payload.RefreshToken
 }
 
+// Listen for SIGINT and SIGTERM signals. When received, persist the access and
+// refresh tokens and shutdown the server.
+func handleInterrupts() {
+	// Gracefully handle interruptions.
+	sigs := make(chan os.Signal)
+	done := make(chan bool)
+
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		sig := <-sigs
+		log.Printf("Received signal to shutdown. %v\n", sig)
+		// Persist existing tokens
+		log.Println("Persisting existing credentials...")
+		msg := readMsg{make(chan state)}
+		chans.reads <- msg
+		state := <-msg.sender
+		kv.Set("monzoAccessToken", state.accessToken)
+		kv.Set("monzoRefreshToken", state.refreshToken)
+		log.Println("Credentials persisted.")
+		done <- true
+	}()
+
+	<-done
+	log.Println("Received signal to shutdown. Exiting...")
+	os.Exit(0)
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Main
 ////////////////////////////////////////////////////////////////////////////////
@@ -135,6 +165,9 @@ func main() {
 		log.Fatal("Cannot start server without access or refresh tokens.")
 	}
 
+	// Gracefully shutdown.
+	go handleInterrupts()
+
 	// Manage application state.
 	go func() {
 		state := &state{accessToken, refreshToken}