From 43c6337a254b9b569089bf30fd5ac489d87b0a3c Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Mon, 16 Jul 2018 12:52:24 -0500
Subject: [PATCH] Structs and validation for spread

---
 nomad/structs/funcs.go   | 26 ++++++++++++++++++
 nomad/structs/structs.go | 59 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+)

diff --git a/nomad/structs/funcs.go b/nomad/structs/funcs.go
index 6a52cd49a..2e4910181 100644
--- a/nomad/structs/funcs.go
+++ b/nomad/structs/funcs.go
@@ -221,6 +221,32 @@ func CopySliceAffinities(s []*Affinity) []*Affinity {
 	return c
 }
 
+func CopySliceSpreads(s []*Spread) []*Spread {
+	l := len(s)
+	if l == 0 {
+		return nil
+	}
+
+	c := make([]*Spread, l)
+	for i, v := range s {
+		c[i] = v.Copy()
+	}
+	return c
+}
+
+func CopySliceSpreadTarget(s []*SpreadTarget) []*SpreadTarget {
+	l := len(s)
+	if l == 0 {
+		return nil
+	}
+
+	c := make([]*SpreadTarget, l)
+	for i, v := range s {
+		c[i] = v.Copy()
+	}
+	return c
+}
+
 // VaultPoliciesSet takes the structure returned by VaultPolicies and returns
 // the set of required policies
 func VaultPoliciesSet(policies map[string]map[string]*Vault) []string {
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index b7fbb3831..c4b33cb9a 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -2008,6 +2008,10 @@ type Job struct {
 	// scheduling preferences that apply to all groups and tasks
 	Affinities []*Affinity
 
+	// Spread can be specified at the job level to express spreading
+	// allocations across a desired attribute, such as datacenter
+	Spreads []*Spread
+
 	// TaskGroups are the collections of task groups that this job needs
 	// to run. Each task group is an atomic unit of scheduling and placement.
 	TaskGroups []*TaskGroup
@@ -3336,6 +3340,10 @@ type TaskGroup struct {
 	// Affinities can be specified at the task group level to express
 	// scheduling preferences.
 	Affinities []*Affinity
+
+	// Spread can be specified at the task group level to express spreading
+	// allocations across a desired attribute, such as datacenter
+	Spreads []*Spread
 }
 
 func (tg *TaskGroup) Copy() *TaskGroup {
@@ -3349,6 +3357,7 @@ func (tg *TaskGroup) Copy() *TaskGroup {
 	ntg.RestartPolicy = ntg.RestartPolicy.Copy()
 	ntg.ReschedulePolicy = ntg.ReschedulePolicy.Copy()
 	ntg.Affinities = CopySliceAffinities(ntg.Affinities)
+	ntg.Spreads = CopySliceSpreads(ntg.Spreads)
 
 	if tg.Tasks != nil {
 		tasks := make([]*Task, len(ntg.Tasks))
@@ -5384,6 +5393,56 @@ func (a *Affinity) Validate() error {
 	return mErr.ErrorOrNil()
 }
 
+type Spread struct {
+	Attribute    string
+	Weight       int
+	SpreadTarget []*SpreadTarget
+	str          string
+}
+
+type SpreadTarget struct {
+	Value string
+	Ratio uint32
+	str   string
+}
+
+func (s *Spread) Copy() *Spread {
+	if s == nil {
+		return nil
+	}
+	ns := new(Spread)
+	*ns = *s
+
+	ns.SpreadTarget = CopySliceSpreadTarget(s.SpreadTarget)
+	return ns
+}
+
+func (s *SpreadTarget) Copy() *SpreadTarget {
+	if s == nil {
+		return nil
+	}
+
+	ns := new(SpreadTarget)
+	*ns = *s
+	return ns
+}
+
+func (s *Spread) String() string {
+	if s.str != "" {
+		return s.str
+	}
+	s.str = fmt.Sprintf("%s %s %v", s.Attribute, s.Weight, s.SpreadTarget)
+	return s.str
+}
+
+func (s *SpreadTarget) String() string {
+	if s.str != "" {
+		return s.str
+	}
+	s.str = fmt.Sprintf("%s %v", s.Value, s.Ratio)
+	return s.str
+}
+
 // EphemeralDisk is an ephemeral disk object
 type EphemeralDisk struct {
 	// Sticky indicates whether the allocation is sticky to a node