Shutdown metrics server when listener exits (#3445)

Author: nikola-jokic

Diff for: cmd/ghalistener/app/app.go

@@ -117,15 +117,19 @@ func (app *App) Run(ctx context.Context) error {
 	}
 
 	g, ctx := errgroup.WithContext(ctx)
+	metricsCtx, cancelMetrics := context.WithCancelCause(ctx)
+
 	g.Go(func() error {
 		app.logger.Info("Starting listener")
-		return app.listener.Listen(ctx, app.worker)
+		listnerErr := app.listener.Listen(ctx, app.worker)
+		cancelMetrics(fmt.Errorf("Listener exited: %w", listnerErr))
+		return listnerErr
 	})
 
 	if app.metrics != nil {
 		g.Go(func() error {
 			app.logger.Info("Starting metrics server")
-			return app.metrics.ListenAndServe(ctx)
+			return app.metrics.ListenAndServe(metricsCtx)
 		})
 	}
 

Diff for: cmd/ghalistener/listener/listener.go

@@ -174,8 +174,8 @@ func (l *Listener) Listen(ctx context.Context, handler Handler) error {
 			continue
 		}
 
-		// New context is created to avoid cancelation during message handling.
-		if err := l.handleMessage(context.Background(), handler, msg); err != nil {
+		// Remove cancellation from the context to avoid cancelling the message handling.
+		if err := l.handleMessage(context.WithoutCancel(ctx), handler, msg); err != nil {
 			return fmt.Errorf("failed to handle message: %w", err)
 		}
 	}

Diff for: cmd/ghalistener/listener/listener_test.go

@@ -484,8 +484,8 @@ func TestListener_Listen(t *testing.T) {
 			).
 			Once()
 
-			// Ensure delete message is called with background context
-		client.On("DeleteMessage", context.Background(), mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
+		// Ensure delete message is called without cancel
+		client.On("DeleteMessage", context.WithoutCancel(ctx), mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
 
 		config.Client = client
 

Diff for: cmd/ghalistener/metrics/metrics.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"strconv"
+	"time"
 
 	"github.com/actions/actions-runner-controller/github/actions"
 	"github.com/go-logr/logr"
@@ -338,7 +339,9 @@ func (e *exporter) ListenAndServe(ctx context.Context) error {
 	e.logger.Info("starting metrics server", "addr", e.srv.Addr)
 	go func() {
 		<-ctx.Done()
-		e.logger.Info("stopping metrics server")
+		e.logger.Info("stopping metrics server", "err", ctx.Err())
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
 		e.srv.Shutdown(ctx)
 	}()
 	return e.srv.ListenAndServe()