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

This commit is contained in:
Michael Quigley 2024-08-23 12:35:19 -04:00
parent 23796be138
commit 3ba179673e
No known key found for this signature in database
GPG Key ID: 9B60314A9DD20A62
2 changed files with 30 additions and 3 deletions

View File

@ -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)
}
}

View File

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