// 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
}
}
}
}
