// CalculateGroups is a function designed to ingest a list of group objects, a collection of agents, and
// return a map that contains the resulting group for each agent UUID.
func CalculateGroups(cfg config.Config) {
tdb, err := db.NewToddDB(cfg)
if err != nil {
log.Fatalf("Error connecting to DB: %v", err)
}
// Retrieve all currently active agents
agents, err := tdb.GetAgents()
if err != nil {
log.Fatalf("Error retrieving agents: %v", err)
}
// Retrieve all objects with type "group"
group_objs, err := tdb.GetObjects("group")
if err != nil {
log.Fatalf("Error retrieving groups: %v", err)
}
// Cast retrieved slice of ToddObject interfaces to actual GroupObjects
groups := make([]objects.GroupObject, len(group_objs))
for i, gobj := range group_objs {
groups[i] = gobj.(objects.GroupObject)
}
// groupmap contains the uuid-to-groupname mappings to be used for test runs
groupmap := map[string]string{}
// This slice will hold all of the agents that are sad because they didn't get into a group
var lonelyAgents []defs.AgentAdvert
next:
for x := range agents {
for i := range groups {
// See if this agent is in this group
if isInGroup(groups[i].Spec.Matches, agents[x].Facts) {
// Insert this group name ("Label") into groupmap under the key of the UUID for the agent that belongs to it
log.Debugf("Agent %s is in group %s\n", agents[x].Uuid, groups[i].Label)
groupmap[agents[x].Uuid] = groups[i].Label
continue next
}
}
// The "continue next" should prohibit all agents that have a group from getting to this point,
// so the only ones left do not have a group.
lonelyAgents = append(lonelyAgents, agents[x])
}
// Write results to database
err = tdb.SetGroupMap(groupmap)
if err != nil {
log.Fatalf("Error setting group map: %v", err)
}
// Send notifications to each agent to let them know what group they're in, so they can cache it
var tc = comms.NewToDDComms(cfg)
for uuid, groupName := range groupmap {
setGroupTask := tasks.SetGroupTask{
GroupName: groupName,
}
setGroupTask.Type = "SetGroup" //TODO(mierdin): Apparently this is necessary because inner type promotion doesn't apply for struct literals?
tc.CommsPackage.SendTask(uuid, setGroupTask)
}
// need to send a message to all agents that weren't in groupmap to set their group to nothing
for x := range lonelyAgents {
setGroupTask := tasks.SetGroupTask{
GroupName: "",
}
setGroupTask.Type = "SetGroup" //TODO(mierdin): Apparently this is necessary because inner type promotion doesn't apply for struct literals?
tc.CommsPackage.SendTask(lonelyAgents[x].Uuid, setGroupTask)
}
}
// ListenForTasks is a method that recieves task notices from the server
func (rmq rabbitMQComms) ListenForTasks(uuid string) error {
// Connect to RabbitMQ with retry logic
conn, err := connectRabbitMQ(rmq.queueUrl)
if err != nil {
log.Error("(AdvertiseAgent) Failed to connect to RabbitMQ")
return err
}
defer conn.Close()
ch, err := conn.Channel()
if err != nil {
log.Error("Failed to open a channel")
os.Exit(1)
}
defer ch.Close()
_, err = ch.QueueDeclare(
uuid, // name
false, // durable
false, // delete when usused
false, // exclusive
false, // no-wait
nil, // arguments
)
if err != nil {
log.Error("Failed to declare a queue")
os.Exit(1)
}
msgs, err := ch.Consume(
uuid, // queue
uuid, // consumer
true, // auto-ack
false, // exclusive
false, // no-local
false, // no-wait
nil, // args
)
if err != nil {
log.Error("Failed to register a consumer")
os.Exit(1)
}
forever := make(chan bool)
go func() {
for d := range msgs {
// Unmarshal into BaseTaskMessage to determine type
var base_msg tasks.BaseTask
err = json.Unmarshal(d.Body, &base_msg)
// TODO(mierdin): Need to handle this error
log.Debugf("Agent task received: %s", d.Body)
// call agent task method based on type
switch base_msg.Type {
case "DownloadAsset":
downloadAssetTask := tasks.DownloadAssetTask{
HTTPClient: &http.Client{},
Fs: tasks.OsFS{},
Ios: tasks.IoSys{},
CollectorDir: fmt.Sprintf("%s/assets/factcollectors", rmq.config.LocalResources.OptDir),
TestletDir: fmt.Sprintf("%s/assets/testlets", rmq.config.LocalResources.OptDir),
}
err = json.Unmarshal(d.Body, &downloadAssetTask)
// TODO(mierdin): Need to handle this error
err = downloadAssetTask.Run()
if err != nil {
log.Warning("The KeyValue task failed to initialize")
}
case "KeyValue":
kv_task := tasks.KeyValueTask{
Config: rmq.config,
}
err = json.Unmarshal(d.Body, &kv_task)
// TODO(mierdin): Need to handle this error
err = kv_task.Run()
if err != nil {
log.Warning("The KeyValue task failed to initialize")
}
case "SetGroup":
sg_task := tasks.SetGroupTask{
Config: rmq.config,
}
err = json.Unmarshal(d.Body, &sg_task)
// TODO(mierdin): Need to handle this error
err = sg_task.Run()
if err != nil {
log.Warning("The SetGroup task failed to initialize")
}
case "DeleteTestData":
dtdt_task := tasks.DeleteTestDataTask{
Config: rmq.config,
}
err = json.Unmarshal(d.Body, &dtdt_task)
// TODO(mierdin): Need to handle this error
err = dtdt_task.Run()
if err != nil {
log.Warning("The DeleteTestData task failed to initialize")
}
case "InstallTestRun":
// Retrieve UUID
var ac = cache.NewAgentCache(rmq.config)
uuid := ac.GetKeyValue("uuid")
itr_task := tasks.InstallTestRunTask{
Config: rmq.config,
}
err = json.Unmarshal(d.Body, &itr_task)
// TODO(mierdin): Need to handle this error
var response responses.SetAgentStatusResponse
response.Type = "AgentStatus" //TODO(mierdin): This is an extra step. Maybe a factory function for the task could help here?
response.AgentUuid = uuid
response.TestUuid = itr_task.Tr.Uuid
err = itr_task.Run()
if err != nil {
log.Warning("The InstallTestRun task failed to initialize")
response.Status = "fail"
} else {
response.Status = "ready"
}
rmq.SendResponse(response)
case "ExecuteTestRun":
// Retrieve UUID
var ac = cache.NewAgentCache(rmq.config)
uuid := ac.GetKeyValue("uuid")
etr_task := tasks.ExecuteTestRunTask{
Config: rmq.config,
}
err = json.Unmarshal(d.Body, &etr_task)
// TODO(mierdin): Need to handle this error
// Send status that the testing has begun, right now.
response := responses.SetAgentStatusResponse{
TestUuid: etr_task.TestUuid,
Status: "testing",
}
response.AgentUuid = uuid // TODO(mierdin): Can't declare this in the literal, it's that embedding behavior again. Need to figure this out.
response.Type = "AgentStatus" //TODO(mierdin): This is an extra step. Maybe a factory function for the task could help here?
rmq.SendResponse(response)
err = etr_task.Run()
if err != nil {
log.Warning("The ExecuteTestRun task failed to initialize")
response.Status = "fail"
rmq.SendResponse(response)
}
default:
log.Errorf(fmt.Sprintf("Unexpected type value for received task: %s", base_msg.Type))
}
}
}()
log.Infof(" [*] Waiting for messages. To exit press CTRL+C")
<-forever
return nil
}