// WatchForFinishedTestRuns simply watches the local cache for any test runs that have test data. // It will periodically look at the table and send any present test data back to the server as a response. // When the server has successfully received this data, it will send a task back to this specific agent // to delete this row from the cache. func WatchForFinishedTestRuns(cfg config.Config) { var ac = cache.NewAgentCache(cfg) agentUuid := ac.GetKeyValue("uuid") for { time.Sleep(5000 * time.Millisecond) testruns, err := ac.GetFinishedTestRuns() if err != nil { log.Error("Problem retrieving finished test runs") os.Exit(1) } for testUuid, testData := range testruns { log.Debug("Found ripe testrun: ", testUuid) var utdr = responses.UploadTestDataResponse{ TestUuid: testUuid, TestData: testData, } utdr.AgentUuid = agentUuid utdr.Type = "TestData" //TODO(mierdin): This is an extra step. Maybe a factory function for the task could help here? var tc = comms.NewToDDComms(cfg) tc.CommsPackage.SendResponse(utdr) } } }
// Run contains the logic necessary to perform this task on the agent. This particular task // will simply pass a key/value pair to the agent cache to be set func (kvt KeyValueTask) Run() error { var ac = cache.NewAgentCache(kvt.Config) err := ac.SetKeyValue(kvt.Key, kvt.Value) if err != nil { return errors.New(fmt.Sprintf("KeyValueTask failed - %s:%s", kvt.Key, kvt.Value)) } log.Info(fmt.Sprintf("KeyValueTask successful - %s:%s", kvt.Key, kvt.Value)) return nil }
// WatchForGroup should be run as a goroutine, like other background services. This is because it will itself spawn a goroutine to // listen for tasks that are sent to groups, and this goroutine can be restarted when group membership changes func (rmq rabbitMQComms) WatchForGroup() { var ac = cache.NewAgentCache(rmq.config) // dereg is a channel that allows us to instruct the goroutine that's listening for tests to stop. This allows us to re-register to a new command dereg := make(chan bool) rereg: group := ac.GetKeyValue("group") // if the group is nothing, rewrite to "mull". This is being done for now so that we don't have to worry if the goroutine was started or not // This way, it's always running, but if the agent is not in a group, it's listening on the "null" queue, which never has anything on it. // This is a minor waste of resources on the agent, so TODO(mierdin): you should probably fix this at some point and figure out how to only run // the goroutine when needed, but at the same time prevent the dereg channel from blocking unnecessarily in that case // // This will also handle the cases when the agent first starts up, and the key for this group isn't present in the cache, and therefore is "". if group == "" { group = "null" } go func() { for { err := rmq.ListenForGroupTasks(group, dereg) if err != nil { log.Warn("ListenForGroupTasks reported a failure. Trying again...") } } }() // Loop until the unackedGroup flag is set for { time.Sleep(2 * time.Second) // The key "unackedGroup" stores a "true" or "false" to indicate that there has been a group change that we need to acknowledge (handle) if ac.GetKeyValue("unackedGroup") == "true" { // This will kill the underlying goroutine, and in effect stop listening to the old queue. dereg <- true // Finally, set the "unackedGroup" to indicate that we've acknowledged the group change, and go back to the "rereg" label // to re-register onto the new group name ac.SetKeyValue("unackedGroup", "false") goto rereg } } }
// Run contains the logic necessary to perform this task on the agent. // This particular task will install, but not execute a testrun on this agent. // The installation procedure will first run the referenced testlet in check mode // to help ensure that the actual testrun execution can take place. If that // succeeds, the testrun is installed in the agent cache. func (itt InstallTestRunTask) Run() error { if itt.Tr.Testlet == "" { log.Error("Testlet parameter for this testrun is null") return errors.New("Testlet parameter for this testrun is null") } // Generate path to testlet and make sure it exists. testlet_path := fmt.Sprintf("%s/assets/testlets/%s", itt.Config.LocalResources.OptDir, itt.Tr.Testlet) if _, err := os.Stat(testlet_path); os.IsNotExist(err) { log.Errorf("Testlet %s does not exist on this agent", itt.Tr.Testlet) return errors.New("Error installing testrun - testlet doesn't exist on this agent.") } // Run the testlet in check mode to verify that everything is okay to run this test log.Debug("Running testlet in check mode: ", testlet_path) cmd := exec.Command(testlet_path, "check") // Stdout buffer cmdOutput := &bytes.Buffer{} // Attach buffer to command cmd.Stdout = cmdOutput // Execute collector cmd.Run() // This is probably the best cross-platform way to see if check mode passed. if strings.Contains(string(cmdOutput.Bytes()), "Check mode PASSED") { log.Debugf("Check mode for %s passed", testlet_path) } else { log.Error("Testlet returned an error during check mode: ", string(cmdOutput.Bytes())) return errors.New("Testlet returned an error during check mode") } // Insert testrun into agent cache var ac = cache.NewAgentCache(itt.Config) err := ac.InsertTestRun(itt.Tr) if err != nil { log.Error(err) return errors.New("Problem installing test run into agent cache") } return nil }
// Run contains the logic necessary to perform this task on the agent. func (sgt SetGroupTask) Run() error { var ac = cache.NewAgentCache(sgt.Config) // First, see what the current group is. If it matches what this task is instructing, we don't need to do anything. if ac.GetKeyValue("group") != sgt.GroupName { err := ac.SetKeyValue("group", sgt.GroupName) if err != nil { return errors.New(fmt.Sprintf("Failed to set keyvalue pair - %s:%s", "group", sgt.GroupName)) } err = ac.SetKeyValue("unackedGroup", "true") if err != nil { return errors.New(fmt.Sprintf("Failed to set keyvalue pair - %s:%s", "unackedGroup", "true")) } } else { log.Info("Already in the group being dictated by the server: ", sgt.GroupName) } return nil }
func main() { cfg := config.GetConfig(arg_config) // Set up cache var ac = cache.NewAgentCache(cfg) ac.Init() // Generate UUID uuid := hostresources.GenerateUuid() ac.SetKeyValue("uuid", uuid) log.Infof("ToDD Agent Activated: %s", uuid) // Start test data reporting service go testing.WatchForFinishedTestRuns(cfg) // Construct comms package var tc = comms.NewToDDComms(cfg) // Spawn goroutine to listen for tasks issued by server go func() { for { err := tc.CommsPackage.ListenForTasks(uuid) if err != nil { log.Warn("ListenForTasks reported a failure. Trying again...") } } }() // Watch for changes to group membership go tc.CommsPackage.WatchForGroup() // Continually advertise agent status into message queue for { // Gather assets here as a map, and refer to a key in that map in the below struct gatheredAssets := GetLocalAssets(cfg) var defaultaddr string if cfg.LocalResources.IPAddrOverride != "" { defaultaddr = cfg.LocalResources.IPAddrOverride } else { defaultaddr = hostresources.GetIPOfInt(cfg.LocalResources.DefaultInterface).String() } // Create an AgentAdvert instance to represent this particular agent me := defs.AgentAdvert{ Uuid: uuid, DefaultAddr: defaultaddr, FactCollectors: gatheredAssets["factcollectors"], Testlets: gatheredAssets["testlets"], Facts: facts.GetFacts(cfg), LocalTime: time.Now().UTC(), } // Advertise this agent err := tc.CommsPackage.AdvertiseAgent(me) if err != nil { log.Error("Failed to advertise agent after several retries") } time.Sleep(15 * time.Second) // TODO(moswalt): make configurable } }
// 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 }