From 973df356d703ebf06f49bc37935a79614c5f9c5e Mon Sep 17 00:00:00 2001
From: David Colburn <xero73@gmail.com>
Date: Wed, 8 May 2024 13:06:38 -0700
Subject: [PATCH] fix race on shutdown

---
 pkg/service/process.go     |  2 +-
 pkg/service/service.go     | 14 +++++++++++---
 pkg/service/service_rpc.go |  1 +
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/pkg/service/process.go b/pkg/service/process.go
index 59c793b3..a2ba4ff9 100644
--- a/pkg/service/process.go
+++ b/pkg/service/process.go
@@ -70,7 +70,7 @@ func (p *Process) Gather() ([]*dto.MetricFamily, error) {
 	// Get the metrics from the handler via IPC
 	metricsResponse, err := p.ipcHandlerClient.GetMetrics(context.Background(), &ipc.MetricsRequest{})
 	if err != nil {
-		logger.Warnw("failed to obtain metrics from handler", err, "egress_id", p.req.EgressId)
+		logger.Warnw("failed to obtain metrics from handler", err, "egressID", p.req.EgressId)
 		return make([]*dto.MetricFamily, 0), nil // don't return an error, just skip this handler
 	}
 
diff --git a/pkg/service/service.go b/pkg/service/service.go
index 8f9b84c6..eeb99d43 100644
--- a/pkg/service/service.go
+++ b/pkg/service/service.go
@@ -203,9 +203,17 @@ func (s *Service) killProcess(egressID string, maxUsage float64) {
 }
 
 func (s *Service) Close() {
-	for s.GetRequestCount() > 0 {
+	for {
+		s.mu.RLock()
+		isIdle := len(s.activeHandlers) == 0
+		s.mu.RUnlock()
+
+		if isIdle {
+			logger.Infow("closing server")
+			s.psrpcServer.Shutdown()
+			return
+		}
+
 		time.Sleep(time.Second)
 	}
-	logger.Infow("closing server")
-	s.psrpcServer.Shutdown()
 }
diff --git a/pkg/service/service_rpc.go b/pkg/service/service_rpc.go
index 576375d7..482649cd 100644
--- a/pkg/service/service_rpc.go
+++ b/pkg/service/service_rpc.go
@@ -161,6 +161,7 @@ func (s *Service) AddHandler(egressID string, p *Process) error {
 		}()
 
 	case <-time.After(10 * time.Second):
+		logger.Warnw("no response from handler", nil, "egressID", egressID)
 		_ = p.cmd.Process.Kill()
 		s.processEnded(p, errors.ErrEgressNotFound)
 	}