示例#1
0
文件: grouping.go 项目: twm1010/todd
// 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)
	}
}
示例#2
0
文件: rabbitmq.go 项目: twm1010/todd
// 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
}