// ListenForResponses listens for responses from an agent func (rmq rabbitMQComms) ListenForResponses(stopListeningForResponses *chan bool) { queueName := "agentresponses" conn, err := amqp.Dial(rmq.queueUrl) if err != nil { log.Error(err) log.Error("Failed to connect to RabbitMQ") os.Exit(1) } 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( queueName, // 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( queueName, // queue queueName, // 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) } tdb, err := db.NewToddDB(rmq.config) // TODO(vcabbage): Consider moving this into the rabbitMQComms struct if err != nil { log.Error("Failed to connect to DB") os.Exit(1) } go func() { for d := range msgs { // Unmarshal into BaseResponse to determine type var base_msg responses.BaseResponse err = json.Unmarshal(d.Body, &base_msg) // TODO(mierdin): Need to handle this error log.Debugf("Agent response received: %s", d.Body) // call agent response method based on type switch base_msg.Type { case "AgentStatus": var sasr responses.SetAgentStatusResponse err = json.Unmarshal(d.Body, &sasr) // TODO(mierdin): Need to handle this error log.Debugf("Agent %s is '%s' regarding test %s. Writing to DB.", sasr.AgentUuid, sasr.Status, sasr.TestUuid) err := tdb.SetAgentTestStatus(sasr.TestUuid, sasr.AgentUuid, sasr.Status) if err != nil { log.Errorf("Error writing agent status to DB: %v", err) } case "TestData": var utdr responses.UploadTestDataResponse err = json.Unmarshal(d.Body, &utdr) // TODO(mierdin): Need to handle this error err = tdb.SetAgentTestData(utdr.TestUuid, utdr.AgentUuid, utdr.TestData) // TODO(mierdin): Need to handle this error // Send task to the agent that says to delete the entry var dtdt tasks.DeleteTestDataTask dtdt.Type = "DeleteTestData" //TODO(mierdin): This is an extra step. Maybe a factory function for the task could help here? dtdt.TestUuid = utdr.TestUuid rmq.SendTask(utdr.AgentUuid, dtdt) // Finally, set the status for this agent in the test to "finished" err := tdb.SetAgentTestStatus(dtdt.TestUuid, utdr.AgentUuid, "finished") if err != nil { log.Errorf("Error writing agent status to DB: %v", err) } default: log.Errorf(fmt.Sprintf("Unexpected type value for received response: %s", base_msg.Type)) } } }() log.Infof(" [*] Waiting for messages. To exit press CTRL+C") <-*stopListeningForResponses }
// 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 }