shutdown hook to clean up agent at exit (#463)

agent/agent.go

@@ -7,12 +7,14 @@ import (
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 	"net"
+	"os"
 )
 
 type Agent struct {
-	root     env_core.Root
-	shares   map[string]*share
-	accesses map[string]*access
+	root        env_core.Root
+	agentSocket string
+	shares      map[string]*share
+	accesses    map[string]*access
 }
 
 func NewAgent(root env_core.Root) (*Agent, error) {
@@ -28,6 +30,7 @@ func NewAgent(root env_core.Root) (*Agent, error) {
 
 func (a *Agent) Run() error {
 	logrus.Infof("started")
+
 	agentSocket, err := a.root.AgentSocket()
 	if err != nil {
 		return err
@@ -36,10 +39,19 @@ func (a *Agent) Run() error {
 	if err != nil {
 		return err
 	}
+	a.agentSocket = agentSocket
+
 	srv := grpc.NewServer()
 	agentGrpc.RegisterAgentServer(srv, &agentGrpcImpl{})
 	if err := srv.Serve(l); err != nil {
 		return err
 	}
+
 	return nil
 }
+
+func (a *Agent) Shutdown() {
+	if err := os.Remove(a.agentSocket); err != nil {
+		logrus.Warnf("unable to remove agent socket: %v", err)
+	}
+}

cmd/zrok/agentStart.go

@@ -5,6 +5,9 @@ import (
 	"github.com/openziti/zrok/environment"
 	"github.com/openziti/zrok/tui"
 	"github.com/spf13/cobra"
+	"os"
+	"os/signal"
+	"syscall"
 )
 
 func init() {
@@ -41,7 +44,19 @@ func (cmd *agentStartCommand) run(_ *cobra.Command, _ []string) {
 		tui.Error("error creating agent", err)
 	}
 
+	c := make(chan os.Signal)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-c
+		cmd.shutdown(a)
+		os.Exit(0)
+	}()
+
 	if err := a.Run(); err != nil {
 		tui.Error("agent aborted", err)
 	}
 }
+
+func (cmd *agentStartCommand) shutdown(a *agent.Agent) {
+	a.Shutdown()
+}