can I just do all of this in AllocsFit?

Author: tehut

client/client.go

@@ -1628,7 +1628,17 @@ func (c *Client) setupNode() error {
 	if _, ok := node.Meta[envoy.DefaultTransparentProxyOutboundPortParam]; !ok {
 		node.Meta[envoy.DefaultTransparentProxyOutboundPortParam] = envoy.DefaultTransparentProxyOutboundPort
 	}
-
+	// Set NodeMaxAllocs before dynamic configuration is set
+	if node.NodeAllocationTracker == nil {
+		if newConfig.NodeMaxAllocs >= 1 {
+			node.NodeAllocationTracker = &structs.NodeAllocationTracker{
+				HasMaxAlloc:       true,
+				NodeAvailable:     true,
+				NodeMaxAllocs:     newConfig.NodeMaxAllocs,
+				NodeCurrentAllocs: 0,
+			}
+		}
+	}
 	// Since node.Meta will get dynamic metadata merged in, save static metadata
 	// here.
 	c.metaStatic = maps.Clone(node.Meta)

client/config/config.go

@@ -375,6 +375,10 @@ type Config struct {
 
 	// ExtraAllocHooks are run with other allocation hooks, mainly for testing.
 	ExtraAllocHooks []interfaces.RunnerHook
+
+	// NodeMaxAllocs is an optional field that sets the maximum number of
+	// allocations a node can be assigned. Defaults to 0 and ignored if unset.
+	NodeMaxAllocs int
 }
 
 type APIListenerRegistrar interface {

nomad/structs/funcs.go

@@ -141,7 +141,11 @@ func (a TerminalByNodeByName) Get(nodeID, name string) (*Allocation, bool) {
 func AllocsFit(node *Node, allocs []*Allocation, netIdx *NetworkIndex, checkDevices bool) (bool, string, *ComparableResources, error) {
 	// Compute the allocs' utilization from zero
 	used := new(ComparableResources)
-
+	if node.NodeAllocationTracker != nil {
+		if node.NodeAllocationTracker.NodeMaxAllocs <= len(allocs) {
+			return false, "max allocation exceeded", used, nil //not sure what used is here but everyone else is doing it
+		}
+	}
 	reservedCores := map[uint16]struct{}{}
 	var coreOverlap bool
 

nomad/structs/structs.go

@@ -2174,6 +2174,11 @@ type Node struct {
 	// Raft Indexes
 	CreateIndex uint64
 	ModifyIndex uint64
+
+	// NodeAllocationTracker holds NodeMaxAllocs value, if configured,
+	// and CurrentNodeAllocations to help the scheduler to block excess
+	// allocations.
+	NodeAllocationTracker *NodeAllocationTracker
 }
 
 // GetID is a helper for getting the ID when the object may be nil and is
@@ -2394,6 +2399,18 @@ type NodeStubFields struct {
 	OS        bool
 }
 
+// NodeAllocationTracker retains awareness of a client's NodeMaxAllocs
+// value, the current number of Allocations and whether the node can
+// accept new Allocations
+type NodeAllocationTracker struct {
+	HasMaxAlloc   bool
+	NodeAvailable bool
+	// TODO: review node concurrency, think allocations are either locked by
+	// default or "single threaded" by virtue of single scheduler.
+	NodeCurrentAllocs int // NOTE: Keeping track of this here might be too ambitious and unneeded
+	NodeMaxAllocs     int
+}
+
 // Resources is used to define the resources available
 // on a client
 type Resources struct {

scheduler/generic_sched.go

@@ -381,6 +381,10 @@ func (s *GenericScheduler) computeJobAllocs() error {
 	// nodes to lost, but only if the scheduler has already marked them
 	updateNonTerminalAllocsToLost(s.plan, tainted, allocs)
 
+	////TODO refactor into helper func
+	//for _, alloc := range allocs {
+
+	//}
 	reconciler := NewAllocReconciler(s.logger,
 		genericAllocUpdateFn(s.ctx, s.stack, s.eval.ID),
 		s.batch, s.eval.JobID, s.job, s.deployment, allocs, tainted, s.eval.ID,