From f0330d6df1bb0a47215f62c00d62f3cb50c4e903 Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Mon, 30 Oct 2023 10:54:22 -0400
Subject: [PATCH] `identity_hook`: implement PreKill hook, not TaskStop hook
 (#18913)

The allocrunner's `identity_hook` implements the interface for TaskStop, but
this interface is only ever called for task-level hooks. This results in a
leaked goroutine that tries to periodically renew WIs until the client shuts
down gracefully.

Add an implementation for the allocrunner's `PreKill` and `Destroy` hooks, so
that whenever an allocation is stopped or garbage collected we stop renewing its
Workload Identities. This also requires making the `Shutdown` method of `WIDMgr`
safe to call multiple times.
---
 client/allocrunner/identity_hook.go      | 15 +++++++++------
 client/allocrunner/identity_hook_test.go |  8 +++++---
 client/widmgr/widmgr.go                  |  3 +++
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/client/allocrunner/identity_hook.go b/client/allocrunner/identity_hook.go
index c9842fcad..043840351 100644
--- a/client/allocrunner/identity_hook.go
+++ b/client/allocrunner/identity_hook.go
@@ -4,10 +4,7 @@
 package allocrunner
 
 import (
-	"context"
-
 	log "github.com/hashicorp/go-hclog"
-	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
 	"github.com/hashicorp/nomad/client/widmgr"
 )
 
@@ -37,13 +34,19 @@ func (h *identityHook) Prerun() error {
 	return nil
 }
 
-// Stop implements interfaces.TaskStopHook
-func (h *identityHook) Stop(context.Context, *interfaces.TaskStopRequest, *interfaces.TaskStopResponse) error {
+// PreKill implements interfaces.PreKill and is called on allocation stop
+func (h *identityHook) PreKill() {
+	h.widmgr.Shutdown()
+}
+
+// Destroy implements interfaces.Destroy and is called on allocation GC
+func (h *identityHook) Destroy() error {
 	h.widmgr.Shutdown()
 	return nil
 }
 
-// Shutdown implements interfaces.ShutdownHook
+// Shutdown implements interfaces.ShutdownHook and is called when the client
+// gracefully shuts down
 func (h *identityHook) Shutdown() {
 	h.widmgr.Shutdown()
 }
diff --git a/client/allocrunner/identity_hook_test.go b/client/allocrunner/identity_hook_test.go
index c812552f6..5571079b4 100644
--- a/client/allocrunner/identity_hook_test.go
+++ b/client/allocrunner/identity_hook_test.go
@@ -4,7 +4,6 @@
 package allocrunner
 
 import (
-	"context"
 	"testing"
 	"time"
 
@@ -21,7 +20,8 @@ import (
 // statically assert network hook implements the expected interfaces
 var _ interfaces.RunnerPrerunHook = (*identityHook)(nil)
 var _ interfaces.ShutdownHook = (*identityHook)(nil)
-var _ interfaces.TaskStopHook = (*identityHook)(nil)
+var _ interfaces.RunnerPreKillHook = (*identityHook)(nil)
+var _ interfaces.RunnerDestroyHook = (*identityHook)(nil)
 
 func TestIdentityHook_Prerun(t *testing.T) {
 	ci.Parallel(t)
@@ -86,5 +86,7 @@ func TestIdentityHook_Prerun(t *testing.T) {
 		start.Add(ttl).Add(1*time.Second).Unix(),
 	)
 
-	must.NoError(t, hook.Stop(context.Background(), nil, nil))
+	// shutting down twice must not panic
+	hook.PreKill()
+	hook.PreKill()
 }
diff --git a/client/widmgr/widmgr.go b/client/widmgr/widmgr.go
index d26b2b373..1ff8bfe7a 100644
--- a/client/widmgr/widmgr.go
+++ b/client/widmgr/widmgr.go
@@ -206,6 +206,9 @@ func (m *WIDMgr) Shutdown() {
 			close(c)
 		}
 	}
+
+	// ensure it's safe to call Shutdown multiple times
+	m.watchers = map[structs.WIHandle][]chan *structs.SignedWorkloadIdentity{}
 }
 
 // restoreStoredIdentities recreates the state of the WIDMgr from a previously