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

@@ -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()
 }

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()
 }

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