Merge pull request #9995 from tarat44/fix-close-h2ping-client-connections

Add graceful shutdown of client connection and WaitGroup to H2PING

agent/checks/check.go

@@ -519,6 +519,18 @@ type CheckH2PING struct {
 	stop     bool
 	stopCh   chan struct{}
 	stopLock sync.Mutex
+	stopWg   sync.WaitGroup
+}
+
+func shutdownHTTP2ClientConn(clientConn *http2.ClientConn, timeout time.Duration, checkIDString string, logger hclog.Logger) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout/2)
+	defer cancel()
+	err := clientConn.Shutdown(ctx)
+	if err != nil {
+		logger.Warn("Shutdown of H2Ping check client connection gave an error",
+			"check", checkIDString,
+			"error", err)
+	}
 }
 
 func (c *CheckH2PING) check() {
@@ -539,6 +551,7 @@ func (c *CheckH2PING) check() {
 		c.StatusHandler.updateCheck(c.CheckID, api.HealthCritical, message)
 		return
 	}
+	defer shutdownHTTP2ClientConn(clientConn, c.Timeout, c.CheckID.String(), c.Logger)
 	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout)
 	defer cancel()
 	err = clientConn.Ping(ctx)
@@ -558,9 +571,11 @@ func (c *CheckH2PING) Stop() {
 		c.stop = true
 		close(c.stopCh)
 	}
+	c.stopWg.Wait()
 }
 
 func (c *CheckH2PING) run() {
+	defer c.stopWg.Done()
 	// Get the randomized initial pause time
 	initialPauseTime := lib.RandomStagger(c.Interval)
 	next := time.After(initialPauseTime)
@@ -583,6 +598,7 @@ func (c *CheckH2PING) Start() {
 	}
 	c.stop = false
 	c.stopCh = make(chan struct{})
+	c.stopWg.Add(1)
 	go c.run()
 }