// 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
}
// ListenForAgent will listen on the message queue for new agent advertisements.
// It is meant to be run as a goroutine
func (rmq rabbitMQComms) ListenForAgent(assets map[string]map[string]string) {
// TODO(mierdin): does func param need to be a pointer?
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()
q, err := ch.QueueDeclare(
"agentadvert", // name
false, // durable
false, // delete when unused
false, // exclusive
false, // no-wait
nil, // arguments
)
if err != nil {
log.Error("Failed to declare a queue")
os.Exit(1)
}
msgs, err := ch.Consume(
q.Name, // queue
"agentadvert", // 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 {
log.Debugf("Agent advertisement recieved: %s", d.Body)
var agent defs.AgentAdvert
err = json.Unmarshal(d.Body, &agent)
// TODO(mierdin): Need to handle this error
// assetList is a slice that will contain any URLs that need to be sent to an
// agent as a response to an incorrect or incomplete list of assets
var assetList []string
// assets is the asset map from the SERVER's perspective
for asset_type, asset_hashes := range assets {
var agentAssets map[string]string
// agentAssets is the asset map from the AGENT's perspective
if asset_type == "factcollectors" {
agentAssets = agent.FactCollectors
} else if asset_type == "testlets" {
agentAssets = agent.Testlets
}
for name, hash := range asset_hashes {
// See if the hashes match (a missing asset will also result in False)
if agentAssets[name] != hash {
// hashes do not match, so we need to append the asset download URL to the remediate list
var default_ip string
if rmq.config.LocalResources.IPAddrOverride != "" {
default_ip = rmq.config.LocalResources.IPAddrOverride
} else {
default_ip = hostresources.GetIPOfInt(rmq.config.LocalResources.DefaultInterface).String()
}
asset_url := fmt.Sprintf("http://%s:%s/%s/%s", default_ip, rmq.config.Assets.Port, asset_type, name)
assetList = append(assetList, asset_url)
}
}
}
// Asset list is empty, so we can continue
if len(assetList) == 0 {
var tdb, _ = db.NewToddDB(rmq.config)
tdb.SetAgent(agent)
// This block of code checked that the agent time was within a certain range of the server time. If there was a large enough
// time skew, the agent advertisement would be rejected.
// I have disabled this for now - My plan was to use this to synchronize testrun execution amongst agents, but I have
// a solution to that for now. May revisit this later.
//
// Determine difference between server and agent time
// t1 := time.Now()
// var diff float64
// diff = t1.Sub(agent.LocalTime).Seconds()
//
// // If Agent is within half a second of server time, add insert to database
// if diff < 0.5 && diff > -0.5 {
// } else {
// // We don't want to register an agent if there is a severe time difference,
// // in order to ensure continuity during tests. So, just print log message.
// log.Warn("Agent time not within boundaries.")
// }
} else {
log.Warnf("Agent %s did not have the required asset files. This advertisement is ignored.", agent.Uuid)
var task tasks.DownloadAssetTask
task.Type = "DownloadAsset" //TODO(mierdin): This is an extra step. Maybe a factory function for the task could help here?
task.Assets = assetList
rmq.SendTask(agent.Uuid, task)
}
}
}()
log.Infof(" [*] Waiting for messages. To exit press CTRL+C")
<-forever
}