Merge pull request #822 from hashicorp/f-filteralloc

Use compound index on node allocations to include terminal status

nomad/plan_apply.go

@@ -256,15 +256,12 @@ func evaluateNodePlan(snap *state.StateSnapshot, plan *structs.Plan, nodeID stri
 		return false, nil
 	}
 
-	// Get the existing allocations
+	// Get the existing allocations that are non-terminal
-	existingAlloc, err := snap.AllocsByNode(nodeID)
+	existingAlloc, err := snap.AllocsByNodeTerminal(nodeID, false)
 	if err != nil {
 		return false, fmt.Errorf("failed to get existing allocations for '%s': %v", nodeID, err)
 	}
 
-	// Filter on alloc state
-	existingAlloc = structs.FilterTerminalAllocs(existingAlloc)
-
 	// Determine the proposed allocation by first removing allocations
 	// that are planned evictions and adding the new allocations.
 	proposed := existingAlloc

nomad/state/schema.go

@@ -222,10 +222,28 @@ func allocTableSchema() *memdb.TableSchema {
 				Name:         "node",
 				AllowMissing: true, // Missing is allow for failed allocations
 				Unique:       false,
-				Indexer: &memdb.StringFieldIndex{
+				Indexer: &memdb.CompoundIndex{
+					Indexes: []memdb.Indexer{
+						&memdb.StringFieldIndex{
 							Field:     "NodeID",
 							Lowercase: true,
 						},
+
+						// Conditional indexer on if allocation is terminal
+						&memdb.ConditionalIndex{
+							Conditional: func(obj interface{}) (bool, error) {
+								// Cast to allocation
+								alloc, ok := obj.(*structs.Allocation)
+								if !ok {
+									return false, fmt.Errorf("wrong type, got %t should be Allocation", obj)
+								}
+
+								// Check if the allocation is terminal
+								return alloc.TerminalStatus(), nil
+							},
+						},
+					},
+				},
 			},
 
 			// Job index is used to lookup allocations by job

nomad/state/state_store.go

@@ -865,8 +865,30 @@ func (s *StateStore) AllocsByIDPrefix(id string) (memdb.ResultIterator, error) {
 func (s *StateStore) AllocsByNode(node string) ([]*structs.Allocation, error) {
 	txn := s.db.Txn(false)
 
+	// Get an iterator over the node allocations, using only the
+	// node prefix which ignores the terminal status
+	iter, err := txn.Get("allocs", "node_prefix", node)
+	if err != nil {
+		return nil, err
+	}
+
+	var out []*structs.Allocation
+	for {
+		raw := iter.Next()
+		if raw == nil {
+			break
+		}
+		out = append(out, raw.(*structs.Allocation))
+	}
+	return out, nil
+}
+
+// AllocsByNode returns all the allocations by node and terminal status
+func (s *StateStore) AllocsByNodeTerminal(node string, terminal bool) ([]*structs.Allocation, error) {
+	txn := s.db.Txn(false)
+
 	// Get an iterator over the node allocations
-	iter, err := txn.Get("allocs", "node", node)
+	iter, err := txn.Get("allocs", "node", node, terminal)
 	if err != nil {
 		return nil, err
 	}

nomad/state/state_store_test.go

@@ -1581,6 +1581,54 @@ func TestStateStore_AllocsByNode(t *testing.T) {
 	}
 }
 
+func TestStateStore_AllocsByNodeTerminal(t *testing.T) {
+	state := testStateStore(t)
+	var allocs, term, nonterm []*structs.Allocation
+
+	for i := 0; i < 10; i++ {
+		alloc := mock.Alloc()
+		alloc.NodeID = "foo"
+		if i%2 == 0 {
+			alloc.DesiredStatus = structs.AllocDesiredStatusStop
+			term = append(term, alloc)
+		} else {
+			nonterm = append(nonterm, alloc)
+		}
+		allocs = append(allocs, alloc)
+	}
+
+	err := state.UpsertAllocs(1000, allocs)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Verify the terminal allocs
+	out, err := state.AllocsByNodeTerminal("foo", true)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	sort.Sort(AllocIDSort(term))
+	sort.Sort(AllocIDSort(out))
+
+	if !reflect.DeepEqual(term, out) {
+		t.Fatalf("bad: %#v %#v", term, out)
+	}
+
+	// Verify the non-terminal allocs
+	out, err = state.AllocsByNodeTerminal("foo", false)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	sort.Sort(AllocIDSort(nonterm))
+	sort.Sort(AllocIDSort(out))
+
+	if !reflect.DeepEqual(nonterm, out) {
+		t.Fatalf("bad: %#v %#v", nonterm, out)
+	}
+}
+
 func TestStateStore_AllocsByJob(t *testing.T) {
 	state := testStateStore(t)
 	var allocs []*structs.Allocation

scheduler/context.go

@@ -107,15 +107,12 @@ func (e *EvalContext) Reset() {
 }
 
 func (e *EvalContext) ProposedAllocs(nodeID string) ([]*structs.Allocation, error) {
-	// Get the existing allocations
+	// Get the existing allocations that are non-terminal
-	existingAlloc, err := e.state.AllocsByNode(nodeID)
+	existingAlloc, err := e.state.AllocsByNodeTerminal(nodeID, false)
 	if err != nil {
 		return nil, err
 	}
 
-	// Filter on alloc state
-	existingAlloc = structs.FilterTerminalAllocs(existingAlloc)
-
 	// Determine the proposed allocation by first removing allocations
 	// that are planned evictions and adding the new allocations.
 	proposed := existingAlloc

scheduler/scheduler.go

@@ -63,6 +63,9 @@ type State interface {
 	// AllocsByNode returns all the allocations by node
 	AllocsByNode(node string) ([]*structs.Allocation, error)
 
+	// AllocsByNodeTerminal returns all the allocations by node filtering by terminal status
+	AllocsByNodeTerminal(node string, terminal bool) ([]*structs.Allocation, error)
+
 	// GetNodeByID is used to lookup a node by ID
 	NodeByID(nodeID string) (*structs.Node, error)