// Assignments is a stream of assignments for a node. Each message contains // either full list of tasks and secrets for the node, or an incremental update. func (d *Dispatcher) Assignments(r *api.AssignmentsRequest, stream api.Dispatcher_AssignmentsServer) error { nodeInfo, err := ca.RemoteNode(stream.Context()) if err != nil { return err } nodeID := nodeInfo.NodeID if err := d.isRunningLocked(); err != nil { return err } fields := logrus.Fields{ "node.id": nodeID, "node.session": r.SessionID, "method": "(*Dispatcher).Assignments", } if nodeInfo.ForwardedBy != nil { fields["forwarder.id"] = nodeInfo.ForwardedBy.NodeID } log := log.G(stream.Context()).WithFields(fields) log.Debugf("") if _, err = d.nodes.GetWithSession(nodeID, r.SessionID); err != nil { return err } var ( sequence int64 appliesTo string initial api.AssignmentsMessage ) tasksMap := make(map[string]*api.Task) sendMessage := func(msg api.AssignmentsMessage, assignmentType api.AssignmentsMessage_Type) error { sequence++ msg.AppliesTo = appliesTo msg.ResultsIn = strconv.FormatInt(sequence, 10) appliesTo = msg.ResultsIn msg.Type = assignmentType if err := stream.Send(&msg); err != nil { return err } return nil } // TODO(aaronl): Also send node secrets that should be exposed to // this node. nodeTasks, cancel, err := store.ViewAndWatch( d.store, func(readTx store.ReadTx) error { tasks, err := store.FindTasks(readTx, store.ByNodeID(nodeID)) if err != nil { return err } for _, t := range tasks { // We only care about tasks that are ASSIGNED or // higher. If the state is below ASSIGNED, the // task may not meet the constraints for this // node, so we have to be careful about sending // secrets associated with it. if t.Status.State < api.TaskStateAssigned { continue } tasksMap[t.ID] = t initial.UpdateTasks = append(initial.UpdateTasks, t) } return nil }, state.EventUpdateTask{Task: &api.Task{NodeID: nodeID}, Checks: []state.TaskCheckFunc{state.TaskCheckNodeID}}, state.EventDeleteTask{Task: &api.Task{NodeID: nodeID}, Checks: []state.TaskCheckFunc{state.TaskCheckNodeID}}, ) if err != nil { return err } defer cancel() if err := sendMessage(initial, api.AssignmentsMessage_COMPLETE); err != nil { return err } for { // Check for session expiration if _, err := d.nodes.GetWithSession(nodeID, r.SessionID); err != nil { return err } // bursty events should be processed in batches and sent out together var ( update api.AssignmentsMessage modificationCnt int batchingTimer *time.Timer batchingTimeout <-chan time.Time updateTasks = make(map[string]*api.Task) removeTasks = make(map[string]struct{}) ) oneModification := func() { modificationCnt++ if batchingTimer != nil { batchingTimer.Reset(batchingWaitTime) } else { batchingTimer = time.NewTimer(batchingWaitTime) batchingTimeout = batchingTimer.C } } // The batching loop waits for 50 ms after the most recent // change, or until modificationBatchLimit is reached. The // worst case latency is modificationBatchLimit * batchingWaitTime, // which is 10 seconds. batchingLoop: for modificationCnt < modificationBatchLimit { select { case event := <-nodeTasks: switch v := event.(type) { // We don't monitor EventCreateTask because tasks are // never created in the ASSIGNED state. First tasks are // created by the orchestrator, then the scheduler moves // them to ASSIGNED. If this ever changes, we will need // to monitor task creations as well. case state.EventUpdateTask: // We only care about tasks that are ASSIGNED or // higher. if v.Task.Status.State < api.TaskStateAssigned { continue } if oldTask, exists := tasksMap[v.Task.ID]; exists { // States ASSIGNED and below are set by the orchestrator/scheduler, // not the agent, so tasks in these states need to be sent to the // agent even if nothing else has changed. if equality.TasksEqualStable(oldTask, v.Task) && v.Task.Status.State > api.TaskStateAssigned { // this update should not trigger a task change for the agent tasksMap[v.Task.ID] = v.Task continue } } tasksMap[v.Task.ID] = v.Task updateTasks[v.Task.ID] = v.Task oneModification() case state.EventDeleteTask: if _, exists := tasksMap[v.Task.ID]; !exists { continue } removeTasks[v.Task.ID] = struct{}{} delete(tasksMap, v.Task.ID) oneModification() } case <-batchingTimeout: break batchingLoop case <-stream.Context().Done(): return stream.Context().Err() case <-d.ctx.Done(): return d.ctx.Err() } } if batchingTimer != nil { batchingTimer.Stop() } if modificationCnt > 0 { for id, task := range updateTasks { if _, ok := removeTasks[id]; !ok { update.UpdateTasks = append(update.UpdateTasks, task) } } for id := range removeTasks { update.RemoveTasks = append(update.RemoveTasks, id) } if err := sendMessage(update, api.AssignmentsMessage_INCREMENTAL); err != nil { return err } } } }