Add support to run weep as a service (fixes #28) (#29)

Netflix · Dec 11, 2020 · 63b0e6c · 63b0e6c
1 parent 26efd22
commit 63b0e6c
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 13 deletions.
diff --git a/cmd/ecs_credential_provider.go b/cmd/ecs_credential_provider.go
@@ -21,8 +21,6 @@ import (
 	"net"
 	"net/http"
 	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/spf13/viper"
 
@@ -65,10 +63,10 @@ func runEcsMetadata(cmd *cobra.Command, args []string) error {
 	}()
 
 	// Check for interrupt signal and exit cleanly
-	quit := make(chan os.Signal, 1)
-	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
-	<-quit
+	<-shutdown
 	log.Print("Shutdown signal received, exiting weep...")
+	// Send a signal to show that we're done shutting down
+	done <- 0
 
 	return nil
 }
diff --git a/cmd/metadata.go b/cmd/metadata.go
@@ -21,8 +21,6 @@ import (
 	"net"
 	"net/http"
 	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/netflix/weep/cache"
 
@@ -86,10 +84,10 @@ func runMetadata(cmd *cobra.Command, args []string) error {
 	}()
 
 	// Check for interrupt signal and exit cleanly
-	quit := make(chan os.Signal, 1)
-	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
-	<-quit
+	<-shutdown
 	log.Print("Shutdown signal received, exiting weep meta-data service...")
+	// Send a signal to show that we're done shutting down
+	done <- 0
 
 	return nil
 }
diff --git a/cmd/root.go b/cmd/root.go
@@ -41,6 +41,8 @@ var (
 		Long:              "Weep is a CLI tool that manages AWS access via ConsoleMe for local development.",
 		DisableAutoGenTag: true,
 	}
+	shutdown chan os.Signal
+	done     chan int
 )
 
 func init() {
@@ -54,7 +56,9 @@ func init() {
 
 }
 
-func Execute() {
+func Execute(shutdownSignal chan os.Signal, doneSignal chan int) {
+	shutdown = shutdownSignal
+	done = doneSignal
 	if err := rootCmd.Execute(); err != nil {
 		_, _ = fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)

diff --git a/go.mod b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/fsnotify/fsnotify v1.4.9 // indirect
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
 	github.com/gorilla/mux v1.7.4
+	github.com/kardianos/service v1.2.0
 	github.com/magiconair/properties v1.8.4 // indirect
 	github.com/manifoldco/promptui v0.8.0
 	github.com/markbates/pkger v0.17.1

diff --git a/go.sum b/go.sum
@@ -191,6 +191,8 @@ github.com/juju/ansiterm v0.0.0-20180109212912-720a0952cc2a h1:FaWFmfWdAUKbSCtOU
 github.com/juju/ansiterm v0.0.0-20180109212912-720a0952cc2a/go.mod h1:UJSiEoRfvx3hP73CvoARgeLjaIOjybY9vj8PUPPFGeU=
 github.com/julienschmidt/httprouter v1.2.0 h1:TDTW5Yz1mjftljbcKqRcrYhd4XeOoI98t+9HbQbYf7g=
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
+github.com/kardianos/service v1.2.0 h1:bGuZ/epo3vrt8IPC7mnKQolqFeYJb7Cs8Rk4PSOBB/g=
+github.com/kardianos/service v1.2.0/go.mod h1:CIMRFEJVL+0DS1a3Nx06NaMn4Dz63Ng6O7dl0qH0zVM=
 github.com/kisielk/errcheck v1.1.0 h1:ZqfnKyx9KGpRcW04j5nnPDgRgoXUeLh2YFBeFzphcA0=
 github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvWXihfKN4Q=
 github.com/kisielk/gotool v1.0.0 h1:AV2c/EiW3KqPNT9ZKl07ehoAGi4C5/01Cfbblndcapg=
@@ -432,6 +434,7 @@ golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20201015000850-e3ed0017c211/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201029080932-201ba4db2418 h1:HlFl4V6pEMziuLXyRkm5BIYq1y1GAbb02pRlWvI54OM=
 golang.org/x/sys v0.0.0-20201029080932-201ba4db2418/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=

diff --git a/main.go b/main.go
@@ -18,17 +18,59 @@ package main
 
 import (
 	"os"
+	"os/signal"
+	"syscall"
 
+	"github.com/kardianos/service"
 	"github.com/netflix/weep/cmd"
 	log "github.com/sirupsen/logrus"
 )
 
+var svcLogger service.Logger
+var done chan int
+
 func init() {
 	// Output to stdout instead of the default stderr
 	log.SetOutput(os.Stdout)
-	log.SetReportCaller(true)
+}
+
+type program struct{}
+
+func (p *program) Start(s service.Service) error {
+	go p.run()
+	return nil
+}
+
+func (p *program) run() {
+	shutdown := make(chan os.Signal, 1)
+	done = make(chan int, 1)
+	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)
+	cmd.Execute(shutdown, done)
+}
+
+func (p *program) Stop(s service.Service) error {
+	<-done
+	return nil
 }
 
 func main() {
-	cmd.Execute()
+	svcConfig := &service.Config{
+		Name:        "Weep",
+		DisplayName: "Weep",
+		Description: "The ConsoleMe CLI",
+	}
+
+	prg := &program{}
+	s, err := service.New(prg, svcConfig)
+	if err != nil {
+		log.Fatal(err)
+	}
+	svcLogger, err = s.Logger(nil)
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = s.Run()
+	if err != nil {
+		_ = svcLogger.Error(err)
+	}
 }