csi: plugins track jobs in addition to allocations, and use job information to set expected counts

api/csi.go

@@ -214,8 +214,10 @@ type CSIPlugin struct {
 	Allocations         []*AllocationListStub
 	ControllersHealthy  int
 	ControllersExpected int
+	ControllerJobs      []JobDescription
 	NodesHealthy        int
 	NodesExpected       int
+	NodeJobs            []JobDescription
 	CreateIndex         uint64
 	ModifyIndex         uint64
 }
@@ -246,6 +248,12 @@ func (v CSIPluginIndexSort) Swap(i, j int) {
 	v[i], v[j] = v[j], v[i]
 }
 
+type JobDescription struct {
+	Namespace string
+	ID        string
+	Version   uint64
+}
+
 // CSIPlugins returns a handle on the CSIPlugins endpoint
 func (c *Client) CSIPlugins() *CSIPlugins {
 	return &CSIPlugins{client: c}

api/nodes.go

@@ -568,6 +568,7 @@ type CSIControllerInfo struct {
 // as plugin health changes on the node.
 type CSIInfo struct {
 	PluginID                 string
+	AllocID                  string
 	Healthy                  bool
 	HealthDescription        string
 	UpdateTime               time.Time

command/agent/csi_endpoint.go

@@ -278,11 +278,13 @@ func structsCSIPluginToApi(plug *structs.CSIPlugin) *api.CSIPlugin {
 		Allocations:         make([]*api.AllocationListStub, 0, len(plug.Allocations)),
 		ControllerRequired:  plug.ControllerRequired,
 		ControllersHealthy:  plug.ControllersHealthy,
-		ControllersExpected: len(plug.Controllers),
+		ControllersExpected: plug.ControllersExpected,
 		Controllers:         make(map[string]*api.CSIInfo, len(plug.Controllers)),
+		ControllerJobs:      make([]api.JobDescription, 0, plug.ControllersExpected),
 		NodesHealthy:        plug.NodesHealthy,
-		NodesExpected:       len(plug.Nodes),
+		NodesExpected:       plug.NodesExpected,
 		Nodes:               make(map[string]*api.CSIInfo, len(plug.Nodes)),
+		NodeJobs:            make([]api.JobDescription, 0, plug.NodesExpected),
 		CreateIndex:         plug.CreateIndex,
 		ModifyIndex:         plug.ModifyIndex,
 	}
@@ -299,9 +301,26 @@ func structsCSIPluginToApi(plug *structs.CSIPlugin) *api.CSIPlugin {
 		out.Allocations = append(out.Allocations, structsAllocListStubToApi(a))
 	}
 
+	for _, jd := range plug.ControllerJobs.AsSlice() {
+		out.ControllerJobs = append(out.ControllerJobs, structsJobDescriptionToApi(jd))
+	}
+
+	for _, jd := range plug.NodeJobs.AsSlice() {
+		out.NodeJobs = append(out.NodeJobs, structsJobDescriptionToApi(jd))
+	}
+
 	return out
 }
 
+// structsJobDescriptionToApi converts the struct
+func structsJobDescriptionToApi(desc structs.JobDescription) api.JobDescription {
+	return api.JobDescription{
+		Namespace: desc.Namespace,
+		ID:        desc.ID,
+		Version:   desc.Version,
+	}
+}
+
 // structsCSIVolumeToApi converts CSIVolume, creating the allocation array
 func structsCSIVolumeToApi(vol *structs.CSIVolume) *api.CSIVolume {
 	if vol == nil {
@@ -358,6 +377,7 @@ func structsCSIInfoToApi(info *structs.CSIInfo) *api.CSIInfo {
 	}
 	out := &api.CSIInfo{
 		PluginID:                 info.PluginID,
+		AllocID:                  info.AllocID,
 		Healthy:                  info.Healthy,
 		HealthDescription:        info.HealthDescription,
 		UpdateTime:               info.UpdateTime,

command/agent/csi_endpoint_test.go

@@ -10,7 +10,6 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/state"
 	"github.com/hashicorp/nomad/nomad/structs"
-	"github.com/kr/pretty"
 	"github.com/stretchr/testify/require"
 )
 
@@ -33,11 +32,13 @@ func TestHTTP_CSIEndpointPlugin(t *testing.T) {
 		out, ok := obj.(*api.CSIPlugin)
 		require.True(t, ok)
 
-		require.Equal(t, 1, out.ControllersExpected)
+		// ControllersExpected is 0 because this plugin was created without a job,
+		// which sets expected
+		require.Equal(t, 0, out.ControllersExpected)
 		require.Equal(t, 1, out.ControllersHealthy)
 		require.Len(t, out.Controllers, 1)
 
-		require.Equal(t, 2, out.NodesExpected)
+		require.Equal(t, 0, out.NodesExpected)
 		require.Equal(t, 2, out.NodesHealthy)
 		require.Len(t, out.Nodes, 2)
 	})
@@ -92,11 +93,7 @@ func TestHTTP_CSIEndpointVolume(t *testing.T) {
 		out, ok := raw.(*api.CSIVolume)
 		require.True(t, ok)
 
-		pretty.Log(out)
-
-		require.Equal(t, 1, out.ControllersExpected)
 		require.Equal(t, 1, out.ControllersHealthy)
-		require.Equal(t, 2, out.NodesExpected)
 		require.Equal(t, 2, out.NodesHealthy)
 	})
 }

nomad/state/state_store.go

@@ -1050,10 +1050,14 @@ func upsertNodeCSIPlugins(txn *memdb.Txn, node *structs.Node, index uint64) erro
 				return nil
 			}
 			plug = structs.NewCSIPlugin(info.PluginID, index)
-			plug.Provider = info.Provider
-			plug.Version = info.ProviderVersion
 		}
 
+		// the plugin may have been created by the job being updated, in which case
+		// this data will not be configured, it's only available to the fingerprint
+		// system
+		plug.Provider = info.Provider
+		plug.Version = info.ProviderVersion
+
 		err = plug.AddPlugin(node.ID, info)
 		if err != nil {
 			return err
@@ -1203,10 +1207,15 @@ func updateOrGCPlugin(index uint64, txn *memdb.Txn, plug *structs.CSIPlugin) err
 	return nil
 }
 
-// deleteJobFromPlugin removes the allocations of this job from any plugins the job is
+// deleteJobFromPlugins removes the allocations of this job from any plugins the job is
 // running, possibly deleting the plugin if it's no longer in use. It's called in DeleteJobTxn
-func (s *StateStore) deleteJobFromPlugin(index uint64, txn *memdb.Txn, job *structs.Job) error {
+func (s *StateStore) deleteJobFromPlugins(index uint64, txn *memdb.Txn, job *structs.Job) error {
 	ws := memdb.NewWatchSet()
+	summary, err := s.JobSummaryByID(ws, job.Namespace, job.ID)
+	if err != nil {
+		return fmt.Errorf("error gettting job summary: %v", err)
+	}
+
 	allocs, err := s.AllocsByJob(ws, job.Namespace, job.ID, false)
 	if err != nil {
 		return fmt.Errorf("error getting allocations: %v", err)
@@ -1218,7 +1227,6 @@ func (s *StateStore) deleteJobFromPlugin(index uint64, txn *memdb.Txn, job *stru
 	}
 
 	plugAllocs := []*pair{}
-	plugins := map[string]*structs.CSIPlugin{}
 
 	for _, a := range allocs {
 		tg := a.Job.LookupTaskGroup(a.TaskGroup)
@@ -1232,6 +1240,8 @@ func (s *StateStore) deleteJobFromPlugin(index uint64, txn *memdb.Txn, job *stru
 		}
 	}
 
+	plugins := map[string]*structs.CSIPlugin{}
+
 	for _, x := range plugAllocs {
 		plug, ok := plugins[x.pluginID]
 
@@ -1255,6 +1265,7 @@ func (s *StateStore) deleteJobFromPlugin(index uint64, txn *memdb.Txn, job *stru
 	}
 
 	for _, plug := range plugins {
+		plug.DeleteJob(job, summary)
 		err = updateOrGCPlugin(index, txn, plug)
 		if err != nil {
 			return err
@@ -1354,6 +1365,7 @@ func (s *StateStore) upsertJobImpl(index uint64, job *structs.Job, keepVersion b
 
 	// Check if the job already exists
 	existing, err := txn.First("jobs", "id", job.Namespace, job.ID)
+	var existingJob *structs.Job
 	if err != nil {
 		return fmt.Errorf("job lookup failed: %v", err)
 	}
@@ -1363,7 +1375,7 @@ func (s *StateStore) upsertJobImpl(index uint64, job *structs.Job, keepVersion b
 		job.CreateIndex = existing.(*structs.Job).CreateIndex
 		job.ModifyIndex = index
 
-		existingJob := existing.(*structs.Job)
+		existingJob = existing.(*structs.Job)
 
 		// Bump the version unless asked to keep it. This should only be done
 		// when changing an internal field such as Stable. A spec change should
@@ -1413,6 +1425,10 @@ func (s *StateStore) upsertJobImpl(index uint64, job *structs.Job, keepVersion b
 		return fmt.Errorf("unable to update job scaling policies: %v", err)
 	}
 
+	if err := s.updateJobCSIPlugins(index, job, existingJob, txn); err != nil {
+		return fmt.Errorf("unable to update job scaling policies: %v", err)
+	}
+
 	// Insert the job
 	if err := txn.Insert("jobs", job); err != nil {
 		return fmt.Errorf("job insert failed: %v", err)
@@ -1507,6 +1523,12 @@ func (s *StateStore) DeleteJobTxn(index uint64, namespace, jobID string, txn Txn
 		return err
 	}
 
+	// Cleanup plugins registered by this job, before we delete the summary
+	err = s.deleteJobFromPlugins(index, txn, job)
+	if err != nil {
+		return fmt.Errorf("deleting job from plugin: %v", err)
+	}
+
 	// Delete the job summary
 	if _, err = txn.DeleteAll("job_summary", "id", namespace, jobID); err != nil {
 		return fmt.Errorf("deleting job summary failed: %v", err)
@@ -1528,12 +1550,6 @@ func (s *StateStore) DeleteJobTxn(index uint64, namespace, jobID string, txn Txn
 		return fmt.Errorf("index update failed: %v", err)
 	}
 
-	// Cleanup plugins registered by this job
-	err = s.deleteJobFromPlugin(index, txn, job)
-	if err != nil {
-		return fmt.Errorf("deleting job from plugin: %v", err)
-	}
-
 	return nil
 }
 
@@ -2222,10 +2238,10 @@ func (s *StateStore) CSIVolumeDenormalizePlugins(ws memdb.WatchSet, vol *structs
 	vol.ControllerRequired = plug.ControllerRequired
 	vol.ControllersHealthy = plug.ControllersHealthy
 	vol.NodesHealthy = plug.NodesHealthy
-	// This number is incorrect! The expected number of node plugins is actually this +
-	// the number of blocked evaluations for the jobs controlling these plugins
-	vol.ControllersExpected = len(plug.Controllers)
-	vol.NodesExpected = len(plug.Nodes)
+
+	// This value may be stale, but stale is ok
+	vol.ControllersExpected = plug.ControllersExpected
+	vol.NodesExpected = plug.NodesExpected
 
 	vol.Schedulable = vol.NodesHealthy > 0
 	if vol.ControllerRequired {
@@ -2327,7 +2343,7 @@ func (s *StateStore) CSIPluginByID(ws memdb.WatchSet, id string) (*structs.CSIPl
 	return plug, nil
 }
 
-// CSIPluginDenormalize returns a CSIPlugin with allocation details
+// CSIPluginDenormalize returns a CSIPlugin with allocation details. Always called on a copy of the plugin.
 func (s *StateStore) CSIPluginDenormalize(ws memdb.WatchSet, plug *structs.CSIPlugin) (*structs.CSIPlugin, error) {
 	if plug == nil {
 		return nil, nil
@@ -4498,6 +4514,70 @@ func (s *StateStore) updateJobScalingPolicies(index uint64, job *structs.Job, tx
 	return nil
 }
 
+// updateJobCSIPlugins runs on job update indexes the job in the plugin
+func (s *StateStore) updateJobCSIPlugins(index uint64, job, prev *structs.Job, txn *memdb.Txn) error {
+	ws := memdb.NewWatchSet()
+	plugIns := make(map[string]*structs.CSIPlugin)
+
+	loop := func(job *structs.Job, delete bool) error {
+		for _, tg := range job.TaskGroups {
+			for _, t := range tg.Tasks {
+				if t.CSIPluginConfig == nil {
+					continue
+				}
+
+				plugIn, ok := plugIns[t.CSIPluginConfig.ID]
+				if !ok {
+					p, err := s.CSIPluginByID(ws, t.CSIPluginConfig.ID)
+					if err != nil {
+						return fmt.Errorf("%v", err)
+					}
+					if p == nil {
+						plugIn = structs.NewCSIPlugin(t.CSIPluginConfig.ID, index)
+					} else {
+						plugIn = p.Copy()
+						plugIn.ModifyIndex = index
+					}
+					plugIns[plugIn.ID] = plugIn
+				}
+
+				if delete {
+					plugIn.DeleteJob(job, nil)
+				} else {
+					plugIn.AddJob(job, nil)
+				}
+			}
+		}
+
+		return nil
+	}
+
+	if prev != nil {
+		err := loop(prev, true)
+		if err != nil {
+			return fmt.Errorf("%v", err)
+		}
+	}
+
+	err := loop(job, false)
+	if err != nil {
+		return fmt.Errorf("%v", err)
+	}
+
+	for _, plugIn := range plugIns {
+		err = txn.Insert("csi_plugins", plugIn)
+		if err != nil {
+			return fmt.Errorf("csi_plugins insert error: %v", err)
+		}
+	}
+
+	if err := txn.Insert("index", &IndexEntry{"csi_plugins", index}); err != nil {
+		return fmt.Errorf("index update failed: %v", err)
+	}
+
+	return nil
+}
+
 // updateDeploymentWithAlloc is used to update the deployment state associated
 // with the given allocation. The passed alloc may be updated if the deployment
 // status has changed to capture the modify index at which it has changed.
@@ -4709,6 +4789,8 @@ func (s *StateStore) updateSummaryWithAlloc(index uint64, alloc *structs.Allocat
 	if summaryChanged {
 		jobSummary.ModifyIndex = index
 
+		s.updatePluginWithJobSummary(index, jobSummary, alloc, txn)
+
 		// Update the indexes table for job summary
 		if err := txn.Insert("index", &IndexEntry{"job_summary", index}); err != nil {
 			return fmt.Errorf("index update failed: %v", err)
@@ -4745,6 +4827,7 @@ func (s *StateStore) updatePluginWithAlloc(index uint64, alloc *structs.Allocati
 				return nil
 			}
 			plug = plug.Copy()
+
 			err = plug.DeleteAlloc(alloc.ID, alloc.NodeID)
 			if err != nil {
 				return err
@@ -4759,6 +4842,41 @@ func (s *StateStore) updatePluginWithAlloc(index uint64, alloc *structs.Allocati
 	return nil
 }
 
+// updatePluginWithJobSummary updates the CSI plugins for a job when the
+// job summary is updated by an alloc
+func (s *StateStore) updatePluginWithJobSummary(index uint64, summary *structs.JobSummary, alloc *structs.Allocation,
+	txn *memdb.Txn) error {
+
+	ws := memdb.NewWatchSet()
+	tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
+	if tg == nil {
+		return nil
+	}
+
+	for _, t := range tg.Tasks {
+		if t.CSIPluginConfig != nil {
+			pluginID := t.CSIPluginConfig.ID
+			plug, err := s.CSIPluginByID(ws, pluginID)
+			if err != nil {
+				return err
+			}
+			if plug == nil {
+				plug = structs.NewCSIPlugin(pluginID, index)
+			} else {
+				plug = plug.Copy()
+			}
+
+			plug.UpdateExpectedWithJob(alloc.Job, summary, alloc.ServerTerminalStatus())
+			err = updateOrGCPlugin(index, txn, plug)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
 // UpsertACLPolicies is used to create or update a set of ACL policies
 func (s *StateStore) UpsertACLPolicies(index uint64, policies []*structs.ACLPolicy) error {
 	txn := s.db.Txn(true)