Set MaxAllocations in client config

Add NodeAllocationTracker struct to Node struct
Evaluate MaxAllocations in AllocsFit function

Author: tehut

client/client.go

@@ -1628,7 +1628,14 @@ 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{
+				NodeMaxAllocs: newConfig.NodeMaxAllocs,
+			}
+		}
+	}
 	// 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, fmt.Errorf("plan exceeds max allocation")
+		}
+	}
 	reservedCores := map[uint16]struct{}{}
 	var coreOverlap bool
 

nomad/structs/funcs_test.go

@@ -716,6 +716,82 @@ func TestScoreFitBinPack(t *testing.T) {
 	}
 }
 
+func TestAllocsFit_MaxNodeAllocs(t *testing.T) {
+	ci.Parallel(t)
+	baseAlloc := &Allocation{
+		AllocatedResources: &AllocatedResources{
+			Tasks: map[string]*AllocatedTaskResources{
+				"web": {
+					Cpu: AllocatedCpuResources{
+						CpuShares:     1000,
+						ReservedCores: []uint16{},
+					},
+					Memory: AllocatedMemoryResources{
+						MemoryMB: 1024,
+					},
+				},
+			},
+			Shared: AllocatedSharedResources{
+				DiskMB: 5000,
+				Networks: Networks{
+					{
+						Mode:          "host",
+						IP:            "10.0.0.1",
+						ReservedPorts: []Port{{Label: "main", Value: 8000}},
+					},
+				},
+				Ports: AllocatedPorts{
+					{
+						Label:  "main",
+						Value:  8000,
+						HostIP: "10.0.0.1",
+					},
+				},
+			},
+		},
+	}
+
+	testCases := []struct {
+		name        string
+		allocations []*Allocation
+		expectErr   bool
+		maxAllocs   int
+	}{
+		{
+			name:        "happy_path",
+			allocations: []*Allocation{baseAlloc},
+			expectErr:   false,
+			maxAllocs:   2,
+		},
+		{
+			name:        "too many allocs",
+			allocations: []*Allocation{baseAlloc, baseAlloc, baseAlloc},
+			expectErr:   true,
+			maxAllocs:   2,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			n := node2k()
+			n.NodeAllocationTracker = &NodeAllocationTracker{false, tc.maxAllocs}
+			fit, dim, used, err := AllocsFit(n, tc.allocations, nil, false)
+			if !tc.expectErr {
+				must.NoError(t, err)
+				must.True(t, fit)
+				must.Eq(t, 1000, used.Flattened.Cpu.CpuShares)
+				must.Eq(t, 1024, used.Flattened.Memory.MemoryMB)
+			} else {
+				must.Error(t, err)
+				must.False(t, fit)
+				must.StrContains(t, dim, "max allocation exceeded")
+				must.StrContains(t, err.Error(), "plan exceeds max allocation")
+				must.Eq(t, 0, used.Flattened.Cpu.CpuShares)
+				must.Eq(t, 0, used.Flattened.Memory.MemoryMB)
+			}
+		})
+	}
+}
 func TestACLPolicyListHash(t *testing.T) {
 	ci.Parallel(t)
 

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,14 @@ 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 {
+	atMax         bool
+	NodeMaxAllocs int
+}
+
 // Resources is used to define the resources available
 // on a client
 type Resources struct {