func TestNodeAvailability(t *testing.T) {
store := store.NewMemoryStore(nil)
assert.NotNil(t, store)
defer store.Close()
watch, cancel := state.Watch(store.WatchQueue())
defer cancel()
SetupCluster(t, store, watch)
node1.Status.State = api.NodeStatus_READY
node1.Spec.Availability = api.NodeAvailabilityActive
// set node1 to drain
updateNodeAvailability(t, store, node1, api.NodeAvailabilityDrain)
// task should be set to dead
observedTask1 := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, observedTask1.ServiceAnnotations.Name, "name1")
assert.Equal(t, observedTask1.NodeID, "id1")
// set node1 to active
updateNodeAvailability(t, store, node1, api.NodeAvailabilityActive)
// task should be added back
observedTask2 := testutils.WatchTaskCreate(t, watch)
assert.Equal(t, observedTask2.Status.State, api.TaskStateNew)
assert.Equal(t, observedTask2.ServiceAnnotations.Name, "name1")
assert.Equal(t, observedTask2.NodeID, "id1")
}
func TestDeleteNode(t *testing.T) {
store := store.NewMemoryStore(nil)
assert.NotNil(t, store)
defer store.Close()
watch, cancel := state.Watch(store.WatchQueue())
defer cancel()
SetupCluster(t, store, watch)
deleteNode(t, store, node1)
// task should be set to dead
observedTask := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, observedTask.ServiceAnnotations.Name, "name1")
assert.Equal(t, observedTask.NodeID, "id1")
}
func TestDrain(t *testing.T) {
ctx := context.Background()
initialService := &api.Service{
ID: "id1",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "name1",
},
Task: api.TaskSpec{
Runtime: &api.TaskSpec_Container{
Container: &api.ContainerSpec{},
},
Restart: &api.RestartPolicy{
Condition: api.RestartOnNone,
},
},
Mode: &api.ServiceSpec_Replicated{
Replicated: &api.ReplicatedService{
Replicas: 6,
},
},
},
}
initialNodeSet := []*api.Node{
{
ID: "id1",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name1",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id2",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name2",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_DOWN,
},
},
// We should NOT kick out tasks on UNKNOWN nodes.
{
ID: "id3",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name3",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_UNKNOWN,
},
},
{
ID: "id4",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name4",
},
Availability: api.NodeAvailabilityPause,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id5",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name5",
},
Availability: api.NodeAvailabilityDrain,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
}
initialTaskSet := []*api.Task{
// Task not assigned to any node
{
ID: "id0",
DesiredState: api.TaskStateRunning,
Spec: initialService.Spec.Task,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
Slot: 1,
ServiceAnnotations: api.Annotations{
Name: "name0",
},
ServiceID: "id1",
},
// Tasks assigned to the nodes defined above
{
ID: "id1",
DesiredState: api.TaskStateRunning,
Spec: initialService.Spec.Task,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
Slot: 2,
ServiceAnnotations: api.Annotations{
Name: "name1",
},
ServiceID: "id1",
NodeID: "id1",
},
{
ID: "id2",
DesiredState: api.TaskStateRunning,
Spec: initialService.Spec.Task,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
Slot: 3,
ServiceAnnotations: api.Annotations{
Name: "name2",
},
ServiceID: "id1",
NodeID: "id2",
},
{
ID: "id3",
DesiredState: api.TaskStateRunning,
Spec: initialService.Spec.Task,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
Slot: 4,
ServiceAnnotations: api.Annotations{
Name: "name3",
},
ServiceID: "id1",
NodeID: "id3",
},
{
ID: "id4",
DesiredState: api.TaskStateRunning,
Spec: initialService.Spec.Task,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
Slot: 5,
ServiceAnnotations: api.Annotations{
Name: "name4",
},
ServiceID: "id1",
NodeID: "id4",
},
{
ID: "id5",
DesiredState: api.TaskStateRunning,
Spec: initialService.Spec.Task,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
Slot: 6,
ServiceAnnotations: api.Annotations{
Name: "name5",
},
ServiceID: "id1",
NodeID: "id5",
},
}
s := store.NewMemoryStore(nil)
assert.NotNil(t, s)
defer s.Close()
err := s.Update(func(tx store.Tx) error {
// Prepopulate service
assert.NoError(t, store.CreateService(tx, initialService))
// Prepoulate nodes
for _, n := range initialNodeSet {
assert.NoError(t, store.CreateNode(tx, n))
}
// Prepopulate tasks
for _, task := range initialTaskSet {
assert.NoError(t, store.CreateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
watch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateTask{})
defer cancel()
orchestrator := NewReplicatedOrchestrator(s)
defer orchestrator.Stop()
go func() {
assert.NoError(t, orchestrator.Run(ctx))
}()
// id2 and id5 should be killed immediately
deletion1 := testutils.WatchShutdownTask(t, watch)
deletion2 := testutils.WatchShutdownTask(t, watch)
assert.Regexp(t, "id(2|5)", deletion1.ID)
assert.Regexp(t, "id(2|5)", deletion1.NodeID)
assert.Regexp(t, "id(2|5)", deletion2.ID)
assert.Regexp(t, "id(2|5)", deletion2.NodeID)
// Create a new task, assigned to node id2
err = s.Update(func(tx store.Tx) error {
task := initialTaskSet[2].Copy()
task.ID = "newtask"
task.NodeID = "id2"
assert.NoError(t, store.CreateTask(tx, task))
return nil
})
assert.NoError(t, err)
deletion3 := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, "newtask", deletion3.ID)
assert.Equal(t, "id2", deletion3.NodeID)
// Set node id4 to the DRAINED state
err = s.Update(func(tx store.Tx) error {
n := initialNodeSet[3].Copy()
n.Spec.Availability = api.NodeAvailabilityDrain
assert.NoError(t, store.UpdateNode(tx, n))
return nil
})
assert.NoError(t, err)
deletion4 := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, "id4", deletion4.ID)
assert.Equal(t, "id4", deletion4.NodeID)
// Delete node id1
err = s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteNode(tx, "id1"))
return nil
})
assert.NoError(t, err)
deletion5 := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, "id1", deletion5.ID)
assert.Equal(t, "id1", deletion5.NodeID)
}
func TestConstraintEnforcer(t *testing.T) {
nodes := []*api.Node{
{
ID: "id1",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name1",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
Role: api.NodeRoleWorker,
},
{
ID: "id2",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name2",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
Description: &api.NodeDescription{
Resources: &api.Resources{
NanoCPUs: 1e9,
MemoryBytes: 1e9,
},
},
},
}
tasks := []*api.Task{
{
ID: "id0",
DesiredState: api.TaskStateRunning,
Spec: api.TaskSpec{
Placement: &api.Placement{
Constraints: []string{"node.role == manager"},
},
},
Status: api.TaskStatus{
State: api.TaskStateNew,
},
NodeID: "id1",
},
{
ID: "id1",
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
NodeID: "id1",
},
{
ID: "id2",
DesiredState: api.TaskStateRunning,
Spec: api.TaskSpec{
Placement: &api.Placement{
Constraints: []string{"node.role == worker"},
},
},
Status: api.TaskStatus{
State: api.TaskStateRunning,
},
NodeID: "id1",
},
{
ID: "id3",
DesiredState: api.TaskStateNew,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
NodeID: "id2",
},
{
ID: "id4",
DesiredState: api.TaskStateReady,
Spec: api.TaskSpec{
Resources: &api.ResourceRequirements{
Reservations: &api.Resources{
MemoryBytes: 9e8,
},
},
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
NodeID: "id2",
},
}
s := store.NewMemoryStore(nil)
assert.NotNil(t, s)
defer s.Close()
err := s.Update(func(tx store.Tx) error {
// Prepoulate nodes
for _, n := range nodes {
assert.NoError(t, store.CreateNode(tx, n))
}
// Prepopulate tasks
for _, task := range tasks {
assert.NoError(t, store.CreateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
watch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateTask{})
defer cancel()
constraintEnforcer := New(s)
defer constraintEnforcer.Stop()
go constraintEnforcer.Run()
// id0 should be killed immediately
shutdown1 := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, "id0", shutdown1.ID)
// Change node id1 to a manager
err = s.Update(func(tx store.Tx) error {
node := store.GetNode(tx, "id1")
if node == nil {
t.Fatal("could not get node id1")
}
node.Role = api.NodeRoleManager
assert.NoError(t, store.UpdateNode(tx, node))
return nil
})
assert.NoError(t, err)
shutdown2 := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, "id2", shutdown2.ID)
// Change resources on node id2
err = s.Update(func(tx store.Tx) error {
node := store.GetNode(tx, "id2")
if node == nil {
t.Fatal("could not get node id2")
}
node.Description.Resources.MemoryBytes = 5e8
assert.NoError(t, store.UpdateNode(tx, node))
return nil
})
assert.NoError(t, err)
shutdown3 := testutils.WatchShutdownTask(t, watch)
assert.Equal(t, "id4", shutdown3.ID)
}