<commit_before>package registry

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"path"
	"time"

	"github.com/coreos/go-etcd/etcd"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
)

const (
	keyPrefix      = "/coreos.com/coreinit/"
	lockPrefix     = "/locks/"
	machinePrefix  = "/machines/"
	requestPrefix  = "/request/"
	schedulePrefix = "/schedule/"
	jobWatchPrefix = "/watch/"
	statePrefix    = "/state/"
)

type Registry struct {
	etcd *etcd.Client
}

func New() (registry *Registry) {
	etcdC := etcd.NewClient(nil)
	etcdC.SetConsistency(etcd.WEAK_CONSISTENCY)
	return &Registry{etcdC}
}

// Describe the list of all known Machines
func (r *Registry) GetActiveMachines() []machine.Machine {
	key := path.Join(keyPrefix, machinePrefix)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return make([]machine.Machine, 0)
	}

	machines := make([]machine.Machine, 0)
	for _, kv := range resp.Kvs {
		_, bootId := path.Split(kv.Key)
		machine := machine.New(bootId)

		// This is a hacky way of telling if a Machine is reporting state
		addrs := r.getMachineAddrs(machine)
		if len(addrs) > 0 {
			machines = append(machines, *machine)
		}
	}

	return machines
}

func (r *Registry) getMachineAddrs(m *machine.Machine) []machine.IPAddress {
	key := path.Join(keyPrefix, machinePrefix, m.BootId, "addrs")
	resp, err := r.etcd.Get(key, false, true)

	addrs := make([]machine.IPAddress, 0)

	// Assume this is KeyNotFound and return an empty data structure
	if err != nil {
		return addrs
	}

	//TODO: Handle the error generated by unmarshal
	unmarshal(resp.Value, &addrs)

	return addrs
}

func (r *Registry) SetMachineAddrs(machine *machine.Machine, addrs []machine.IPAddress, ttl time.Duration) {
	//TODO: Handle the error generated by marshal
	json, _ := marshal(addrs)
	key := path.Join(keyPrefix, machinePrefix, machine.BootId, "addrs")
	r.etcd.Set(key, json, uint64(ttl.Seconds()))
}

func (r *Registry) AddRequest(req *job.JobRequest) {
	key := path.Join(keyPrefix, requestPrefix, req.ID.String())
	//TODO: Handle the error generated by marshal
	json, _ := marshal(req)
	r.etcd.Set(key, json, 0)
}

func (r *Registry) ResolveRequest(req *job.JobRequest) {
	key := path.Join(keyPrefix, requestPrefix, req.ID.String())
	r.etcd.Delete(key, true)
}

// Describe the list of jobs a given Machine is scheduled to run
func (r *Registry) GetMachineJobs(machine *machine.Machine) map[string]job.Job {
	key := path.Join(keyPrefix, machinePrefix, machine.BootId, schedulePrefix)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return make(map[string]job.Job, 0)
	}

	jobs := make(map[string]job.Job, len(resp.Kvs))
	for _, kv := range resp.Kvs {
		name := path.Base(kv.Key)

		var payload job.JobPayload
		err := unmarshal(kv.Value, &payload)

		if err == nil {
			j, _ := job.NewJob(name, nil, &payload)
			//FIXME: This will hide duplicate jobs!
			jobs[j.Name] = *j
		} else {
			log.Print(err)
		}
	}
	return jobs
}

func (r *Registry) GetScheduledJobs() map[string]job.Job {
	machines := r.GetActiveMachines()
	jobs := map[string]job.Job{}
	for _, mach := range machines {
		for name, j := range r.GetMachineJobs(&mach) {
			//FIXME: This will hide duplicate jobs!
			jobs[name] = j
		}
	}
	return jobs
}

func (r *Registry) GetJobPayload(j *job.Job) *job.JobPayload {
	key := path.Join(keyPrefix, schedulePrefix, j.Name)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return nil
	}

	var payload job.JobPayload
	//TODO: Handle the error generated by unmarshal
	unmarshal(resp.Value, &payload)
	return &payload
}

func (r *Registry) GetJobState(j *job.Job) *job.JobState {
	key := path.Join(keyPrefix, statePrefix, j.Name)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return nil
	}

	var state job.JobState
	//TODO: Handle the error generated by unmarshal
	unmarshal(resp.Value, &state)
	return &state
}

func (r *Registry) SaveJobWatch(watch *job.JobWatch) {
	key := path.Join(keyPrefix, jobWatchPrefix, watch.Payload.Name)
	//TODO: Handle the error generated by marshal
	json, _ := marshal(watch)
	r.etcd.Set(key, json, 0)
}

func (r *Registry) ScheduleMachineJob(job *job.Job, machine *machine.Machine) {
	key := path.Join(keyPrefix, machinePrefix, machine.BootId, schedulePrefix, job.Name)
	//TODO: Handle the error generated by marshal
	json, _ := marshal(job.Payload)
	log.Printf("Registry: setting key %s to value %s", key, json)
	r.etcd.Set(key, json, 0)
}

// StopJob removes the Job from any Machine's schedule. It also removes any
// relevant JobWatch objects.
func (r *Registry) StopJob(job *job.Job) {
	key := path.Join(keyPrefix, jobWatchPrefix, job.Name)
	r.etcd.Delete(key, true)

	for _, m := range r.GetActiveMachines() {
		name := fmt.Sprintf("%s.%s", m.BootId, job.Name)
		key := path.Join(keyPrefix, machinePrefix, m.BootId, schedulePrefix, name)
		r.etcd.Delete(key, true)<commit_msg>refactor(registry): Add method documentation to Registry<commit_after>package registry

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"path"
	"time"

	"github.com/coreos/go-etcd/etcd"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
)

const (
	keyPrefix      = "/coreos.com/coreinit/"
	lockPrefix     = "/locks/"
	machinePrefix  = "/machines/"
	requestPrefix  = "/request/"
	schedulePrefix = "/schedule/"
	jobWatchPrefix = "/watch/"
	statePrefix    = "/state/"
)

type Registry struct {
	etcd *etcd.Client
}

func New() (registry *Registry) {
	etcdC := etcd.NewClient(nil)
	etcdC.SetConsistency(etcd.WEAK_CONSISTENCY)
	return &Registry{etcdC}
}

// Describe all active Machines
func (r *Registry) GetActiveMachines() []machine.Machine {
	key := path.Join(keyPrefix, machinePrefix)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return make([]machine.Machine, 0)
	}

	machines := make([]machine.Machine, 0)
	for _, kv := range resp.Kvs {
		_, bootId := path.Split(kv.Key)
		machine := machine.New(bootId)

		// This is a hacky way of telling if a Machine is reporting state
		addrs := r.getMachineAddrs(machine)
		if len(addrs) > 0 {
			machines = append(machines, *machine)
		}
	}

	return machines
}

func (r *Registry) getMachineAddrs(m *machine.Machine) []machine.IPAddress {
	key := path.Join(keyPrefix, machinePrefix, m.BootId, "addrs")
	resp, err := r.etcd.Get(key, false, true)

	addrs := make([]machine.IPAddress, 0)

	// Assume this is KeyNotFound and return an empty data structure
	if err != nil {
		return addrs
	}

	//TODO: Handle the error generated by unmarshal
	unmarshal(resp.Value, &addrs)

	return addrs
}

// Push Machine Addr data to coreinit
func (r *Registry) SetMachineAddrs(machine *machine.Machine, addrs []machine.IPAddress, ttl time.Duration) {
	//TODO: Handle the error generated by marshal
	json, _ := marshal(addrs)
	key := path.Join(keyPrefix, machinePrefix, machine.BootId, "addrs")
	r.etcd.Set(key, json, uint64(ttl.Seconds()))
}

// Submit a new JobRequest to coreinit
func (r *Registry) AddRequest(req *job.JobRequest) {
	key := path.Join(keyPrefix, requestPrefix, req.ID.String())
	//TODO: Handle the error generated by marshal
	json, _ := marshal(req)
	r.etcd.Set(key, json, 0)
}

// Remove a given JobRequest from coreinit
func (r *Registry) ResolveRequest(req *job.JobRequest) {
	key := path.Join(keyPrefix, requestPrefix, req.ID.String())
	r.etcd.Delete(key, true)
}

// List the jobs a given Machine is scheduled to run
func (r *Registry) GetMachineJobs(machine *machine.Machine) map[string]job.Job {
	key := path.Join(keyPrefix, machinePrefix, machine.BootId, schedulePrefix)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return make(map[string]job.Job, 0)
	}

	jobs := make(map[string]job.Job, len(resp.Kvs))
	for _, kv := range resp.Kvs {
		name := path.Base(kv.Key)

		var payload job.JobPayload
		err := unmarshal(kv.Value, &payload)

		if err == nil {
			j, _ := job.NewJob(name, nil, &payload)
			//FIXME: This will hide duplicate jobs!
			jobs[j.Name] = *j
		} else {
			log.Print(err)
		}
	}
	return jobs
}

// List the jobs all Machines are scheduled to run
func (r *Registry) GetScheduledJobs() map[string]job.Job {
	machines := r.GetActiveMachines()
	jobs := map[string]job.Job{}
	for _, mach := range machines {
		for name, j := range r.GetMachineJobs(&mach) {
			//FIXME: This will hide duplicate jobs!
			jobs[name] = j
		}
	}
	return jobs
}

func (r *Registry) GetJobPayload(j *job.Job) *job.JobPayload {
	key := path.Join(keyPrefix, schedulePrefix, j.Name)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return nil
	}

	var payload job.JobPayload
	//TODO: Handle the error generated by unmarshal
	unmarshal(resp.Value, &payload)
	return &payload
}

// Get the current JobState of the provided Job
func (r *Registry) GetJobState(j *job.Job) *job.JobState {
	key := path.Join(keyPrefix, statePrefix, j.Name)
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return nil
	}

	var state job.JobState
	//TODO: Handle the error generated by unmarshal
	unmarshal(resp.Value, &state)
	return &state
}

// Add a new JobWatch to coreinit
func (r *Registry) SaveJobWatch(watch *job.JobWatch) {
	key := path.Join(keyPrefix, jobWatchPrefix, watch.Payload.Name)
	//TODO: Handle the error generated by marshal
	json, _ := marshal(watch)
	r.etcd.Set(key, json, 0)
}

// Schedule a Job to a given Machine
func (r *Registry) ScheduleMachineJob(job *job.Job, machine *machine.Machine) {
	key := path.Join(keyPrefix, machinePrefix, machine.BootId, schedulePrefix, job.Name)
	//TODO: Handle the error generated by marshal
	json, _ := marshal(job.Payload)
	log.Printf("Registry: setting key %s to value %s", key, json)
	r.etcd.Set(key, json, 0)
}

// StopJob removes the Job from any Machine's schedule. It also removes any
// relevant JobWatch objects.
func (r *Registry) StopJob(job *job.Job) {
	key := path.Join(keyPrefix, jobWatchPrefix, job.Name)
	r.etcd.Delete(key, true)

	for _, m := range r.GetActiveMachines() {
		name := fmt.Sprintf("%s.%s", m.BootId, job.Name)
		key := path.Join(keyPrefix, machinePrefix, m.BootId, schedulePrefix, name)
		r.etcd.Delete(key, true)
{'fraction_non_alphanumeric': 0.1357504215851602, 'fraction_numerical': 0.001873711823121604, 'mean_word_length': 4.035377358490566, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3701572', 'n_tokens_mistral': 3892, 'n_tokens_neox': 3715, 'n_words': 1065}