func TestOldTasksBatch(t *testing.T) {
gd, err := startDispatcher(DefaultConfig())
assert.NoError(t, err)
defer gd.Close()
var expectedSessionID string
var nodeID string
{
stream, err := gd.Clients[0].Session(context.Background(), &api.SessionRequest{})
assert.NoError(t, err)
defer stream.CloseSend()
resp, err := stream.Recv()
assert.NoError(t, err)
assert.NotEmpty(t, resp.SessionID)
expectedSessionID = resp.SessionID
nodeID = resp.Node.ID
}
testTask1 := &api.Task{
NodeID: nodeID,
ID: "testTask1",
Status: api.TaskStatus{State: api.TaskStateAssigned},
}
testTask2 := &api.Task{
NodeID: nodeID,
ID: "testTask2",
Status: api.TaskStatus{State: api.TaskStateAssigned},
}
stream, err := gd.Clients[0].Tasks(context.Background(), &api.TasksRequest{SessionID: expectedSessionID})
assert.NoError(t, err)
resp, err := stream.Recv()
assert.NoError(t, err)
// initially no tasks
assert.Equal(t, 0, len(resp.Tasks))
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateTask(tx, testTask1))
assert.NoError(t, store.CreateTask(tx, testTask2))
return nil
})
assert.NoError(t, err)
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteTask(tx, testTask1.ID))
assert.NoError(t, store.DeleteTask(tx, testTask2.ID))
return nil
})
assert.NoError(t, err)
resp, err = stream.Recv()
assert.NoError(t, err)
// all tasks have been deleted
assert.Equal(t, len(resp.Tasks), 0)
}
func deleteServiceTasks(ctx context.Context, s *store.MemoryStore, service *api.Service) {
var (
tasks []*api.Task
err error
)
s.View(func(tx store.ReadTx) {
tasks, err = store.FindTasks(tx, store.ByServiceID(service.ID))
})
if err != nil {
log.G(ctx).WithError(err).Errorf("failed to list tasks")
return
}
_, err = s.Batch(func(batch *store.Batch) error {
for _, t := range tasks {
err := batch.Update(func(tx store.Tx) error {
if err := store.DeleteTask(tx, t.ID); err != nil {
log.G(ctx).WithError(err).Errorf("failed to delete task")
}
return nil
})
if err != nil {
return err
}
}
return nil
})
if err != nil {
log.G(ctx).WithError(err).Errorf("task search transaction failed")
}
}
// DetachNetwork allows the node to request the release of
// the resources associated to the network attachment.
// - Returns `InvalidArgument` if attachment ID is not provided.
// - Returns `NotFound` if the attachment is not found.
// - Returns an error if the deletion fails.
func (ra *ResourceAllocator) DetachNetwork(ctx context.Context, request *api.DetachNetworkRequest) (*api.DetachNetworkResponse, error) {
if request.AttachmentID == "" {
return nil, grpc.Errorf(codes.InvalidArgument, errInvalidArgument.Error())
}
nodeInfo, err := ca.RemoteNode(ctx)
if err != nil {
return nil, err
}
if err := ra.store.Update(func(tx store.Tx) error {
t := store.GetTask(tx, request.AttachmentID)
if t == nil {
return grpc.Errorf(codes.NotFound, "attachment %s not found", request.AttachmentID)
}
if t.NodeID != nodeInfo.NodeID {
return grpc.Errorf(codes.PermissionDenied, "attachment %s doesn't belong to this node", request.AttachmentID)
}
return store.DeleteTask(tx, request.AttachmentID)
}); err != nil {
return nil, err
}
return &api.DetachNetworkResponse{}, nil
}
func (r *ReplicatedOrchestrator) deleteTask(ctx context.Context, batch *store.Batch, t *api.Task) {
err := batch.Update(func(tx store.Tx) error {
return store.DeleteTask(tx, t.ID)
})
if err != nil {
log.G(ctx).WithError(err).Errorf("deleting task %s failed", t.ID)
}
}
// RemoveTask removes a Task referenced by TaskID.
// - Returns `InvalidArgument` if TaskID is not provided.
// - Returns `NotFound` if the Task is not found.
// - Returns an error if the deletion fails.
func (s *Server) RemoveTask(ctx context.Context, request *api.RemoveTaskRequest) (*api.RemoveTaskResponse, error) {
if request.TaskID == "" {
return nil, grpc.Errorf(codes.InvalidArgument, errInvalidArgument.Error())
}
err := s.store.Update(func(tx store.Tx) error {
return store.DeleteTask(tx, request.TaskID)
})
if err != nil {
if err == store.ErrNotExist {
return nil, grpc.Errorf(codes.NotFound, "task %s not found", request.TaskID)
}
return nil, err
}
return &api.RemoveTaskResponse{}, nil
}
func (r *Orchestrator) initTasks(ctx context.Context, readTx store.ReadTx) error {
tasks, err := store.FindTasks(readTx, store.All)
if err != nil {
return err
}
for _, t := range tasks {
if t.NodeID != "" {
n := store.GetNode(readTx, t.NodeID)
if invalidNode(n) && t.Status.State <= api.TaskStateRunning && t.DesiredState <= api.TaskStateRunning {
r.restartTasks[t.ID] = struct{}{}
}
}
}
_, err = r.store.Batch(func(batch *store.Batch) error {
for _, t := range tasks {
if t.ServiceID == "" {
continue
}
// TODO(aluzzardi): We should NOT retrieve the service here.
service := store.GetService(readTx, t.ServiceID)
if service == nil {
// Service was deleted
err := batch.Update(func(tx store.Tx) error {
return store.DeleteTask(tx, t.ID)
})
if err != nil {
log.G(ctx).WithError(err).Error("failed to set task desired state to dead")
}
continue
}
// TODO(aluzzardi): This is shady. We should have a more generic condition.
if t.DesiredState != api.TaskStateReady || !orchestrator.IsReplicatedService(service) {
continue
}
restartDelay := orchestrator.DefaultRestartDelay
if t.Spec.Restart != nil && t.Spec.Restart.Delay != nil {
var err error
restartDelay, err = gogotypes.DurationFromProto(t.Spec.Restart.Delay)
if err != nil {
log.G(ctx).WithError(err).Error("invalid restart delay")
restartDelay = orchestrator.DefaultRestartDelay
}
}
if restartDelay != 0 {
timestamp, err := gogotypes.TimestampFromProto(t.Status.Timestamp)
if err == nil {
restartTime := timestamp.Add(restartDelay)
calculatedRestartDelay := restartTime.Sub(time.Now())
if calculatedRestartDelay < restartDelay {
restartDelay = calculatedRestartDelay
}
if restartDelay > 0 {
_ = batch.Update(func(tx store.Tx) error {
t := store.GetTask(tx, t.ID)
// TODO(aluzzardi): This is shady as well. We should have a more generic condition.
if t == nil || t.DesiredState != api.TaskStateReady {
return nil
}
r.restarts.DelayStart(ctx, tx, nil, t.ID, restartDelay, true)
return nil
})
continue
}
} else {
log.G(ctx).WithError(err).Error("invalid status timestamp")
}
}
// Start now
err := batch.Update(func(tx store.Tx) error {
return r.restarts.StartNow(tx, t.ID)
})
if err != nil {
log.G(ctx).WithError(err).WithField("task.id", t.ID).Error("moving task out of delayed state failed")
}
}
return nil
})
return err
}
func TestAllocator(t *testing.T) {
s := store.NewMemoryStore(nil)
assert.NotNil(t, s)
defer s.Close()
a, err := New(s)
assert.NoError(t, err)
assert.NotNil(t, a)
// Try adding some objects to store before allocator is started
assert.NoError(t, s.Update(func(tx store.Tx) error {
n1 := &api.Network{
ID: "testID1",
Spec: api.NetworkSpec{
Annotations: api.Annotations{
Name: "test1",
},
},
}
assert.NoError(t, store.CreateNetwork(tx, n1))
s1 := &api.Service{
ID: "testServiceID1",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "service1",
},
Task: api.TaskSpec{
Networks: []*api.NetworkAttachmentConfig{
{
Target: "testID1",
},
},
},
Endpoint: &api.EndpointSpec{},
},
}
assert.NoError(t, store.CreateService(tx, s1))
t1 := &api.Task{
ID: "testTaskID1",
Status: api.TaskStatus{
State: api.TaskStateNew,
},
Networks: []*api.NetworkAttachment{
{
Network: n1,
},
},
}
assert.NoError(t, store.CreateTask(tx, t1))
return nil
}))
netWatch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateNetwork{}, state.EventDeleteNetwork{})
defer cancel()
taskWatch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateTask{}, state.EventDeleteTask{})
defer cancel()
serviceWatch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateService{}, state.EventDeleteService{})
defer cancel()
// Start allocator
go func() {
assert.NoError(t, a.Run(context.Background()))
}()
// Now verify if we get network and tasks updated properly
watchNetwork(t, netWatch, false, isValidNetwork)
watchTask(t, s, taskWatch, false, isValidTask)
watchService(t, serviceWatch, false, nil)
// Add new networks/tasks/services after allocator is started.
assert.NoError(t, s.Update(func(tx store.Tx) error {
n2 := &api.Network{
ID: "testID2",
Spec: api.NetworkSpec{
Annotations: api.Annotations{
Name: "test2",
},
},
}
assert.NoError(t, store.CreateNetwork(tx, n2))
return nil
}))
watchNetwork(t, netWatch, false, isValidNetwork)
assert.NoError(t, s.Update(func(tx store.Tx) error {
s2 := &api.Service{
ID: "testServiceID2",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "service2",
},
Networks: []*api.NetworkAttachmentConfig{
{
Target: "testID2",
},
},
Endpoint: &api.EndpointSpec{},
},
}
assert.NoError(t, store.CreateService(tx, s2))
return nil
}))
watchService(t, serviceWatch, false, nil)
assert.NoError(t, s.Update(func(tx store.Tx) error {
t2 := &api.Task{
ID: "testTaskID2",
Status: api.TaskStatus{
State: api.TaskStateNew,
},
ServiceID: "testServiceID2",
DesiredState: api.TaskStateRunning,
}
assert.NoError(t, store.CreateTask(tx, t2))
return nil
}))
watchTask(t, s, taskWatch, false, isValidTask)
// Now try adding a task which depends on a network before adding the network.
n3 := &api.Network{
ID: "testID3",
Spec: api.NetworkSpec{
Annotations: api.Annotations{
Name: "test3",
},
},
}
assert.NoError(t, s.Update(func(tx store.Tx) error {
t3 := &api.Task{
ID: "testTaskID3",
Status: api.TaskStatus{
State: api.TaskStateNew,
},
DesiredState: api.TaskStateRunning,
Networks: []*api.NetworkAttachment{
{
Network: n3,
},
},
}
assert.NoError(t, store.CreateTask(tx, t3))
return nil
}))
// Wait for a little bit of time before adding network just to
// test network is not available while task allocation is
// going through
time.Sleep(10 * time.Millisecond)
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateNetwork(tx, n3))
return nil
}))
watchNetwork(t, netWatch, false, isValidNetwork)
watchTask(t, s, taskWatch, false, isValidTask)
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteTask(tx, "testTaskID3"))
return nil
}))
watchTask(t, s, taskWatch, false, isValidTask)
assert.NoError(t, s.Update(func(tx store.Tx) error {
t5 := &api.Task{
ID: "testTaskID5",
Spec: api.TaskSpec{
Networks: []*api.NetworkAttachmentConfig{
{
Target: "testID2",
},
},
},
Status: api.TaskStatus{
State: api.TaskStateNew,
},
DesiredState: api.TaskStateRunning,
ServiceID: "testServiceID2",
}
assert.NoError(t, store.CreateTask(tx, t5))
return nil
}))
watchTask(t, s, taskWatch, false, isValidTask)
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteNetwork(tx, "testID3"))
return nil
}))
watchNetwork(t, netWatch, false, isValidNetwork)
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteService(tx, "testServiceID2"))
return nil
}))
watchService(t, serviceWatch, false, nil)
// Try to create a task with no network attachments and test
// that it moves to ALLOCATED state.
assert.NoError(t, s.Update(func(tx store.Tx) error {
t4 := &api.Task{
ID: "testTaskID4",
Status: api.TaskStatus{
State: api.TaskStateNew,
},
DesiredState: api.TaskStateRunning,
}
assert.NoError(t, store.CreateTask(tx, t4))
return nil
}))
watchTask(t, s, taskWatch, false, isValidTask)
assert.NoError(t, s.Update(func(tx store.Tx) error {
n2 := store.GetNetwork(tx, "testID2")
require.NotEqual(t, nil, n2)
assert.NoError(t, store.UpdateNetwork(tx, n2))
return nil
}))
watchNetwork(t, netWatch, false, isValidNetwork)
watchNetwork(t, netWatch, true, nil)
// Try updating task which is already allocated
assert.NoError(t, s.Update(func(tx store.Tx) error {
t2 := store.GetTask(tx, "testTaskID2")
require.NotEqual(t, nil, t2)
assert.NoError(t, store.UpdateTask(tx, t2))
return nil
}))
watchTask(t, s, taskWatch, false, isValidTask)
watchTask(t, s, taskWatch, true, nil)
// Try adding networks with conflicting network resources and
// add task which attaches to a network which gets allocated
// later and verify if task reconciles and moves to ALLOCATED.
n4 := &api.Network{
ID: "testID4",
Spec: api.NetworkSpec{
Annotations: api.Annotations{
Name: "test4",
},
DriverConfig: &api.Driver{
Name: "overlay",
Options: map[string]string{
"com.docker.network.driver.overlay.vxlanid_list": "328",
},
},
},
}
n5 := n4.Copy()
n5.ID = "testID5"
n5.Spec.Annotations.Name = "test5"
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateNetwork(tx, n4))
return nil
}))
watchNetwork(t, netWatch, false, isValidNetwork)
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateNetwork(tx, n5))
return nil
}))
watchNetwork(t, netWatch, true, nil)
assert.NoError(t, s.Update(func(tx store.Tx) error {
t6 := &api.Task{
ID: "testTaskID6",
Status: api.TaskStatus{
State: api.TaskStateNew,
},
DesiredState: api.TaskStateRunning,
Networks: []*api.NetworkAttachment{
{
Network: n5,
},
},
}
assert.NoError(t, store.CreateTask(tx, t6))
return nil
}))
watchTask(t, s, taskWatch, true, nil)
// Now remove the conflicting network.
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteNetwork(tx, n4.ID))
return nil
}))
watchNetwork(t, netWatch, false, isValidNetwork)
watchTask(t, s, taskWatch, false, isValidTask)
// Try adding services with conflicting port configs and add
// task which is part of the service whose allocation hasn't
// happened and when that happens later and verify if task
// reconciles and moves to ALLOCATED.
s3 := &api.Service{
ID: "testServiceID3",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "service3",
},
Endpoint: &api.EndpointSpec{
Ports: []*api.PortConfig{
{
Name: "http",
TargetPort: 80,
PublishedPort: 8080,
},
},
},
},
}
s4 := s3.Copy()
s4.ID = "testServiceID4"
s4.Spec.Annotations.Name = "service4"
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateService(tx, s3))
return nil
}))
watchService(t, serviceWatch, false, nil)
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateService(tx, s4))
return nil
}))
watchService(t, serviceWatch, true, nil)
assert.NoError(t, s.Update(func(tx store.Tx) error {
t7 := &api.Task{
ID: "testTaskID7",
Status: api.TaskStatus{
State: api.TaskStateNew,
},
ServiceID: "testServiceID4",
DesiredState: api.TaskStateRunning,
}
assert.NoError(t, store.CreateTask(tx, t7))
return nil
}))
watchTask(t, s, taskWatch, true, nil)
// Now remove the conflicting service.
assert.NoError(t, s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteService(tx, s3.ID))
return nil
}))
watchService(t, serviceWatch, false, nil)
watchTask(t, s, taskWatch, false, isValidTask)
a.Stop()
}
func TestHA(t *testing.T) {
ctx := context.Background()
initialNodeSet := []*api.Node{
{
ID: "id1",
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id2",
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id3",
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id4",
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id5",
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
}
taskTemplate1 := &api.Task{
DesiredState: api.TaskStateRunning,
ServiceID: "service1",
Spec: api.TaskSpec{
Runtime: &api.TaskSpec_Container{
Container: &api.ContainerSpec{
Image: "v:1",
},
},
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
taskTemplate2 := &api.Task{
DesiredState: api.TaskStateRunning,
ServiceID: "service2",
Spec: api.TaskSpec{
Runtime: &api.TaskSpec_Container{
Container: &api.ContainerSpec{
Image: "v:2",
},
},
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
s := store.NewMemoryStore(nil)
assert.NotNil(t, s)
defer s.Close()
t1Instances := 18
err := s.Update(func(tx store.Tx) error {
// Prepoulate nodes
for _, n := range initialNodeSet {
assert.NoError(t, store.CreateNode(tx, n))
}
// Prepopulate tasks from template 1
for i := 0; i != t1Instances; i++ {
taskTemplate1.ID = fmt.Sprintf("t1id%d", i)
assert.NoError(t, store.CreateTask(tx, taskTemplate1))
}
return nil
})
assert.NoError(t, err)
scheduler := New(s)
watch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateTask{})
defer cancel()
go func() {
assert.NoError(t, scheduler.Run(ctx))
}()
defer scheduler.Stop()
t1Assignments := make(map[string]int)
for i := 0; i != t1Instances; i++ {
assignment := watchAssignment(t, watch)
if !strings.HasPrefix(assignment.ID, "t1") {
t.Fatal("got assignment for different kind of task")
}
t1Assignments[assignment.NodeID]++
}
assert.Len(t, t1Assignments, 5)
nodesWith3T1Tasks := 0
nodesWith4T1Tasks := 0
for nodeID, taskCount := range t1Assignments {
if taskCount == 3 {
nodesWith3T1Tasks++
} else if taskCount == 4 {
nodesWith4T1Tasks++
} else {
t.Fatalf("unexpected number of tasks %d on node %s", taskCount, nodeID)
}
}
assert.Equal(t, 3, nodesWith4T1Tasks)
assert.Equal(t, 2, nodesWith3T1Tasks)
t2Instances := 2
// Add a new service with two instances. They should fill the nodes
// that only have two tasks.
err = s.Update(func(tx store.Tx) error {
for i := 0; i != t2Instances; i++ {
taskTemplate2.ID = fmt.Sprintf("t2id%d", i)
assert.NoError(t, store.CreateTask(tx, taskTemplate2))
}
return nil
})
assert.NoError(t, err)
t2Assignments := make(map[string]int)
for i := 0; i != t2Instances; i++ {
assignment := watchAssignment(t, watch)
if !strings.HasPrefix(assignment.ID, "t2") {
t.Fatal("got assignment for different kind of task")
}
t2Assignments[assignment.NodeID]++
}
assert.Len(t, t2Assignments, 2)
for nodeID := range t2Assignments {
assert.Equal(t, 3, t1Assignments[nodeID])
}
// Scale up service 1 to 21 tasks. It should cover the two nodes that
// service 2 was assigned to, and also one other node.
err = s.Update(func(tx store.Tx) error {
for i := t1Instances; i != t1Instances+3; i++ {
taskTemplate1.ID = fmt.Sprintf("t1id%d", i)
assert.NoError(t, store.CreateTask(tx, taskTemplate1))
}
return nil
})
assert.NoError(t, err)
var sharedNodes [2]string
for i := 0; i != 3; i++ {
assignment := watchAssignment(t, watch)
if !strings.HasPrefix(assignment.ID, "t1") {
t.Fatal("got assignment for different kind of task")
}
if t1Assignments[assignment.NodeID] == 5 {
t.Fatal("more than one new task assigned to the same node")
}
t1Assignments[assignment.NodeID]++
if t2Assignments[assignment.NodeID] != 0 {
if sharedNodes[0] == "" {
sharedNodes[0] = assignment.NodeID
} else if sharedNodes[1] == "" {
sharedNodes[1] = assignment.NodeID
} else {
t.Fatal("all three assignments went to nodes with service2 tasks")
}
}
}
assert.NotEmpty(t, sharedNodes[0])
assert.NotEmpty(t, sharedNodes[1])
assert.NotEqual(t, sharedNodes[0], sharedNodes[1])
nodesWith4T1Tasks = 0
nodesWith5T1Tasks := 0
for nodeID, taskCount := range t1Assignments {
if taskCount == 4 {
nodesWith4T1Tasks++
} else if taskCount == 5 {
nodesWith5T1Tasks++
} else {
t.Fatalf("unexpected number of tasks %d on node %s", taskCount, nodeID)
}
}
assert.Equal(t, 4, nodesWith4T1Tasks)
assert.Equal(t, 1, nodesWith5T1Tasks)
// Add another task from service2. It must not land on the node that
// has 5 service1 tasks.
err = s.Update(func(tx store.Tx) error {
taskTemplate2.ID = "t2id4"
assert.NoError(t, store.CreateTask(tx, taskTemplate2))
return nil
})
assert.NoError(t, err)
assignment := watchAssignment(t, watch)
if assignment.ID != "t2id4" {
t.Fatal("got assignment for different task")
}
if t2Assignments[assignment.NodeID] != 0 {
t.Fatal("was scheduled on a node that already has a service2 task")
}
if t1Assignments[assignment.NodeID] == 5 {
t.Fatal("was scheduled on the node that has the most service1 tasks")
}
t2Assignments[assignment.NodeID]++
// Remove all tasks on node id1.
err = s.Update(func(tx store.Tx) error {
tasks, err := store.FindTasks(tx, store.ByNodeID("id1"))
assert.NoError(t, err)
for _, task := range tasks {
assert.NoError(t, store.DeleteTask(tx, task.ID))
}
return nil
})
assert.NoError(t, err)
t1Assignments["id1"] = 0
t2Assignments["id1"] = 0
// Add four instances of service1 and two instances of service2.
// All instances of service1 should land on node "id1", and one
// of the two service2 instances should as well.
// Put these in a map to randomize the order in which they are
// created.
err = s.Update(func(tx store.Tx) error {
tasksMap := make(map[string]*api.Task)
for i := 22; i <= 25; i++ {
taskTemplate1.ID = fmt.Sprintf("t1id%d", i)
tasksMap[taskTemplate1.ID] = taskTemplate1.Copy()
}
for i := 5; i <= 6; i++ {
taskTemplate2.ID = fmt.Sprintf("t2id%d", i)
tasksMap[taskTemplate2.ID] = taskTemplate2.Copy()
}
for _, task := range tasksMap {
assert.NoError(t, store.CreateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
for i := 0; i != 4+2; i++ {
assignment := watchAssignment(t, watch)
if strings.HasPrefix(assignment.ID, "t1") {
t1Assignments[assignment.NodeID]++
} else if strings.HasPrefix(assignment.ID, "t2") {
t2Assignments[assignment.NodeID]++
}
}
assert.Equal(t, 4, t1Assignments["id1"])
assert.Equal(t, 1, t2Assignments["id1"])
}
func TestScheduler(t *testing.T) {
ctx := context.Background()
initialNodeSet := []*api.Node{
{
ID: "id1",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name1",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id2",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name2",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
{
ID: "id3",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name2",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
},
}
initialTaskSet := []*api.Task{
{
ID: "id1",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name1",
},
Status: api.TaskStatus{
State: api.TaskStateAssigned,
},
NodeID: initialNodeSet[0].ID,
},
{
ID: "id2",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name2",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
},
{
ID: "id3",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name2",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
},
}
s := store.NewMemoryStore(nil)
assert.NotNil(t, s)
defer s.Close()
err := s.Update(func(tx store.Tx) error {
// 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)
scheduler := New(s)
watch, cancel := state.Watch(s.WatchQueue(), state.EventUpdateTask{})
defer cancel()
go func() {
assert.NoError(t, scheduler.Run(ctx))
}()
defer scheduler.Stop()
assignment1 := watchAssignment(t, watch)
// must assign to id2 or id3 since id1 already has a task
assert.Regexp(t, assignment1.NodeID, "(id2|id3)")
assignment2 := watchAssignment(t, watch)
// must assign to id2 or id3 since id1 already has a task
if assignment1.NodeID == "id2" {
assert.Equal(t, "id3", assignment2.NodeID)
} else {
assert.Equal(t, "id2", assignment2.NodeID)
}
err = s.Update(func(tx store.Tx) error {
// Update each node to make sure this doesn't mess up the
// scheduler's state.
for _, n := range initialNodeSet {
assert.NoError(t, store.UpdateNode(tx, n))
}
return nil
})
assert.NoError(t, err)
err = s.Update(func(tx store.Tx) error {
// Delete the task associated with node 1 so it's now the most lightly
// loaded node.
assert.NoError(t, store.DeleteTask(tx, "id1"))
// Create a new task. It should get assigned to id1.
t4 := &api.Task{
ID: "id4",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name4",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
assert.NoError(t, store.CreateTask(tx, t4))
return nil
})
assert.NoError(t, err)
assignment3 := watchAssignment(t, watch)
assert.Equal(t, "id1", assignment3.NodeID)
// Update a task to make it unassigned. It should get assigned by the
// scheduler.
err = s.Update(func(tx store.Tx) error {
// Remove assignment from task id4. It should get assigned
// to node id1.
t4 := &api.Task{
ID: "id4",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name4",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
assert.NoError(t, store.UpdateTask(tx, t4))
return nil
})
assert.NoError(t, err)
assignment4 := watchAssignment(t, watch)
assert.Equal(t, "id1", assignment4.NodeID)
err = s.Update(func(tx store.Tx) error {
// Create a ready node, then remove it. No tasks should ever
// be assigned to it.
node := &api.Node{
ID: "removednode",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "removednode",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_DOWN,
},
}
assert.NoError(t, store.CreateNode(tx, node))
assert.NoError(t, store.DeleteNode(tx, node.ID))
// Create an unassigned task.
task := &api.Task{
ID: "removednode",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "removednode",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
assert.NoError(t, store.CreateTask(tx, task))
return nil
})
assert.NoError(t, err)
assignmentRemovedNode := watchAssignment(t, watch)
assert.NotEqual(t, "removednode", assignmentRemovedNode.NodeID)
err = s.Update(func(tx store.Tx) error {
// Create a ready node. It should be used for the next
// assignment.
n4 := &api.Node{
ID: "id4",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name4",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
}
assert.NoError(t, store.CreateNode(tx, n4))
// Create an unassigned task.
t5 := &api.Task{
ID: "id5",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name5",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
assert.NoError(t, store.CreateTask(tx, t5))
return nil
})
assert.NoError(t, err)
assignment5 := watchAssignment(t, watch)
assert.Equal(t, "id4", assignment5.NodeID)
err = s.Update(func(tx store.Tx) error {
// Create a non-ready node. It should NOT be used for the next
// assignment.
n5 := &api.Node{
ID: "id5",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name5",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_DOWN,
},
}
assert.NoError(t, store.CreateNode(tx, n5))
// Create an unassigned task.
t6 := &api.Task{
ID: "id6",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name6",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
assert.NoError(t, store.CreateTask(tx, t6))
return nil
})
assert.NoError(t, err)
assignment6 := watchAssignment(t, watch)
assert.NotEqual(t, "id5", assignment6.NodeID)
err = s.Update(func(tx store.Tx) error {
// Update node id5 to put it in the READY state.
n5 := &api.Node{
ID: "id5",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name5",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
}
assert.NoError(t, store.UpdateNode(tx, n5))
// Create an unassigned task. Should be assigned to the
// now-ready node.
t7 := &api.Task{
ID: "id7",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name7",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
assert.NoError(t, store.CreateTask(tx, t7))
return nil
})
assert.NoError(t, err)
assignment7 := watchAssignment(t, watch)
assert.Equal(t, "id5", assignment7.NodeID)
err = s.Update(func(tx store.Tx) error {
// Create a ready node, then immediately take it down. The next
// unassigned task should NOT be assigned to it.
n6 := &api.Node{
ID: "id6",
Spec: api.NodeSpec{
Annotations: api.Annotations{
Name: "name6",
},
},
Status: api.NodeStatus{
State: api.NodeStatus_READY,
},
}
assert.NoError(t, store.CreateNode(tx, n6))
n6.Status.State = api.NodeStatus_DOWN
assert.NoError(t, store.UpdateNode(tx, n6))
// Create an unassigned task.
t8 := &api.Task{
ID: "id8",
DesiredState: api.TaskStateRunning,
ServiceAnnotations: api.Annotations{
Name: "name8",
},
Status: api.TaskStatus{
State: api.TaskStatePending,
},
}
assert.NoError(t, store.CreateTask(tx, t8))
return nil
})
assert.NoError(t, err)
assignment8 := watchAssignment(t, watch)
assert.NotEqual(t, "id6", assignment8.NodeID)
}
func deleteTask(t *testing.T, s *store.MemoryStore, task *api.Task) {
s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteTask(tx, task.ID))
return nil
})
}
func (tr *TaskReaper) tick() {
if len(tr.dirty) == 0 {
return
}
defer func() {
tr.dirty = make(map[instanceTuple]struct{})
}()
var deleteTasks []string
tr.store.View(func(tx store.ReadTx) {
for dirty := range tr.dirty {
service := store.GetService(tx, dirty.serviceID)
if service == nil {
continue
}
taskHistory := tr.taskHistory
if taskHistory < 0 {
continue
}
var historicTasks []*api.Task
switch service.Spec.GetMode().(type) {
case *api.ServiceSpec_Replicated:
var err error
historicTasks, err = store.FindTasks(tx, store.BySlot(dirty.serviceID, dirty.instance))
if err != nil {
continue
}
case *api.ServiceSpec_Global:
tasksByNode, err := store.FindTasks(tx, store.ByNodeID(dirty.nodeID))
if err != nil {
continue
}
for _, t := range tasksByNode {
if t.ServiceID == dirty.serviceID {
historicTasks = append(historicTasks, t)
}
}
}
if int64(len(historicTasks)) <= taskHistory {
continue
}
// TODO(aaronl): This could filter for non-running tasks and use quickselect
// instead of sorting the whole slice.
sort.Sort(tasksByTimestamp(historicTasks))
for _, t := range historicTasks {
if t.DesiredState <= api.TaskStateRunning {
// Don't delete running tasks
continue
}
deleteTasks = append(deleteTasks, t.ID)
taskHistory++
if int64(len(historicTasks)) <= taskHistory {
break
}
}
}
})
if len(deleteTasks) > 0 {
tr.store.Batch(func(batch *store.Batch) error {
for _, taskID := range deleteTasks {
batch.Update(func(tx store.Tx) error {
return store.DeleteTask(tx, taskID)
})
}
return nil
})
}
}
func TestReplicatedOrchestrator(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.EventCreateTask{}, state.EventUpdateTask{}*/)
defer cancel()
// Create a service with two instances specified before the orchestrator is
// started. This should result in two tasks when the orchestrator
// starts up.
err := s.Update(func(tx store.Tx) error {
s1 := &api.Service{
ID: "id1",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "name1",
},
Task: api.TaskSpec{
Runtime: &api.TaskSpec_Container{
Container: &api.ContainerSpec{},
},
},
Mode: &api.ServiceSpec_Replicated{
Replicated: &api.ReplicatedService{
Replicas: 2,
},
},
},
}
assert.NoError(t, store.CreateService(tx, s1))
return nil
})
assert.NoError(t, err)
// Start the orchestrator.
go func() {
assert.NoError(t, orchestrator.Run(ctx))
}()
observedTask1 := watchTaskCreate(t, watch)
assert.Equal(t, observedTask1.Status.State, api.TaskStateNew)
assert.Equal(t, observedTask1.ServiceAnnotations.Name, "name1")
observedTask2 := watchTaskCreate(t, watch)
assert.Equal(t, observedTask2.Status.State, api.TaskStateNew)
assert.Equal(t, observedTask2.ServiceAnnotations.Name, "name1")
// Create a second service.
err = s.Update(func(tx store.Tx) error {
s2 := &api.Service{
ID: "id2",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "name2",
},
Task: api.TaskSpec{
Runtime: &api.TaskSpec_Container{
Container: &api.ContainerSpec{},
},
},
Mode: &api.ServiceSpec_Replicated{
Replicated: &api.ReplicatedService{
Replicas: 1,
},
},
},
}
assert.NoError(t, store.CreateService(tx, s2))
return nil
})
assert.NoError(t, err)
observedTask3 := watchTaskCreate(t, watch)
assert.Equal(t, observedTask3.Status.State, api.TaskStateNew)
assert.Equal(t, observedTask3.ServiceAnnotations.Name, "name2")
// Update a service to scale it out to 3 instances
err = s.Update(func(tx store.Tx) error {
s2 := &api.Service{
ID: "id2",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "name2",
},
Task: api.TaskSpec{
Runtime: &api.TaskSpec_Container{
Container: &api.ContainerSpec{},
},
},
Mode: &api.ServiceSpec_Replicated{
Replicated: &api.ReplicatedService{
Replicas: 3,
},
},
},
}
assert.NoError(t, store.UpdateService(tx, s2))
return nil
})
assert.NoError(t, err)
observedTask4 := watchTaskCreate(t, watch)
assert.Equal(t, observedTask4.Status.State, api.TaskStateNew)
assert.Equal(t, observedTask4.ServiceAnnotations.Name, "name2")
observedTask5 := watchTaskCreate(t, watch)
assert.Equal(t, observedTask5.Status.State, api.TaskStateNew)
assert.Equal(t, observedTask5.ServiceAnnotations.Name, "name2")
// Now scale it back down to 1 instance
err = s.Update(func(tx store.Tx) error {
s2 := &api.Service{
ID: "id2",
Spec: api.ServiceSpec{
Annotations: api.Annotations{
Name: "name2",
},
Task: api.TaskSpec{
Runtime: &api.TaskSpec_Container{
Container: &api.ContainerSpec{},
},
},
Mode: &api.ServiceSpec_Replicated{
Replicated: &api.ReplicatedService{
Replicas: 1,
},
},
},
}
assert.NoError(t, store.UpdateService(tx, s2))
return nil
})
assert.NoError(t, err)
observedDeletion1 := watchShutdownTask(t, watch)
assert.Equal(t, observedDeletion1.Status.State, api.TaskStateNew)
assert.Equal(t, observedDeletion1.ServiceAnnotations.Name, "name2")
observedDeletion2 := watchShutdownTask(t, watch)
assert.Equal(t, observedDeletion2.Status.State, api.TaskStateNew)
assert.Equal(t, observedDeletion2.ServiceAnnotations.Name, "name2")
// There should be one remaining task attached to service id2/name2.
var liveTasks []*api.Task
s.View(func(readTx store.ReadTx) {
var tasks []*api.Task
tasks, err = store.FindTasks(readTx, store.ByServiceID("id2"))
for _, t := range tasks {
if t.DesiredState == api.TaskStateRunning {
liveTasks = append(liveTasks, t)
}
}
})
assert.NoError(t, err)
assert.Len(t, liveTasks, 1)
// Delete the remaining task directly. It should be recreated by the
// orchestrator.
err = s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteTask(tx, liveTasks[0].ID))
return nil
})
assert.NoError(t, err)
observedTask6 := watchTaskCreate(t, watch)
assert.Equal(t, observedTask6.Status.State, api.TaskStateNew)
assert.Equal(t, observedTask6.ServiceAnnotations.Name, "name2")
// Delete the service. Its remaining task should go away.
err = s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteService(tx, "id2"))
return nil
})
assert.NoError(t, err)
deletedTask := watchTaskDelete(t, watch)
assert.Equal(t, deletedTask.Status.State, api.TaskStateNew)
assert.Equal(t, deletedTask.ServiceAnnotations.Name, "name2")
}
func TestTasks(t *testing.T) {
gd, err := startDispatcher(DefaultConfig())
assert.NoError(t, err)
defer gd.Close()
var expectedSessionID string
var nodeID string
{
stream, err := gd.Clients[0].Session(context.Background(), &api.SessionRequest{})
assert.NoError(t, err)
defer stream.CloseSend()
resp, err := stream.Recv()
assert.NoError(t, err)
assert.NotEmpty(t, resp.SessionID)
expectedSessionID = resp.SessionID
nodeID = resp.Node.ID
}
testTask1 := &api.Task{
NodeID: nodeID,
ID: "testTask1",
Status: api.TaskStatus{State: api.TaskStateAssigned},
}
testTask2 := &api.Task{
NodeID: nodeID,
ID: "testTask2",
Status: api.TaskStatus{State: api.TaskStateAssigned},
}
{
// without correct SessionID should fail
stream, err := gd.Clients[0].Tasks(context.Background(), &api.TasksRequest{})
assert.NoError(t, err)
assert.NotNil(t, stream)
resp, err := stream.Recv()
assert.Nil(t, resp)
assert.Error(t, err)
assert.Equal(t, grpc.Code(err), codes.InvalidArgument)
}
stream, err := gd.Clients[0].Tasks(context.Background(), &api.TasksRequest{SessionID: expectedSessionID})
assert.NoError(t, err)
time.Sleep(100 * time.Millisecond)
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateTask(tx, testTask1))
assert.NoError(t, store.CreateTask(tx, testTask2))
return nil
})
assert.NoError(t, err)
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.UpdateTask(tx, &api.Task{
ID: testTask1.ID,
NodeID: nodeID,
Status: api.TaskStatus{State: api.TaskStateFailed, Err: "1234"},
}))
return nil
})
assert.NoError(t, err)
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteTask(tx, testTask1.ID))
assert.NoError(t, store.DeleteTask(tx, testTask2.ID))
return nil
})
assert.NoError(t, err)
resp, err := stream.Recv()
assert.NoError(t, err)
assert.Equal(t, 0, len(resp.Tasks))
resp, err = stream.Recv()
assert.Equal(t, len(resp.Tasks), 1)
assert.Equal(t, resp.Tasks[0].ID, "testTask1")
resp, err = stream.Recv()
assert.NoError(t, err)
assert.Equal(t, len(resp.Tasks), 2)
assert.True(t, resp.Tasks[0].ID == "testTask1" && resp.Tasks[1].ID == "testTask2" || resp.Tasks[0].ID == "testTask2" && resp.Tasks[1].ID == "testTask1")
resp, err = stream.Recv()
assert.NoError(t, err)
assert.Equal(t, len(resp.Tasks), 2)
for _, task := range resp.Tasks {
if task.ID == "testTask1" {
assert.Equal(t, task.Status.State, api.TaskStateFailed)
assert.Equal(t, task.Status.Err, "1234")
}
}
resp, err = stream.Recv()
assert.NoError(t, err)
assert.Equal(t, len(resp.Tasks), 1)
resp, err = stream.Recv()
assert.NoError(t, err)
assert.Equal(t, len(resp.Tasks), 0)
}
func TestTasksBatch(t *testing.T) {
gd, err := startDispatcher(DefaultConfig())
assert.NoError(t, err)
defer gd.Close()
var expectedSessionID string
var nodeID string
{
stream, err := gd.Clients[0].Session(context.Background(), &api.SessionRequest{})
assert.NoError(t, err)
defer stream.CloseSend()
resp, err := stream.Recv()
assert.NoError(t, err)
assert.NotEmpty(t, resp.SessionID)
expectedSessionID = resp.SessionID
nodeID = resp.Node.ID
}
testTask1 := &api.Task{
NodeID: nodeID,
ID: "testTask1",
Status: api.TaskStatus{State: api.TaskStateAssigned},
}
testTask2 := &api.Task{
NodeID: nodeID,
ID: "testTask2",
Status: api.TaskStatus{State: api.TaskStateAssigned},
}
stream, err := gd.Clients[0].Assignments(context.Background(), &api.AssignmentsRequest{SessionID: expectedSessionID})
assert.NoError(t, err)
resp, err := stream.Recv()
assert.NoError(t, err)
// initially no tasks
assert.Equal(t, 0, len(resp.Changes))
// Create, Update and Delete tasks.
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.CreateTask(tx, testTask1))
assert.NoError(t, store.CreateTask(tx, testTask2))
return nil
})
assert.NoError(t, err)
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.UpdateTask(tx, testTask1))
assert.NoError(t, store.UpdateTask(tx, testTask2))
return nil
})
assert.NoError(t, err)
err = gd.Store.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteTask(tx, testTask1.ID))
assert.NoError(t, store.DeleteTask(tx, testTask2.ID))
return nil
})
assert.NoError(t, err)
resp, err = stream.Recv()
assert.NoError(t, err)
// all tasks have been deleted
tasks, secrets := collectTasksAndSecrets(resp.Changes)
assert.Len(t, tasks, 2)
assert.Len(t, secrets, 0)
assert.Equal(t, api.AssignmentChange_AssignmentActionRemove, resp.Changes[0].Action)
assert.Equal(t, api.AssignmentChange_AssignmentActionRemove, resp.Changes[1].Action)
}
// As tasks are added, assignments will send down tasks > ASSIGNED, and any secrets
// for said tasks that are <= RUNNING (if the secrets exist)
func TestAssignmentsAddingTasks(t *testing.T) {
t.Parallel()
gd, err := startDispatcher(DefaultConfig())
assert.NoError(t, err)
defer gd.Close()
expectedSessionID, nodeID := getSessionAndNodeID(t, gd.Clients[0])
stream, err := gd.Clients[0].Assignments(context.Background(), &api.AssignmentsRequest{SessionID: expectedSessionID})
assert.NoError(t, err)
defer stream.CloseSend()
time.Sleep(100 * time.Millisecond)
// There are no initial tasks or secrets
resp, err := stream.Recv()
assert.NoError(t, err)
assert.Empty(t, resp.Changes)
// create the relevant secrets and tasks and update the tasks
secrets, tasks := makeTasksAndSecrets(t, nodeID)
err = gd.Store.Update(func(tx store.Tx) error {
for _, secret := range secrets[:len(secrets)-1] {
assert.NoError(t, store.CreateSecret(tx, secret))
}
for _, task := range tasks {
assert.NoError(t, store.CreateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
// Nothing happens until we update. Updating all the tasks will send updates for all the tasks >= ASSIGNED (10),
// and secrets for all the tasks >= ASSIGNED and <= RUNNING (6).
err = gd.Store.Update(func(tx store.Tx) error {
for _, task := range tasks {
assert.NoError(t, store.UpdateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
resp, err = stream.Recv()
assert.NoError(t, err)
// FIXME(aaronl): This is hard to maintain.
assert.Equal(t, 10+6, len(resp.Changes))
taskChanges, secretChanges := collectTasksAndSecrets(resp.Changes)
assert.Len(t, taskChanges, 10)
for _, task := range tasks[2:] {
assert.NotNil(t, taskChanges[idAndAction{id: task.ID, action: api.AssignmentChange_AssignmentActionUpdate}])
}
assert.Len(t, secretChanges, 6)
// all the secrets for tasks >= ASSIGNED and <= RUNNING
for _, secret := range secrets[2:8] {
assert.NotNil(t, secretChanges[idAndAction{id: secret.ID, action: api.AssignmentChange_AssignmentActionUpdate}])
}
// deleting the tasks removes all the secrets for every single task, no matter
// what state it's in
err = gd.Store.Update(func(tx store.Tx) error {
for _, task := range tasks {
assert.NoError(t, store.DeleteTask(tx, task.ID))
}
return nil
})
assert.NoError(t, err)
// updates for all the tasks >= ASSIGNMENT, and remove secrets for all of them, even ones that don't exist
// (there will be 2 tasks changes that won't be sent down)
resp, err = stream.Recv()
assert.NoError(t, err)
assert.Equal(t, len(tasks)-2+len(secrets)-2, len(resp.Changes))
taskChanges, secretChanges = collectTasksAndSecrets(resp.Changes)
assert.Len(t, taskChanges, len(tasks)-2)
for _, task := range tasks[2:] {
assert.NotNil(t, taskChanges[idAndAction{id: task.ID, action: api.AssignmentChange_AssignmentActionRemove}])
}
assert.Len(t, secretChanges, len(secrets)-2)
for _, secret := range secrets[2:] {
assert.NotNil(t, secretChanges[idAndAction{id: secret.ID, action: api.AssignmentChange_AssignmentActionRemove}])
}
}
// Assignments will send down any existing node tasks > ASSIGNED, and any secrets
// for said tasks that are <= RUNNING (if the secrets exist)
func TestAssignmentsInitialNodeTasks(t *testing.T) {
t.Parallel()
gd, err := startDispatcher(DefaultConfig())
assert.NoError(t, err)
defer gd.Close()
expectedSessionID, nodeID := getSessionAndNodeID(t, gd.Clients[0])
// create the relevant secrets and tasks
secrets, tasks := makeTasksAndSecrets(t, nodeID)
err = gd.Store.Update(func(tx store.Tx) error {
for _, secret := range secrets[:] {
assert.NoError(t, store.CreateSecret(tx, secret))
}
for _, task := range tasks {
assert.NoError(t, store.CreateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
stream, err := gd.Clients[0].Assignments(context.Background(), &api.AssignmentsRequest{SessionID: expectedSessionID})
assert.NoError(t, err)
defer stream.CloseSend()
time.Sleep(100 * time.Millisecond)
// check the initial task and secret stream
resp, err := stream.Recv()
assert.NoError(t, err)
// FIXME(aaronl): This is hard to maintain.
assert.Equal(t, 17, len(resp.Changes))
taskChanges, secretChanges := collectTasksAndSecrets(resp.Changes)
assert.Len(t, taskChanges, 10) // 10 types of task states >= assigned, 2 types < assigned
for _, task := range tasks[2:] {
assert.NotNil(t, taskChanges[idAndAction{id: task.ID, action: api.AssignmentChange_AssignmentActionUpdate}])
}
assert.Len(t, secretChanges, 7) // 6 different secrets for states between assigned and running inclusive plus secret12
for _, secret := range secrets[2:8] {
assert.NotNil(t, secretChanges[idAndAction{id: secret.ID, action: api.AssignmentChange_AssignmentActionUpdate}])
}
// updating all the tasks will attempt to remove all the secrets for the tasks that are in state > running
err = gd.Store.Update(func(tx store.Tx) error {
for _, task := range tasks {
assert.NoError(t, store.UpdateTask(tx, task))
}
return nil
})
assert.NoError(t, err)
// updates for all the tasks, remove secret sent for the 4 types of states > running
resp, err = stream.Recv()
assert.NoError(t, err)
assert.Equal(t, 5, len(resp.Changes))
taskChanges, secretChanges = collectTasksAndSecrets(resp.Changes)
assert.Len(t, taskChanges, 1)
assert.NotNil(t, taskChanges[idAndAction{id: tasks[2].ID, action: api.AssignmentChange_AssignmentActionUpdate}]) // this is the task in ASSIGNED
assert.Len(t, secretChanges, 4) // these are the secrets for states > running
for _, secret := range secrets[9 : len(secrets)-1] {
assert.NotNil(t, secretChanges[idAndAction{id: secret.ID, action: api.AssignmentChange_AssignmentActionRemove}])
}
// deleting the tasks removes all the secrets for every single task, no matter
// what state it's in
err = gd.Store.Update(func(tx store.Tx) error {
for _, task := range tasks {
assert.NoError(t, store.DeleteTask(tx, task.ID))
}
return nil
})
assert.NoError(t, err)
// updates for all the tasks >= ASSIGNMENT, and remove secrets for all of them,
// (there will be 2 tasks changes that won't be sent down)
resp, err = stream.Recv()
assert.NoError(t, err)
assert.Equal(t, len(tasks)-2+len(secrets)-2, len(resp.Changes))
taskChanges, secretChanges = collectTasksAndSecrets(resp.Changes)
assert.Len(t, taskChanges, len(tasks)-2)
for _, task := range tasks[2:] {
assert.NotNil(t, taskChanges[idAndAction{id: task.ID, action: api.AssignmentChange_AssignmentActionRemove}])
}
assert.Len(t, secretChanges, len(secrets)-2)
for _, secret := range secrets[2:] {
assert.NotNil(t, secretChanges[idAndAction{id: secret.ID, action: api.AssignmentChange_AssignmentActionRemove}])
}
}
