From 9653980a30ff496013e8c720a2db9a6b25514135 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Mon, 16 Sep 2024 13:12:12 -0400
Subject: [PATCH] agent shutdown lifecycle management improvements (#463, #748)

---
 agent/agent.go | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 7eb30359..e88f3e84 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -2,6 +2,7 @@ package agent
 
 import (
 	"github.com/openziti/zrok/agent/agentGrpc"
+	"github.com/openziti/zrok/agent/proctree"
 	"github.com/openziti/zrok/environment/env_core"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -61,9 +62,19 @@ func (a *Agent) Run() error {
 }
 
 func (a *Agent) Shutdown() {
+	logrus.Infof("stopping")
+
 	if err := os.Remove(a.agentSocket); err != nil {
 		logrus.Warnf("unable to remove agent socket: %v", err)
 	}
+	for _, shr := range a.shares {
+		logrus.Infof("stopping share '%v'", shr.token)
+		a.outShares <- shr
+	}
+	for _, acc := range a.accesses {
+		logrus.Infof("stopping access '%v'", acc.token)
+		a.outAccesses <- acc
+	}
 }
 
 func (a *Agent) manager() {
@@ -77,8 +88,18 @@ func (a *Agent) manager() {
 			a.shares[inShare.token] = inShare
 
 		case outShare := <-a.outShares:
-			logrus.Infof("removing share '%v'", outShare.token)
-			delete(a.shares, outShare.token)
+			if outShare.token != "" {
+				logrus.Infof("removing share '%v'", outShare.token)
+				if err := proctree.StopChild(outShare.process); err != nil {
+					logrus.Errorf("error stopping share '%v': %v", outShare.token, err)
+				}
+				if err := proctree.WaitChild(outShare.process); err != nil {
+					logrus.Errorf("error joining share '%v': %v", outShare.token, err)
+				}
+				delete(a.shares, outShare.token)
+			} else {
+				logrus.Debug("skipping unidentified (orphaned) share removal")
+			}
 		}
 	}
 }