<commit_before>package empire

import (
	"fmt"
	"time"

	"github.com/coreos/fleet/log"
	"github.com/remind101/empire/empire/scheduler"
)

// Manager is responsible for talking to the scheduler to schedule jobs onto the
// cluster.
type Manager interface {
	// ScheduleRelease schedules a release onto the cluster.
	ScheduleRelease(*Release, *Config, *Slug, Formation) error

	// ScaleRelease scales a release based on a process quantity map.
	ScaleRelease(*Release, *Config, *Slug, Formation, ProcessQuantityMap) error
}

// manager is a base implementation of the Manager interface.
type manager struct {
	JobsService
	ProcessesRepository
}

// ScheduleRelease creates jobs for every process and instance count and
// schedules them onto the cluster.
func (m *manager) ScheduleRelease(release *Release, config *Config, slug *Slug, formation Formation) error {
	// Find any existing jobs that have been scheduled for this app.
	existing, err := m.JobsService.JobsByApp(release.AppName)
	if err != nil {
		return err
	}

	jobs := buildJobs(
		release.AppName,
		release.Ver,
		slug.Image,
		config.Vars,
		formation,
	)

	err = m.JobsService.Schedule(jobs...)
	if err != nil {
		return err
	}

	go func() {
		time.Sleep(time.Second * 60)
		if err := m.JobsService.Unschedule(existing...); err != nil {
			// TODO What to do here?
			log.Errorf("Error unscheduling stale jobs: %s", err)
		}
	}()

	return nil
}

// ScaleRelease takes a release and process quantity map, and
// schedules/unschedules jobs to make the formation match the quantity map
func (m *manager) ScaleRelease(release *Release, config *Config, slug *Slug, formation Formation, qm ProcessQuantityMap) error {
	for t, q := range qm {
		if p, ok := formation[t]; ok {
			if err := m.scaleProcess(release, config, slug, p, q); err != nil {
				return err
			}
		}
	}

	return nil
}

func (m *manager) scaleProcess(release *Release, config *Config, slug *Slug, p *Process, q int) error {
	// Scale up
	if p.Quantity < q {
		for i := p.Quantity + 1; i <= q; i++ {
			err := m.JobsService.Schedule(
				&Job{
					AppName:        release.AppName,
					ReleaseVersion: release.Ver,
					ProcessType:    p.Type,
					Instance:       i,
					Environment:    config.Vars,
					Image:          slug.Image,
					Command:        p.Command,
				},
			)
			if err != nil {
				return err
			}
		}
	}

	// Scale down
	if p.Quantity > q {
		// Find existing jobs for this app
		existing, err := m.JobsService.JobsByApp(release.AppName)
		if err != nil {
			return err
		}

		// Create a map for easy lookup
		jm := make(map[scheduler.JobName]*Job, len(existing))
		for _, j := range existing {
			jm[j.JobName()] = j
		}

		// Unschedule jobs
		for i := p.Quantity; i > q; i-- {
			jobName := newJobName(release.AppName, release.Ver, p.Type, i)
			if j, ok := jm[jobName]; ok {
				m.JobsService.Unschedule(j)
			} else {
				return fmt.Errorf("Job not found to unschedule: %s", jobName)
			}
		}
	}

	// Update quantity for this process in the formation
	p.Quantity = q
	_, err := m.ProcessesRepository.Update(p)
	return err
}

// newJobName returns a new Name with the proper format.
func newJobName(name AppName, v ReleaseVersion, t ProcessType, i int) scheduler.JobName {
	return scheduler.JobName(fmt.Sprintf("%s.%d.%s.%d", name, v, t, i))
}

func buildJobs(name AppName, version ReleaseVersion, image Image, vars Vars, f Formation) []*Job {<commit_msg>Add ScaleDown and ScaleUp functions.<commit_after>package empire

import (
	"fmt"
	"time"

	"github.com/coreos/fleet/log"
	"github.com/remind101/empire/empire/scheduler"
)

// Manager is responsible for talking to the scheduler to schedule jobs onto the
// cluster.
type Manager interface {
	// ScheduleRelease schedules a release onto the cluster.
	ScheduleRelease(*Release, *Config, *Slug, Formation) error

	// ScaleRelease scales a release based on a process quantity map.
	ScaleRelease(*Release, *Config, *Slug, Formation, ProcessQuantityMap) error
}

// manager is a base implementation of the Manager interface.
type manager struct {
	JobsService
	ProcessesRepository
}

// ScheduleRelease creates jobs for every process and instance count and
// schedules them onto the cluster.
func (m *manager) ScheduleRelease(release *Release, config *Config, slug *Slug, formation Formation) error {
	// Find any existing jobs that have been scheduled for this app.
	existing, err := m.JobsService.JobsByApp(release.AppName)
	if err != nil {
		return err
	}

	jobs := buildJobs(
		release.AppName,
		release.Ver,
		slug.Image,
		config.Vars,
		formation,
	)

	err = m.JobsService.Schedule(jobs...)
	if err != nil {
		return err
	}

	go func() {
		time.Sleep(time.Second * 60)
		if err := m.JobsService.Unschedule(existing...); err != nil {
			// TODO What to do here?
			log.Errorf("Error unscheduling stale jobs: %s", err)
		}
	}()

	return nil
}

// ScaleRelease takes a release and process quantity map, and
// schedules/unschedules jobs to make the formation match the quantity map
func (m *manager) ScaleRelease(release *Release, config *Config, slug *Slug, formation Formation, qm ProcessQuantityMap) error {
	for t, q := range qm {
		if p, ok := formation[t]; ok {
			if err := m.scaleProcess(release, config, slug, p, q); err != nil {
				return err
			}
		}
	}

	return nil
}

func (m *manager) scaleProcess(release *Release, config *Config, slug *Slug, p *Process, q int) error {
	// Scale up
	if p.Quantity < q {
		jobs := ScaleUp(release, config, slug, p, q)
		if err := m.JobsService.Schedule(jobs...); err != nil {
			return err
		}
	}

	// Scale down
	if p.Quantity > q {
		// Find existing jobs for this app
		existing, err := m.JobsService.JobsByApp(release.AppName)
		if err != nil {
			return err
		}

		jobs := ScaleDown(existing, release, config, slug, p, q)

		if err := m.JobsService.Unschedule(jobs...); err != nil {
			return err
		}
	}

	// Update quantity for this process in the formation
	p.Quantity = q
	_, err := m.ProcessesRepository.Update(p)
	return err
}

// ScaleUp returns new Jobs to schedule when scaling up.
func ScaleUp(release *Release, config *Config, slug *Slug, p *Process, q int) []*Job {
	var jobs []*Job

	for i := p.Quantity + 1; i <= q; i++ {
		jobs = append(jobs, &Job{
			AppName:        release.AppName,
			ReleaseVersion: release.Ver,
			ProcessType:    p.Type,
			Instance:       i,
			Environment:    config.Vars,
			Image:          slug.Image,
			Command:        p.Command,
		})
	}

	return jobs
}

// ScaleDown returns Jobs to unschedule when scaling down.
func ScaleDown(existing []*Job, release *Release, config *Config, slug *Slug, p *Process, q int) []*Job {
	// Create a map for easy lookup
	jm := make(map[scheduler.JobName]*Job, len(existing))
	for _, j := range existing {
		jm[j.JobName()] = j
	}

	var jobs []*Job

	// Unschedule jobs
	for i := p.Quantity; i > q; i-- {
		jobName := newJobName(release.AppName, release.Ver, p.Type, i)
		if j, ok := jm[jobName]; ok {
			jobs = append(jobs, j)
		}
	}

	return jobs
}

// newJobName returns a new Name with the proper format.
func newJobName(name AppName, v ReleaseVersion, t ProcessType, i int) scheduler.JobName {
	return scheduler.JobName(fmt.Sprintf("%s.%d.%s.%d", name, v, t, i))
}

func buildJobs(name AppName, version ReleaseVersion, image Image, vars Vars, f Formation) []*Job {
{'fraction_non_alphanumeric': 0.12085956747878455, 'fraction_numerical': 0.0016424856282507528, 'mean_word_length': 3.308372641509434, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, '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': '5677223', 'n_tokens_mistral': 2643, 'n_tokens_neox': 2363, 'n_words': 802}