From 8d0a69bb6d57b859d512937c7e639b0ac781689b Mon Sep 17 00:00:00 2001
From: Adam Hicks <adam@luno.com>
Date: Fri, 4 Oct 2024 14:16:09 +0100
Subject: [PATCH] app: Defer cleanup to ensure pid file removal

---
 app.go  | 17 ++++++++---------
 file.go |  4 +++-
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/app.go b/app.go
index 2801fa3..82b4a32 100644
--- a/app.go
+++ b/app.go
@@ -117,7 +117,7 @@ func (a *App) startup(ctx context.Context) error {
 	return ctx.Err()
 }
 
-func (a *App) cleanup(ctx context.Context) error {
+func (a *App) runShutdownHooks(ctx context.Context) error {
 	var errs []error
 	for idx, h := range a.shutdownHooks {
 		if ctx.Err() != nil {
@@ -148,6 +148,7 @@ func (a *App) cleanup(ctx context.Context) error {
 func (a *App) Run() int {
 	ac := NewAppContext(context.Background())
 	defer ac.Stop()
+	defer a.cleanup(ac.TerminationContext)
 
 	ctx := ac.AppContext
 
@@ -166,13 +167,7 @@ func (a *App) Run() int {
 		exit = 1
 	}
 
-	// TODO(adam): Move pid removal into Shutdown
-
-	// This should be called in Shutdown so that clients which call that instead of
-	// Run can get the right behaviour
-	if a.UseProcessFile {
-		removePIDFile(ctx)
-	}
+	log.Info(ctx, "Waiting to terminate", j.MKV{"exit_code": exit})
 
 	// Wait for termination in case we've only been told to quit
 	<-ac.TerminationContext.Done()
@@ -255,7 +250,7 @@ func (a *App) Shutdown() error {
 	defer a.OnEvent(ctx, Event{Type: AppTerminated})
 
 	defer func() {
-		err := a.cleanup(ctx)
+		err := a.runShutdownHooks(ctx)
 		if err != nil {
 			// NoReturnErr: Log
 			log.Error(ctx, errors.Wrap(err, ""))
@@ -325,6 +320,10 @@ func (a *App) RunningProcesses() []string {
 	return ret
 }
 
+func (a *App) cleanup(ctx context.Context) {
+	removePIDFile(ctx)
+}
+
 // Wait is a cancellable wait, it will return either when
 // d has passed or ctx is cancelled.
 // It will return an error if cancelled early.
diff --git a/file.go b/file.go
index 6163950..a11c2c7 100644
--- a/file.go
+++ b/file.go
@@ -35,7 +35,9 @@ func createPIDFile() error {
 
 func removePIDFile(ctx context.Context) {
 	err := os.Remove(fileName)
-	if err != nil {
+	if errors.Is(err, os.ErrNotExist) {
+		// NoReturnErr: File already gone, no worries
+	} else if err != nil {
 		// NoReturnErr: We'll terminate after this so just log
 		log.Error(ctx, errors.Wrap(err, "remove pid file", j.KV("file", fileName)))
 	}