stateful deployments: use TaskGroupVolumeClaim table to associate volume requests with volume IDs

api/allocations.go

@@ -278,7 +278,6 @@ type Allocation struct {
 	Resources             *Resources
 	TaskResources         map[string]*Resources
 	AllocatedResources    *AllocatedResources
-	HostVolumeIDs         []string
 	Services              map[string]string
 	Metrics               *AllocationMetric
 	DesiredStatus         string

nomad/fsm.go

@@ -389,6 +389,8 @@ func (n *nomadFSM) Apply(log *raft.Log) interface{} {
 		return n.applyHostVolumeRegister(msgType, buf[1:], log.Index)
 	case structs.HostVolumeDeleteRequestType:
 		return n.applyHostVolumeDelete(msgType, buf[1:], log.Index)
+		// case structs.TaskGroupVolumeClaimDeleteRequestType:
+		// 	return n.applyTaskVolumeClaimDelete(buf[1:], log.Index)
 	}
 
 	// Check enterprise only message types.
@@ -2450,6 +2452,21 @@ func (n *nomadFSM) applyHostVolumeDelete(msgType structs.MessageType, buf []byte
 	return nil
 }
 
+// func (n *nomadFSM) applyTaskVolumeClaimDelete(buf []byte, index uint64) interface{} {
+// 	defer metrics.MeasureSince([]string{"nomad", "fsm", "apply_task_group_volume_claim_delete"}, time.Now())
+
+// 	var req structs.TaskGroupVolumeClaimDeleteRequest
+// 	if err := structs.Decode(buf, &req); err != nil {
+// 		panic(fmt.Errorf("failed to decode request: %v", err))
+// 	}
+
+// 	if err := n.state.DeleteTaskGroupVolumeClaim(index, req.ClaimID); err != nil {
+// 		n.logger.Error("DeleteTaskGroupVolumeClaim failed", "error", err)
+// 		return err
+// 	}
+// 	return nil
+// }
+
 func (s *nomadSnapshot) Persist(sink raft.SnapshotSink) error {
 	defer metrics.MeasureSince([]string{"nomad", "fsm", "persist"}, time.Now())
 	// Register the nodes

nomad/state/schema.go

@@ -29,6 +29,7 @@ const (
 	TableHostVolumes          = "host_volumes"
 	TableCSIVolumes           = "csi_volumes"
 	TableCSIPlugins           = "csi_plugins"
+	TableTaskGroupVolumeClaim = "task_volume"
 )
 
 const (
@@ -45,6 +46,7 @@ const (
 	indexSigningKey    = "signing_key"
 	indexAuthMethod    = "auth_method"
 	indexNodePool      = "node_pool"
+	indexVolumeID      = "volume_id"
 )
 
 var (
@@ -102,6 +104,7 @@ func init() {
 		aclAuthMethodsTableSchema,
 		bindingRulesTableSchema,
 		hostVolumeTableSchema,
+		taskVolumeClaimSchema,
 	}...)
 }
 
@@ -1706,3 +1709,42 @@ func hostVolumeTableSchema() *memdb.TableSchema {
 		},
 	}
 }
+
+func taskVolumeClaimSchema() *memdb.TableSchema {
+	return &memdb.TableSchema{
+		Name: TableTaskGroupVolumeClaim,
+		Indexes: map[string]*memdb.IndexSchema{
+			indexID: {
+				Name:         indexID,
+				AllowMissing: false,
+				Unique:       true,
+
+				// Use a compound index so the combination of (Namespace, JobID, TaskGroupName)
+				// is uniquely identifying
+				Indexer: &memdb.CompoundIndex{
+					Indexes: []memdb.Indexer{
+						&memdb.StringFieldIndex{
+							Field: "Namespace",
+						},
+
+						&memdb.StringFieldIndex{
+							Field: "JobID",
+						},
+
+						&memdb.StringFieldIndex{
+							Field: "TaskGroupName",
+						},
+					},
+				},
+			},
+			indexVolumeID: {
+				Name:         indexVolumeID,
+				AllowMissing: false,
+				Unique:       false,
+				Indexer: &memdb.StringFieldIndex{
+					Field: "VolumeID",
+				},
+			},
+		},
+	}
+}

nomad/state/state_store.go

@@ -2025,6 +2025,14 @@ func (s *StateStore) DeleteJobTxn(index uint64, namespace, jobID string, txn Txn
 	if _, err = txn.DeleteAll("scaling_event", "id", namespace, jobID); err != nil {
 		return fmt.Errorf("deleting job scaling events failed: %v", err)
 	}
+
+	// Delete task group volume claims
+	for _, tg := range job.TaskGroups {
+		if _, err = txn.DeleteAll(TableTaskGroupVolumeClaim, indexID, namespace, jobID, tg.Name); err != nil {
+			return fmt.Errorf("deleting job volume claims failed: %v", err)
+		}
+	}
+
 	if err := txn.Insert("index", &IndexEntry{"scaling_event", index}); err != nil {
 		return fmt.Errorf("index update failed: %v", err)
 	}
@@ -4125,6 +4133,54 @@ func (s *StateStore) upsertAllocsImpl(index uint64, allocs []*structs.Allocation
 			if alloc.Job == nil {
 				return fmt.Errorf("attempting to upsert allocation %q without a job", alloc.ID)
 			}
+
+			// Check if the alloc requires sticky volumes. If yes, find a node
+			// that has the right volume and update the task group volume
+			// claims table
+			stickyVolumes := []*structs.TaskGroupVolumeClaim{}
+			for _, tg := range alloc.Job.TaskGroups {
+				for _, v := range tg.Volumes {
+					if !v.Sticky {
+						continue
+					}
+					stickyVolumes = append(stickyVolumes, &structs.TaskGroupVolumeClaim{
+						Namespace:     alloc.Namespace,
+						JobID:         alloc.JobID,
+						TaskGroupName: tg.Name,
+						AllocID:       alloc.ID,
+						VolumeName:    v.Source,
+					})
+				}
+			}
+
+			if len(stickyVolumes) > 0 {
+				for _, sv := range stickyVolumes {
+					// has this volume been claimed already?
+					existingClaim, err := s.GetTaskGroupVolumeClaim(nil, sv.Namespace, sv.JobID, sv.TaskGroupName)
+					if err != nil {
+						return err
+					}
+
+					if existingClaim == nil {
+						allocNode, err := s.NodeByID(nil, alloc.NodeID)
+						if err != nil {
+							return err
+						}
+
+						// since there's no existing claim, find a volume and register a claim
+						for _, v := range allocNode.HostVolumes {
+							if v.Name != sv.VolumeName {
+								continue
+							}
+
+							sv.VolumeID = v.ID
+							if err := s.upsertTaskGroupVolumeClaimImpl(index, sv, txn); err != nil {
+								return err
+							}
+						}
+					}
+				}
+			}
 		} else {
 			alloc.CreateIndex = exist.CreateIndex
 			alloc.ModifyIndex = index

nomad/state/state_store_task_group_volume_claims.go

@@ -0,0 +1,132 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package state
+
+import (
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/hashicorp/go-memdb"
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+// upsertTaskGroupVolumeClaimImpl is used to insert a task group volume claim into
+// the state store.
+func (s *StateStore) upsertTaskGroupVolumeClaimImpl(
+	index uint64, claim *structs.TaskGroupVolumeClaim, txn *txn) error {
+
+	existingRaw, err := txn.First(TableTaskGroupVolumeClaim, indexID, claim.Namespace, claim.JobID, claim.TaskGroupName)
+	if err != nil {
+		return fmt.Errorf("Task group volume association lookup failed: %v", err)
+	}
+
+	var existing *structs.TaskGroupVolumeClaim
+	if existingRaw != nil {
+		existing = existingRaw.(*structs.TaskGroupVolumeClaim)
+	}
+
+	now := time.Now().UTC()
+
+	if existing != nil {
+		// do allocation ID and volume ID match?
+		if existing.ClaimedByAlloc(claim) {
+			return nil
+		}
+
+		claim.CreateIndex = existing.CreateIndex
+		claim.ModifyIndex = index
+		claim.CreateTime = existing.CreateTime
+		claim.ModifyTime = now
+	} else {
+		claim.CreateIndex = index
+		claim.ModifyIndex = index
+		claim.CreateTime = now
+		claim.ModifyTime = now
+	}
+
+	// Insert the claim into the table.
+	if err := txn.Insert(TableTaskGroupVolumeClaim, claim); err != nil {
+		return fmt.Errorf("Task group volume claim insert failed: %v", err)
+	}
+
+	// Perform the index table update to mark the new insert.
+	if err := txn.Insert(tableIndex, &IndexEntry{TableTaskGroupVolumeClaim, index}); err != nil {
+		return fmt.Errorf("index update failed: %v", err)
+	}
+
+	return nil
+}
+
+// DeleteTaskGroupVolumeClaim is responsible for deleting volume claims.
+func (s *StateStore) DeleteTaskGroupVolumeClaim(index uint64, namespace, jobID, taskGroupName string) error {
+	txn := s.db.WriteTxnMsgT(structs.TaskGroupVolumeClaimDeleteRequestType, index)
+	defer txn.Abort()
+
+	existing, err := txn.First(TableTaskGroupVolumeClaim, indexID, namespace, jobID, taskGroupName)
+	if err != nil {
+		return fmt.Errorf("Task group volume claim lookup failed: %v", err)
+	}
+	if existing == nil {
+		return errors.New("ACL binding rule not found")
+	}
+
+	// Delete the existing entry from the table.
+	if err := txn.Delete(TableTaskGroupVolumeClaim, existing); err != nil {
+		return fmt.Errorf("Task group volume claim deletion failed: %v", err)
+	}
+
+	// Update the index table to indicate an update has occurred.
+	if err := txn.Insert(tableIndex, &IndexEntry{TableTaskGroupVolumeClaim, index}); err != nil {
+		return fmt.Errorf("index update failed: %v", err)
+	}
+
+	return txn.Commit()
+}
+
+// GetTaskGroupVolumeClaims returns an iterator that contains all task group
+// volume associations stored within state.
+func (s *StateStore) GetTaskGroupVolumeClaims(ws memdb.WatchSet) (memdb.ResultIterator, error) {
+	txn := s.db.ReadTxn()
+
+	iter, err := txn.Get(TableTaskGroupVolumeClaim, indexID)
+	if err != nil {
+		return nil, fmt.Errorf("Task group volume claims lookup failed: %v", err)
+	}
+	ws.Add(iter.WatchCh())
+
+	return iter, nil
+}
+
+// GetTaskGroupVolumeClaimsByVolumeID gets an iterator that contains all task
+// group volume claims that claim a particular volume ID
+func (s *StateStore) GetTaskGroupVolumeClaimsByVolumeID(ws memdb.WatchSet, volID string) (memdb.ResultIterator, error) {
+	txn := s.db.ReadTxn()
+
+	iter, err := txn.Get(TableTaskGroupVolumeClaim, indexVolumeID, volID)
+	if err != nil {
+		return nil, fmt.Errorf("Task group volume claims lookup failed: %v", err)
+	}
+	ws.Add(iter.WatchCh())
+
+	return iter, nil
+}
+
+// GetTaskGroupVolumeClaim returns a volume claim that matches the namespace,
+// job id and task group name (there can be only one)
+func (s *StateStore) GetTaskGroupVolumeClaim(ws memdb.WatchSet, namespace, jobID, taskGroupName string) (*structs.TaskGroupVolumeClaim, error) {
+	txn := s.db.ReadTxn()
+
+	watchCh, existing, err := txn.FirstWatch(TableTaskGroupVolumeClaim, indexID, namespace, jobID, taskGroupName)
+	if err != nil {
+		return nil, fmt.Errorf("Task group volume claim lookup failed: %v", err)
+	}
+	ws.Add(watchCh)
+
+	if existing != nil {
+		return existing.(*structs.TaskGroupVolumeClaim), nil
+	}
+
+	return nil, nil
+}

nomad/structs/host_volumes.go

@@ -428,3 +428,12 @@ type HostVolumeListResponse struct {
 	Volumes []*HostVolumeStub
 	QueryMeta
 }
+
+type TaskGroupVolumeClaimDeleteRequest struct {
+	ClaimID string
+	WriteRequest
+}
+
+type TaskGroupVolumeClaimDeleteResponse struct {
+	WriteMeta
+}