github.com/emate/nomad@v0.8.2-wo-binpacking/scheduler/testing.go (about) 1 package scheduler 2 3 import ( 4 "fmt" 5 "sync" 6 "time" 7 8 memdb "github.com/hashicorp/go-memdb" 9 "github.com/hashicorp/nomad/helper/testlog" 10 "github.com/hashicorp/nomad/nomad/state" 11 "github.com/hashicorp/nomad/nomad/structs" 12 "github.com/mitchellh/go-testing-interface" 13 ) 14 15 // RejectPlan is used to always reject the entire plan and force a state refresh 16 type RejectPlan struct { 17 Harness *Harness 18 } 19 20 func (r *RejectPlan) SubmitPlan(*structs.Plan) (*structs.PlanResult, State, error) { 21 result := new(structs.PlanResult) 22 result.RefreshIndex = r.Harness.NextIndex() 23 return result, r.Harness.State, nil 24 } 25 26 func (r *RejectPlan) UpdateEval(eval *structs.Evaluation) error { 27 return nil 28 } 29 30 func (r *RejectPlan) CreateEval(*structs.Evaluation) error { 31 return nil 32 } 33 34 func (r *RejectPlan) ReblockEval(*structs.Evaluation) error { 35 return nil 36 } 37 38 // Harness is a lightweight testing harness for schedulers. It manages a state 39 // store copy and provides the planner interface. It can be extended for various 40 // testing uses or for invoking the scheduler without side effects. 41 type Harness struct { 42 t testing.T 43 State *state.StateStore 44 45 Planner Planner 46 planLock sync.Mutex 47 48 Plans []*structs.Plan 49 Evals []*structs.Evaluation 50 CreateEvals []*structs.Evaluation 51 ReblockEvals []*structs.Evaluation 52 53 nextIndex uint64 54 nextIndexLock sync.Mutex 55 } 56 57 // NewHarness is used to make a new testing harness 58 func NewHarness(t testing.T) *Harness { 59 state := state.TestStateStore(t) 60 h := &Harness{ 61 t: t, 62 State: state, 63 nextIndex: 1, 64 } 65 return h 66 } 67 68 // NewHarnessWithState creates a new harness with the given state for testing 69 // purposes. 70 func NewHarnessWithState(t testing.T, state *state.StateStore) *Harness { 71 return &Harness{ 72 t: t, 73 State: state, 74 nextIndex: 1, 75 } 76 } 77 78 // SubmitPlan is used to handle plan submission 79 func (h *Harness) SubmitPlan(plan *structs.Plan) (*structs.PlanResult, State, error) { 80 // Ensure sequential plan application 81 h.planLock.Lock() 82 defer h.planLock.Unlock() 83 84 // Store the plan 85 h.Plans = append(h.Plans, plan) 86 87 // Check for custom planner 88 if h.Planner != nil { 89 return h.Planner.SubmitPlan(plan) 90 } 91 92 // Get the index 93 index := h.NextIndex() 94 95 // Prepare the result 96 result := new(structs.PlanResult) 97 result.NodeUpdate = plan.NodeUpdate 98 result.NodeAllocation = plan.NodeAllocation 99 result.AllocIndex = index 100 101 // Flatten evicts and allocs 102 var allocs []*structs.Allocation 103 for _, updateList := range plan.NodeUpdate { 104 allocs = append(allocs, updateList...) 105 } 106 for _, allocList := range plan.NodeAllocation { 107 allocs = append(allocs, allocList...) 108 } 109 110 // Set the time the alloc was applied for the first time. This can be used 111 // to approximate the scheduling time. 112 now := time.Now().UTC().UnixNano() 113 for _, alloc := range allocs { 114 if alloc.CreateTime == 0 { 115 alloc.CreateTime = now 116 } 117 } 118 119 // Setup the update request 120 req := structs.ApplyPlanResultsRequest{ 121 AllocUpdateRequest: structs.AllocUpdateRequest{ 122 Job: plan.Job, 123 Alloc: allocs, 124 }, 125 Deployment: plan.Deployment, 126 DeploymentUpdates: plan.DeploymentUpdates, 127 EvalID: plan.EvalID, 128 } 129 130 // Apply the full plan 131 err := h.State.UpsertPlanResults(index, &req) 132 return result, nil, err 133 } 134 135 func (h *Harness) UpdateEval(eval *structs.Evaluation) error { 136 // Ensure sequential plan application 137 h.planLock.Lock() 138 defer h.planLock.Unlock() 139 140 // Store the eval 141 h.Evals = append(h.Evals, eval) 142 143 // Check for custom planner 144 if h.Planner != nil { 145 return h.Planner.UpdateEval(eval) 146 } 147 return nil 148 } 149 150 func (h *Harness) CreateEval(eval *structs.Evaluation) error { 151 // Ensure sequential plan application 152 h.planLock.Lock() 153 defer h.planLock.Unlock() 154 155 // Store the eval 156 h.CreateEvals = append(h.CreateEvals, eval) 157 158 // Check for custom planner 159 if h.Planner != nil { 160 return h.Planner.CreateEval(eval) 161 } 162 return nil 163 } 164 165 func (h *Harness) ReblockEval(eval *structs.Evaluation) error { 166 // Ensure sequential plan application 167 h.planLock.Lock() 168 defer h.planLock.Unlock() 169 170 // Check that the evaluation was already blocked. 171 ws := memdb.NewWatchSet() 172 old, err := h.State.EvalByID(ws, eval.ID) 173 if err != nil { 174 return err 175 } 176 177 if old == nil { 178 return fmt.Errorf("evaluation does not exist to be reblocked") 179 } 180 if old.Status != structs.EvalStatusBlocked { 181 return fmt.Errorf("evaluation %q is not already in a blocked state", old.ID) 182 } 183 184 h.ReblockEvals = append(h.ReblockEvals, eval) 185 return nil 186 } 187 188 // NextIndex returns the next index 189 func (h *Harness) NextIndex() uint64 { 190 h.nextIndexLock.Lock() 191 defer h.nextIndexLock.Unlock() 192 idx := h.nextIndex 193 h.nextIndex += 1 194 return idx 195 } 196 197 // Snapshot is used to snapshot the current state 198 func (h *Harness) Snapshot() State { 199 snap, _ := h.State.Snapshot() 200 return snap 201 } 202 203 // Scheduler is used to return a new scheduler from 204 // a snapshot of current state using the harness for planning. 205 func (h *Harness) Scheduler(factory Factory) Scheduler { 206 logger := testlog.Logger(h.t) 207 return factory(logger, h.Snapshot(), h) 208 } 209 210 // Process is used to process an evaluation given a factory 211 // function to create the scheduler 212 func (h *Harness) Process(factory Factory, eval *structs.Evaluation) error { 213 sched := h.Scheduler(factory) 214 return sched.Process(eval) 215 } 216 217 func (h *Harness) AssertEvalStatus(t testing.T, state string) { 218 if len(h.Evals) != 1 { 219 t.Fatalf("bad: %#v", h.Evals) 220 } 221 update := h.Evals[0] 222 223 if update.Status != state { 224 t.Fatalf("bad: %#v", update) 225 } 226 }