func TestReplicatedScaleDown(t *testing.T) {
ctx := context.Background()
s := store.NewMemoryStore(nil)
assert.NotNil(t, s)
defer s.Close()
orchestrator := NewReplicatedOrchestrator(s)
defer orchestrator.Stop()
watch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateTask{})
defer cancel()
s1 := &api.Service{
ID: "id1",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "name1",
},
Mode: &api.ServiceSpec_Replicated{
Replicated: &api.ReplicatedService{
Replicas: 6,
},
},
},
}
err := s.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateService(tx, s1))
nodes := []*api.Node{
{
ID: "node1",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name1",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "node2",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name2",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "node3",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name3",
},
Availability: api.NodeAvailabilityActive,
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
}
for _, node := range nodes {
assert.NoError(t, store.CreateNode(tx, node))
}
// task1 is assigned to node1
// task2 - task3 are assigned to node2
// task4 - task6 are assigned to node3
// task7 is unassigned
tasks := []*api.Task{
{
ID: "task1",
Slot: 1,
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateStarting,
},
ServiceAnnotations: api.Annotations{
Name: "task1",
},
ServiceID: "id1",
NodeID: "node1",
},
{
ID: "task2",
Slot: 2,
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateRunning,
},
ServiceAnnotations: api.Annotations{
Name: "task2",
},
ServiceID: "id1",
NodeID: "node2",
},
{
ID: "task3",
Slot: 3,
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateRunning,
},
ServiceAnnotations: api.Annotations{
Name: "task3",
},
ServiceID: "id1",
NodeID: "node2",
},
{
ID: "task4",
Slot: 4,
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateRunning,
},
ServiceAnnotations: api.Annotations{
Name: "task4",
},
ServiceID: "id1",
NodeID: "node3",
},
{
ID: "task5",
Slot: 5,
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateRunning,
},
ServiceAnnotations: api.Annotations{
Name: "task5",
},
ServiceID: "id1",
NodeID: "node3",
},
{
ID: "task6",
Slot: 6,
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateRunning,
},
ServiceAnnotations: api.Annotations{
Name: "task6",
},
ServiceID: "id1",
NodeID: "node3",
},
{
ID: "task7",
Slot: 7,
DesiredState: api.TaskStateRunning,
Status: api.TaskStatus{
State: api.TaskStateNew,
},
ServiceAnnotations: api.Annotations{
Name: "task7",
},
ServiceID: "id1",
},
}
for _, task := range tasks {
assert.NoError(t, store.CreateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
// Start the orchestrator.
go func() {
assert.NoError(t, orchestrator.Run(ctx))
}()
// Replicas was set to 6, but we started with 7 tasks. task7 should
// be the one the orchestrator chose to shut down because it was not
// assigned yet.
observedShutdown := watchShutdownTask(t, watch)
assert.Equal(t, "task7", observedShutdown.ID)
// Now scale down to 2 instances.
err = s.Update(func(tx store.Tx) error {
s1.Spec.Mode = &api.ServiceSpec_Replicated{
Replicated: &api.ReplicatedService{
Replicas: 2,
},
}
assert.NoError(t, store.UpdateService(tx, s1))
return nil
})
assert.NoError(t, err)
// Tasks should be shut down in a way that balances the remaining tasks.
// node2 and node3 should be preferred over node1 because node1's task
// is not running yet.
shutdowns := make(map[string]int)
for i := 0; i != 4; i++ {
observedShutdown := watchShutdownTask(t, watch)
shutdowns[observedShutdown.NodeID]++
}
assert.Equal(t, 1, shutdowns["node1"])
assert.Equal(t, 1, shutdowns["node2"])
assert.Equal(t, 2, shutdowns["node3"])
// There should be remaining tasks on node2 and node3.
s.View(func(readTx store.ReadTx) {
tasks, err := store.FindTasks(readTx, store.ByDesiredState(api.TaskStateRunning))
require.NoError(t, err)
require.Len(t, tasks, 2)
if tasks[0].NodeID == "node2" {
assert.Equal(t, "node3", tasks[1].NodeID)
} else {
assert.Equal(t, "node3", tasks[0].NodeID)
assert.Equal(t, "node2", tasks[1].NodeID)
}
})
}
// ListTasks returns a list of all tasks.
func (s *Server) ListTasks(ctx context.Context, request *api.ListTasksRequest) (*api.ListTasksResponse, error) {
var (
tasks []*api.Task
err error
)
s.store.View(func(tx store.ReadTx) {
switch {
case request.Filters != nil && len(request.Filters.Names) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByName, request.Filters.Names))
case request.Filters != nil && len(request.Filters.NamePrefixes) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByNamePrefix, request.Filters.NamePrefixes))
case request.Filters != nil && len(request.Filters.IDPrefixes) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByIDPrefix, request.Filters.IDPrefixes))
case request.Filters != nil && len(request.Filters.ServiceIDs) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByServiceID, request.Filters.ServiceIDs))
case request.Filters != nil && len(request.Filters.NodeIDs) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByNodeID, request.Filters.NodeIDs))
case request.Filters != nil && len(request.Filters.DesiredStates) > 0:
filters := make([]store.By, 0, len(request.Filters.DesiredStates))
for _, v := range request.Filters.DesiredStates {
filters = append(filters, store.ByDesiredState(v))
}
tasks, err = store.FindTasks(tx, store.Or(filters...))
default:
tasks, err = store.FindTasks(tx, store.All)
}
})
if err != nil {
return nil, err
}
if request.Filters != nil {
tasks = filterTasks(tasks,
func(e *api.Task) bool {
return filterContains(naming.Task(e), request.Filters.Names)
},
func(e *api.Task) bool {
return filterContainsPrefix(naming.Task(e), request.Filters.NamePrefixes)
},
func(e *api.Task) bool {
return filterContainsPrefix(e.ID, request.Filters.IDPrefixes)
},
func(e *api.Task) bool {
return filterMatchLabels(e.ServiceAnnotations.Labels, request.Filters.Labels)
},
func(e *api.Task) bool {
return filterContains(e.ServiceID, request.Filters.ServiceIDs)
},
func(e *api.Task) bool {
return filterContains(e.NodeID, request.Filters.NodeIDs)
},
func(e *api.Task) bool {
if len(request.Filters.DesiredStates) == 0 {
return true
}
for _, c := range request.Filters.DesiredStates {
if c == e.DesiredState {
return true
}
}
return false
},
)
}
return &api.ListTasksResponse{
Tasks: tasks,
}, nil
}
// ListTasks returns a list of all tasks.
func (s *Server) ListTasks(ctx context.Context, request *api.ListTasksRequest) (*api.ListTasksResponse, error) {
var (
tasks []*api.Task
err error
)
s.store.View(func(tx store.ReadTx) {
switch {
case request.Filters != nil && len(request.Filters.Names) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByName, request.Filters.Names))
case request.Filters != nil && len(request.Filters.NamePrefixes) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByNamePrefix, request.Filters.NamePrefixes))
case request.Filters != nil && len(request.Filters.IDPrefixes) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByIDPrefix, request.Filters.IDPrefixes))
case request.Filters != nil && len(request.Filters.ServiceIDs) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByServiceID, request.Filters.ServiceIDs))
case request.Filters != nil && len(request.Filters.NodeIDs) > 0:
tasks, err = store.FindTasks(tx, buildFilters(store.ByNodeID, request.Filters.NodeIDs))
case request.Filters != nil && len(request.Filters.DesiredStates) > 0:
filters := make([]store.By, 0, len(request.Filters.DesiredStates))
for _, v := range request.Filters.DesiredStates {
filters = append(filters, store.ByDesiredState(v))
}
tasks, err = store.FindTasks(tx, store.Or(filters...))
default:
tasks, err = store.FindTasks(tx, store.All)
}
})
if err != nil {
return nil, err
}
if request.Filters != nil {
tasks = filterTasks(tasks,
func(e *api.Task) bool {
name := e.Annotations.Name
if name == "" {
// If Task name is not assigned then calculated name is used like before.
// This might be removed in the future.
name = fmt.Sprintf("%v.%v.%v", e.ServiceAnnotations.Name, e.Slot, e.ID)
}
return filterContains(name, request.Filters.Names)
},
func(e *api.Task) bool {
name := e.Annotations.Name
if name == "" {
// If Task name is not assigned then calculated name is used like before
// This might be removed in the future.
name = fmt.Sprintf("%v.%v.%v", e.ServiceAnnotations.Name, e.Slot, e.ID)
}
return filterContainsPrefix(name, request.Filters.NamePrefixes)
},
func(e *api.Task) bool {
return filterContainsPrefix(e.ID, request.Filters.IDPrefixes)
},
func(e *api.Task) bool {
return filterMatchLabels(e.ServiceAnnotations.Labels, request.Filters.Labels)
},
func(e *api.Task) bool {
return filterContains(e.ServiceID, request.Filters.ServiceIDs)
},
func(e *api.Task) bool {
return filterContains(e.NodeID, request.Filters.NodeIDs)
},
func(e *api.Task) bool {
if len(request.Filters.DesiredStates) == 0 {
return true
}
for _, c := range request.Filters.DesiredStates {
if c == e.DesiredState {
return true
}
}
return false
},
)
}
return &api.ListTasksResponse{
Tasks: tasks,
}, nil
}