From 1a111ab70102ae00639c7f9e051662551a8bf1f7 Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Thu, 6 Aug 2020 14:59:10 -0400
Subject: [PATCH] csi: controller unpublish should check current alloc count

Using the count of node claims from earlier in the `CSIVolume.Unpublish RPC
doesn't correctly account for cases where the RPC was interrupted but
checkpointed. Instead, we'll check the current allocation count and status to
determine whether we need to send a controller unpublish.
---
 nomad/csi_endpoint.go | 47 ++++++++++++++++++++++++-------------------
 1 file changed, 26 insertions(+), 21 deletions(-)

diff --git a/nomad/csi_endpoint.go b/nomad/csi_endpoint.go
index 5d12a4ee3dc2..47025da568ec 100644
--- a/nomad/csi_endpoint.go
+++ b/nomad/csi_endpoint.go
@@ -552,20 +552,6 @@ func (v *CSIVolume) Unpublish(args *structs.CSIVolumeUnpublishRequest, reply *st
 	claim := args.Claim
 	claim.Mode = structs.CSIVolumeClaimRelease
 
-	// we send a controller detach if a Nomad client no longer has
-	// any claim to the volume, so track the counts here
-	var nodeClaims int
-	for _, alloc := range vol.ReadAllocs {
-		if alloc != nil && alloc.NodeID == claim.NodeID {
-			nodeClaims++
-		}
-	}
-	for _, alloc := range vol.WriteAllocs {
-		if alloc != nil && alloc.NodeID == claim.NodeID {
-			nodeClaims++
-		}
-	}
-
 	// previous checkpoints may have set the past claim state already.
 	// in practice we should never see CSIVolumeClaimStateControllerDetached
 	// but having an option for the state makes it easy to add a checkpoint
@@ -584,8 +570,7 @@ func (v *CSIVolume) Unpublish(args *structs.CSIVolumeUnpublishRequest, reply *st
 	}
 
 NODE_DETACHED:
-	nodeClaims--
-	err = v.controllerUnpublishVolume(vol, claim, nodeClaims)
+	err = v.controllerUnpublishVolume(vol, claim)
 	if err != nil {
 		return err
 	}
@@ -630,15 +615,35 @@ func (v *CSIVolume) nodeUnpublishVolume(vol *structs.CSIVolume, claim *structs.C
 	return v.checkpointClaim(vol, claim)
 }
 
-func (v *CSIVolume) controllerUnpublishVolume(vol *structs.CSIVolume, claim *structs.CSIVolumeClaim, nodeClaims int) error {
+func (v *CSIVolume) controllerUnpublishVolume(vol *structs.CSIVolume, claim *structs.CSIVolumeClaim) error {
 
-	// we can drop the claim without sending the controller detach if
-	// another node has a claim on the volume
-	if !vol.ControllerRequired || nodeClaims >= 1 {
+	if !vol.ControllerRequired {
 		claim.State = structs.CSIVolumeClaimStateReadyToFree
 		return nil
 	}
 
+	// we only send a controller detach if a Nomad client no longer has
+	// any claim to the volume, so we need to check the status of claimed
+	// allocations
+	ws := memdb.NewWatchSet()
+	state := v.srv.fsm.State()
+	vol, err := state.CSIVolumeDenormalize(ws, vol)
+	if err != nil {
+		return err
+	}
+	for _, alloc := range vol.ReadAllocs {
+		if alloc != nil && alloc.NodeID == claim.NodeID && !alloc.TerminalStatus() {
+			claim.State = structs.CSIVolumeClaimStateReadyToFree
+			return nil
+		}
+	}
+	for _, alloc := range vol.WriteAllocs {
+		if alloc != nil && alloc.NodeID == claim.NodeID && !alloc.TerminalStatus() {
+			claim.State = structs.CSIVolumeClaimStateReadyToFree
+			return nil
+		}
+	}
+
 	// if the RPC is sent by a client node, it doesn't know the claim's
 	// external node ID.
 	if claim.ExternalNodeID == "" {
@@ -655,7 +660,7 @@ func (v *CSIVolume) controllerUnpublishVolume(vol *structs.CSIVolume, claim *str
 		Secrets:         vol.Secrets,
 	}
 	req.PluginID = vol.PluginID
-	err := v.srv.RPC("ClientCSI.ControllerDetachVolume", req,
+	err = v.srv.RPC("ClientCSI.ControllerDetachVolume", req,
 		&cstructs.ClientCSIControllerDetachVolumeResponse{})
 	if err != nil {
 		// TODO(tgross): need to capture case where ControllerUnpublish previously