From bd50ea9c00f2766caa5edccf9ec72c242eedba2e Mon Sep 17 00:00:00 2001
From: Piotr Kazmierczak <470696+pkazmierczak@users.noreply.github.com>
Date: Wed, 4 Dec 2024 17:17:29 +0100
Subject: [PATCH] super hacky prototype

---
 nomad/structs/structs.go |  6 +++---
 scheduler/feasible.go    | 42 ++++++++++++++++++++++++++++++----------
 scheduler/scheduler.go   |  6 ++++++
 3 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 1894bcf10da..de87e373448 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -11114,9 +11114,9 @@ type Allocation struct {
 	// AllocatedResources is the total resources allocated for the task group.
 	AllocatedResources *AllocatedResources
 
-	// VolumeID is the ID of the host volume that this allocation requires.
-	// FIXME:could be multiple, could be CSI? can't just be a string
-	VolumeID *string
+	// VolumeIDs is a list of volume IDs (host or CSI) that this allocation
+	// requires.
+	VolumeIDs []string
 
 	// Metrics associated with this allocation
 	Metrics *AllocMetric
diff --git a/scheduler/feasible.go b/scheduler/feasible.go
index 7ea834fb6be..74f5c2e491e 100644
--- a/scheduler/feasible.go
+++ b/scheduler/feasible.go
@@ -8,18 +8,21 @@ import (
 	"fmt"
 	"reflect"
 	"regexp"
+	"slices"
 	"strconv"
 	"strings"
 
 	"github.com/hashicorp/go-memdb"
 	"github.com/hashicorp/go-version"
 	"github.com/hashicorp/nomad/helper/constraints/semver"
+	"github.com/hashicorp/nomad/nomad/state"
 	"github.com/hashicorp/nomad/nomad/structs"
 	psstructs "github.com/hashicorp/nomad/plugins/shared/structs"
 )
 
 const (
 	FilterConstraintHostVolumes                    = "missing compatible host volumes"
+	FilterConstraintHostVolumesLookupFailed        = "host volume lookup failed"
 	FilterConstraintCSIPluginTemplate              = "CSI plugin %s is missing from client %s"
 	FilterConstraintCSIPluginUnhealthyTemplate     = "CSI plugin %s is unhealthy on client %s"
 	FilterConstraintCSIPluginMaxVolumesTemplate    = "CSI plugin %s has the maximum number of volumes on client %s"
@@ -161,8 +164,6 @@ func (h *HostVolumeChecker) SetVolumes(allocName string, volumes map[string]*str
 			continue
 		}
 
-		// FIXME: if there's a sticky vol set, adjust this to look for an ID
-
 		if req.PerAlloc {
 			// provide a unique volume source per allocation
 			copied := req.Copy()
@@ -176,32 +177,33 @@ func (h *HostVolumeChecker) SetVolumes(allocName string, volumes map[string]*str
 }
 
 func (h *HostVolumeChecker) Feasible(candidate *structs.Node) bool {
-	if h.hasVolumes(candidate) {
+	ok, failure := h.hasVolumes(candidate)
+	if ok {
 		return true
 	}
 
-	h.ctx.Metrics().FilterNode(candidate, FilterConstraintHostVolumes)
+	h.ctx.Metrics().FilterNode(candidate, failure)
 	return false
 }
 
-func (h *HostVolumeChecker) hasVolumes(n *structs.Node) bool {
+func (h *HostVolumeChecker) hasVolumes(n *structs.Node) (bool, string) {
 	rLen := len(h.volumes)
 	hLen := len(n.HostVolumes)
 
 	// Fast path: Requested no volumes. No need to check further.
 	if rLen == 0 {
-		return true
+		return true, ""
 	}
 
 	// Fast path: Requesting more volumes than the node has, can't meet the criteria.
 	if rLen > hLen {
-		return false
+		return false, FilterConstraintHostVolumes
 	}
 
 	for source, requests := range h.volumes {
 		nodeVolume, ok := n.HostVolumes[source]
 		if !ok {
-			return false
+			return false, FilterConstraintHostVolumes
 		}
 
 		// If the volume supports being mounted as ReadWrite, we do not need to
@@ -210,16 +212,36 @@ func (h *HostVolumeChecker) hasVolumes(n *structs.Node) bool {
 			continue
 		}
 
+		ws := memdb.NewWatchSet()
 		// The Volume can only be mounted ReadOnly, validate that no requests for
 		// it are ReadWrite.
 		for _, req := range requests {
 			if !req.ReadOnly {
-				return false
+				return false, FilterConstraintHostVolumes
+			}
+
+			// Sticky volumes must always be paired with the right allocation;
+			// validate that this node has the right volume ID present.
+			if req.Sticky {
+				volumes, err := h.ctx.State().HostVolumesByNodeID(ws, n.ID, state.SortDefault)
+				if err != nil {
+					return false, FilterConstraintHostVolumesLookupFailed
+				}
+				for raw := volumes.Next(); raw != nil; raw = volumes.Next() {
+					vol := raw.(*structs.HostVolume)
+					if vol.Name != nodeVolume.Name {
+						continue
+					}
+					allocs, _ := h.ctx.ProposedAllocs(n.ID)
+					return slices.ContainsFunc(allocs, func(a *structs.Allocation) bool {
+						return !slices.Contains(a.VolumeIDs, vol.ID)
+					}), FilterConstraintHostVolumes
+				}
 			}
 		}
 	}
 
-	return true
+	return true, ""
 }
 
 type CSIVolumeChecker struct {
diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go
index 7e22070966f..d6416527f60 100644
--- a/scheduler/scheduler.go
+++ b/scheduler/scheduler.go
@@ -118,6 +118,12 @@ type State interface {
 	// CSIVolumeByID fetch CSI volumes, containing controller jobs
 	CSIVolumesByNodeID(memdb.WatchSet, string, string) (memdb.ResultIterator, error)
 
+	// HostVolumesByNodeID returns an iterator over all volumes of the node
+	HostVolumesByNodeID(memdb.WatchSet, string, state.SortOption) (memdb.ResultIterator, error)
+
+	// HostVolumesByName returns an iterator over all volumes by name
+	HostVolumesByName(memdb.WatchSet, string, string, state.SortOption) (memdb.ResultIterator, error)
+
 	// LatestIndex returns the greatest index value for all indexes.
 	LatestIndex() (uint64, error)
 }