<commit_before>package smd

import (
	"math"
	"os"
	"time"

	"github.com/ChristopherRabotin/gokalman"
	"github.com/ChristopherRabotin/ode"
	kitlog "github.com/go-kit/kit/log"
	"github.com/gonum/matrix/mat64"
)

const (
	transitionPhiOnly = true
)

// OrbitEstimate is an ode.Integrable which allows to propagate an orbit via its initial estimate.
type OrbitEstimate struct {
	Φ      *mat64.Dense  // STM
	Orbit  Orbit         // estimated orbit
	Perts  Perturbations // perturbations to account for
	StopDT time.Time     // end time of te integration
	dt     time.Time     // current time of the integration
	step   time.Duration // time step
	logger kitlog.Logger // logger
}

// GetState gets the state.
func (e *OrbitEstimate) GetState() []float64 {
	rΦ, cΦ := e.Φ.Dims()
	s := make([]float64, 6+rΦ*cΦ)
	R, V := e.Orbit.RV()
	s[0] = R[0]
	s[1] = R[1]
	s[2] = R[2]
	s[3] = V[0]
	s[4] = V[1]
	s[5] = V[2]
	// Add the components of Φ
	sIdx := 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			s[sIdx] = e.Φ.At(i, j)
			sIdx++
		}
	}
	return s
}

// SetState sets the next state at time t.
func (e *OrbitEstimate) SetState(t float64, s []float64) {
	R := []float64{s[0], s[1], s[2]}
	V := []float64{s[3], s[4], s[5]}
	e.Orbit = *NewOrbitFromRV(R, V, e.Orbit.Origin)
	// Extract the components of Φ
	sIdx := 6
	rΦ, cΦ := e.Φ.Dims()
	Φk20 := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φk20.Set(i, j, s[sIdx])
			sIdx++
		}
	}
	if transitionPhiOnly {
		// Compute the Φ for this transition
		var Φinv mat64.Dense
		if err := Φinv.Inverse(e.Φ); err != nil {
			panic("could not invert e.Φ")
		}
		e.Φ.Mul(Φk20, &Φinv)
	} else {
		e.Φ = Φk20
	}
	// Increment the time.
	e.dt = e.dt.Add(e.step)
}

// Stop returns whether we should stop the integration.
func (e *OrbitEstimate) Stop(t float64) bool {
	return e.dt.After(e.StopDT)
}

// State returns the latest state
func (e *OrbitEstimate) State() State {
	return State{e.dt, Spacecraft{}, e.Orbit, nil, nil}
}

// Func does the math. Returns a new state.
func (e *OrbitEstimate) Func(t float64, f []float64) (fDot []float64) {
	// XXX: Note that this function is very similar to Mission.Func for a Cartesian propagation.
	// *BUT* we need to add in all the components of Φ, since they have to be integrated too.
	rΦ, cΦ := e.Φ.Dims()
	fDot = make([]float64, 6+rΦ*cΦ) // init return vector
	// Re-create the orbit from the state.
	R := []float64{f[0], f[1], f[2]}
	V := []float64{f[3], f[4], f[5]}
	orbit := NewOrbitFromRV(R, V, e.Orbit.Origin)
	bodyAcc := -orbit.Origin.μ / math.Pow(orbit.RNorm(), 3)
	// d\vec{R}/dt
	fDot[0] = f[3]
	fDot[1] = f[4]
	fDot[2] = f[5]
	// d\vec{V}/dt
	fDot[3] = bodyAcc * f[0]
	fDot[4] = bodyAcc * f[1]
	fDot[5] = bodyAcc * f[2]

	pert := e.Perts.Perturb(*orbit, e.dt)
	for i := 0; i < 6; i++ {
		fDot[i] += pert[i]
	}

	// Extract the components of Φ
	fIdx := 6
	Φ := mat64.NewDense(rΦ, cΦ, nil)
	ΦDot := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φ.Set(i, j, f[fIdx])
			fIdx++
		}
	}

	// Compute the STM.
	A := mat64.NewDense(6, 6, nil)
	// Top right is Identity 3x3
	A.Set(0, 3, 1)
	A.Set(1, 4, 1)
	A.Set(2, 5, 1)
	// Bottom left is where the magix is.
	x := R[0]
	y := R[1]
	z := R[2]
	x2 := math.Pow(R[0], 2)
	y2 := math.Pow(R[1], 2)
	z2 := math.Pow(R[2], 2)
	r2 := x2 + y2 + z2
	r232 := math.Pow(r2, 3/2.)
	r252 := math.Pow(r2, 5/2.)
	// Add the body perturbations

	dAxDx := 3*e.Orbit.Origin.μ*x2/r252 - e.Orbit.Origin.μ/r232
	dAxDy := 3 * e.Orbit.Origin.μ * x * y / r252
	dAxDz := 3 * e.Orbit.Origin.μ * x * z / r252
	dAyDx := 3 * e.Orbit.Origin.μ * x * y / r252
	dAyDy := 3*e.Orbit.Origin.μ*y2/r252 - e.Orbit.Origin.μ/r232
	dAyDz := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDx := 3 * e.Orbit.Origin.μ * x * z / r252
	dAzDy := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDz := 3*e.Orbit.Origin.μ*z2/r252 - e.Orbit.Origin.μ/r232

	A.Set(3, 0, dAxDx)
	A.Set(4, 0, dAyDx)
	A.Set(5, 0, dAzDx)
	A.Set(3, 1, dAxDy)
	A.Set(4, 1, dAyDy)
	A.Set(5, 1, dAzDy)
	A.Set(3, 2, dAxDz)
	A.Set(4, 2, dAyDz)
	A.Set(5, 2, dAzDz)

	// Jn perturbations:
	if e.Perts.Jn > 1 {
		// Ai0 = \frac{\partial a}{\partial x}
		// Ai1 = \frac{\partial a}{\partial y}
		// Ai2 = \frac{\partial a}{\partial z}
		A30 := A.At(3, 0)
		A40 := A.At(4, 0)
		A50 := A.At(5, 0)
		A31 := A.At(3, 1)
		A41 := A.At(4, 1)
		A51 := A.At(5, 1)
		A32 := A.At(3, 2)
		A42 := A.At(4, 2)
		A52 := A.At(5, 2)

		// Notation simplification
		z3 := math.Pow(R[2], 3)
		z4 := math.Pow(R[2], 4)
		// Adding those fractions to avoid forgetting the trailing period which makes them floats.
		f32 := 3 / 2.
		f152 := 15 / 2.
		r272 := math.Pow(r2, 7/2.)
		r292 := math.Pow(r2, 9/2.)
		// J2
		j2fact := e.Orbit.Origin.J(2) * math.Pow(e.Orbit.Origin.Radius, 2) * e.Orbit.Origin.μ
		A30 += -f32 * j2fact * (35*x2*z2/r292 - 5*x2/r272 - 5*z2/r272 + 1/r252) //dAxDx
		A40 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAyDx
		A50 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)                      //dAzDx

		A31 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAxDy
		A41 += -f32 * j2fact * (35*y2*z2/r292 - 5*y2/r272 - 5*z2/r272 + 1/r252) // dAyDy
		A51 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)                      // dAzDy

		A32 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)        //dAxDz
		A42 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)        //dAyDz
		A52 += -f32 * j2fact * (35*z4/r292 - 30*z2/r272 + 3/r252) // dAzDz

		// J3
		if e.Perts.Jn > 2 {
			z5 := math.Pow(R[2], 5)
			r2112 := math.Pow(r2, 11/2.)
			f52 := 5 / 2.
			f1052 := 105 / 2.
			j3fact := e.Orbit.Origin.J(3) * math.Pow(e.Orbit.Origin.Radius, 3) * e.Orbit.Origin.μ
			A30 += -f52 * j3fact * (63*x2*z3/r2112 - 21*x2*z/r292 - 7*z3/r292 + 3*z/r272) //dAxDx
			A40 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAyDx
			A50 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272)               //dAzDx

			A31 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAxDy
			A41 += -f52 * j3fact * (63*y2*z3/r2112 - 21*y2*z/r292 - 7*z3/r292 + 3*z/r272) // dAyDy
			A51 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272)               // dAzDy

			A32 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272) //dAxDz
			A42 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272) //dAyDz
			A52 += -f52 * j3fact * (63*z5/r2112 - 70*z3/r292 + 15*z/r272)   // dAzDz
		}
		// \frac{\partial a}{\partial x}
		A.Set(3, 0, A30)
		A.Set(4, 0, A40)
		A.Set(5, 0, A50)
		// \partial a/\partial y
		A.Set(3, 1, A31)
		A.Set(4, 1, A41)
		A.Set(5, 1, A51)
		// \partial a/\partial z
		A.Set(3, 2, A32)
		A.Set(4, 2, A42)
		A.Set(5, 2, A52)
	}
	ΦDot.Mul(A, Φ)

	// Store ΦDot in fDot
	fIdx = 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			fDot[fIdx] = ΦDot.At(i, j)
			fIdx++
		}
	}
	return fDot
}

// PropagateUntil propagates until the given time is reached.
func (e *OrbitEstimate) PropagateUntil(dt time.Time) {
	e.StopDT = dt
	ode.NewRK4(0, e.step.Seconds(), e).Solve() // Blocking.
}

// NewOrbitEstimate returns a new Estimate of an orbit given the perturbations to be taken into account.
// The only supported state is [\vec{r} \vec{v}]T (for now at least).
func NewOrbitEstimate(n string, o Orbit, p Perturbations, epoch time.Time, step time.Duration) *OrbitEstimate {
	// The initial previous STM is identity.
	klog := kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stdout))
	klog = kitlog.With(klog, "estimate", n)
	stopDT := epoch
	// XXX: We add the step for consistency with Mission. Mission is broken: it skips the first step because the time addition
	// happens in the Stop function instead of the SetState function, the former being called at the start of the integration.
	return &OrbitEstimate{gokalman.DenseIdentity(6), o, p, stopDT, epoch.Add(step), step, klog}
}
<commit_msg>Reset OrbitEstimate STM computation to be between start and end (instead of per time step).<commit_after>package smd

import (
	"math"
	"os"
	"time"

	"github.com/ChristopherRabotin/gokalman"
	"github.com/ChristopherRabotin/ode"
	kitlog "github.com/go-kit/kit/log"
	"github.com/gonum/matrix/mat64"
)

const (
	transitionPhiOnly = false
)

// OrbitEstimate is an ode.Integrable which allows to propagate an orbit via its initial estimate.
type OrbitEstimate struct {
	Φ      *mat64.Dense  // STM
	Orbit  Orbit         // estimated orbit
	Perts  Perturbations // perturbations to account for
	StopDT time.Time     // end time of te integration
	dt     time.Time     // current time of the integration
	step   time.Duration // time step
	logger kitlog.Logger // logger
}

// GetState gets the state.
func (e *OrbitEstimate) GetState() []float64 {
	rΦ, cΦ := e.Φ.Dims()
	s := make([]float64, 6+rΦ*cΦ)
	R, V := e.Orbit.RV()
	s[0] = R[0]
	s[1] = R[1]
	s[2] = R[2]
	s[3] = V[0]
	s[4] = V[1]
	s[5] = V[2]
	// Add the components of Φ
	sIdx := 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			s[sIdx] = e.Φ.At(i, j)
			sIdx++
		}
	}
	return s
}

// SetState sets the next state at time t.
func (e *OrbitEstimate) SetState(t float64, s []float64) {
	R := []float64{s[0], s[1], s[2]}
	V := []float64{s[3], s[4], s[5]}
	e.Orbit = *NewOrbitFromRV(R, V, e.Orbit.Origin)
	// Extract the components of Φ
	sIdx := 6
	rΦ, cΦ := e.Φ.Dims()
	Φk20 := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φk20.Set(i, j, s[sIdx])
			sIdx++
		}
	}
	if transitionPhiOnly {
		// Compute the Φ for this transition
		var Φinv mat64.Dense
		if err := Φinv.Inverse(e.Φ); err != nil {
			panic("could not invert e.Φ")
		}
		e.Φ.Mul(Φk20, &Φinv)
	} else {
		e.Φ = Φk20
	}
	// Increment the time.
	e.dt = e.dt.Add(e.step)
}

// Stop returns whether we should stop the integration.
func (e *OrbitEstimate) Stop(t float64) bool {
	return e.dt.After(e.StopDT)
}

// State returns the latest state
func (e *OrbitEstimate) State() State {
	return State{e.dt, Spacecraft{}, e.Orbit, nil, nil}
}

// Func does the math. Returns a new state.
func (e *OrbitEstimate) Func(t float64, f []float64) (fDot []float64) {
	// XXX: Note that this function is very similar to Mission.Func for a Cartesian propagation.
	// *BUT* we need to add in all the components of Φ, since they have to be integrated too.
	rΦ, cΦ := e.Φ.Dims()
	fDot = make([]float64, 6+rΦ*cΦ) // init return vector
	// Re-create the orbit from the state.
	R := []float64{f[0], f[1], f[2]}
	V := []float64{f[3], f[4], f[5]}
	orbit := NewOrbitFromRV(R, V, e.Orbit.Origin)
	bodyAcc := -orbit.Origin.μ / math.Pow(orbit.RNorm(), 3)
	// d\vec{R}/dt
	fDot[0] = f[3]
	fDot[1] = f[4]
	fDot[2] = f[5]
	// d\vec{V}/dt
	fDot[3] = bodyAcc * f[0]
	fDot[4] = bodyAcc * f[1]
	fDot[5] = bodyAcc * f[2]

	pert := e.Perts.Perturb(*orbit, e.dt)
	for i := 0; i < 6; i++ {
		fDot[i] += pert[i]
	}

	// Extract the components of Φ
	fIdx := 6
	Φ := mat64.NewDense(rΦ, cΦ, nil)
	ΦDot := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φ.Set(i, j, f[fIdx])
			fIdx++
		}
	}

	// Compute the STM.
	A := mat64.NewDense(6, 6, nil)
	// Top right is Identity 3x3
	A.Set(0, 3, 1)
	A.Set(1, 4, 1)
	A.Set(2, 5, 1)
	// Bottom left is where the magix is.
	x := R[0]
	y := R[1]
	z := R[2]
	x2 := math.Pow(R[0], 2)
	y2 := math.Pow(R[1], 2)
	z2 := math.Pow(R[2], 2)
	r2 := x2 + y2 + z2
	r232 := math.Pow(r2, 3/2.)
	r252 := math.Pow(r2, 5/2.)
	// Add the body perturbations

	dAxDx := 3*e.Orbit.Origin.μ*x2/r252 - e.Orbit.Origin.μ/r232
	dAxDy := 3 * e.Orbit.Origin.μ * x * y / r252
	dAxDz := 3 * e.Orbit.Origin.μ * x * z / r252
	dAyDx := 3 * e.Orbit.Origin.μ * x * y / r252
	dAyDy := 3*e.Orbit.Origin.μ*y2/r252 - e.Orbit.Origin.μ/r232
	dAyDz := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDx := 3 * e.Orbit.Origin.μ * x * z / r252
	dAzDy := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDz := 3*e.Orbit.Origin.μ*z2/r252 - e.Orbit.Origin.μ/r232

	A.Set(3, 0, dAxDx)
	A.Set(4, 0, dAyDx)
	A.Set(5, 0, dAzDx)
	A.Set(3, 1, dAxDy)
	A.Set(4, 1, dAyDy)
	A.Set(5, 1, dAzDy)
	A.Set(3, 2, dAxDz)
	A.Set(4, 2, dAyDz)
	A.Set(5, 2, dAzDz)

	// Jn perturbations:
	if e.Perts.Jn > 1 {
		// Ai0 = \frac{\partial a}{\partial x}
		// Ai1 = \frac{\partial a}{\partial y}
		// Ai2 = \frac{\partial a}{\partial z}
		A30 := A.At(3, 0)
		A40 := A.At(4, 0)
		A50 := A.At(5, 0)
		A31 := A.At(3, 1)
		A41 := A.At(4, 1)
		A51 := A.At(5, 1)
		A32 := A.At(3, 2)
		A42 := A.At(4, 2)
		A52 := A.At(5, 2)

		// Notation simplification
		z3 := math.Pow(R[2], 3)
		z4 := math.Pow(R[2], 4)
		// Adding those fractions to avoid forgetting the trailing period which makes them floats.
		f32 := 3 / 2.
		f152 := 15 / 2.
		r272 := math.Pow(r2, 7/2.)
		r292 := math.Pow(r2, 9/2.)
		// J2
		j2fact := e.Orbit.Origin.J(2) * math.Pow(e.Orbit.Origin.Radius, 2) * e.Orbit.Origin.μ
		A30 += -f32 * j2fact * (35*x2*z2/r292 - 5*x2/r272 - 5*z2/r272 + 1/r252) //dAxDx
		A40 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAyDx
		A50 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)                      //dAzDx

		A31 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAxDy
		A41 += -f32 * j2fact * (35*y2*z2/r292 - 5*y2/r272 - 5*z2/r272 + 1/r252) // dAyDy
		A51 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)                      // dAzDy

		A32 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)        //dAxDz
		A42 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)        //dAyDz
		A52 += -f32 * j2fact * (35*z4/r292 - 30*z2/r272 + 3/r252) // dAzDz

		// J3
		if e.Perts.Jn > 2 {
			z5 := math.Pow(R[2], 5)
			r2112 := math.Pow(r2, 11/2.)
			f52 := 5 / 2.
			f1052 := 105 / 2.
			j3fact := e.Orbit.Origin.J(3) * math.Pow(e.Orbit.Origin.Radius, 3) * e.Orbit.Origin.μ
			A30 += -f52 * j3fact * (63*x2*z3/r2112 - 21*x2*z/r292 - 7*z3/r292 + 3*z/r272) //dAxDx
			A40 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAyDx
			A50 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272)               //dAzDx

			A31 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAxDy
			A41 += -f52 * j3fact * (63*y2*z3/r2112 - 21*y2*z/r292 - 7*z3/r292 + 3*z/r272) // dAyDy
			A51 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272)               // dAzDy

			A32 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272) //dAxDz
			A42 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272) //dAyDz
			A52 += -f52 * j3fact * (63*z5/r2112 - 70*z3/r292 + 15*z/r272)   // dAzDz
		}
		// \frac{\partial a}{\partial x}
		A.Set(3, 0, A30)
		A.Set(4, 0, A40)
		A.Set(5, 0, A50)
		// \partial a/\partial y
		A.Set(3, 1, A31)
		A.Set(4, 1, A41)
		A.Set(5, 1, A51)
		// \partial a/\partial z
		A.Set(3, 2, A32)
		A.Set(4, 2, A42)
		A.Set(5, 2, A52)
	}
	ΦDot.Mul(A, Φ)

	// Store ΦDot in fDot
	fIdx = 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			fDot[fIdx] = ΦDot.At(i, j)
			fIdx++
		}
	}
	return fDot
}

// PropagateUntil propagates until the given time is reached.
func (e *OrbitEstimate) PropagateUntil(dt time.Time) {
	e.StopDT = dt
	ode.NewRK4(0, e.step.Seconds(), e).Solve() // Blocking.
}

// NewOrbitEstimate returns a new Estimate of an orbit given the perturbations to be taken into account.
// The only supported state is [\vec{r} \vec{v}]T (for now at least).
func NewOrbitEstimate(n string, o Orbit, p Perturbations, epoch time.Time, step time.Duration) *OrbitEstimate {
	// The initial previous STM is identity.
	klog := kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stdout))
	klog = kitlog.With(klog, "estimate", n)
	stopDT := epoch
	// XXX: We add the step for consistency with Mission. Mission is broken: it skips the first step because the time addition
	// happens in the Stop function instead of the SetState function, the former being called at the start of the integration.
	return &OrbitEstimate{gokalman.DenseIdentity(6), o, p, stopDT, epoch.Add(step), step, klog}
}
<|endoftext|>


<commit_before>package exchequer

import (
	"fmt"
	"reflect"
)

type I interface{}
type A []interface{}
type M map[string]interface{}

type PathDoesntExist string
type TypeCastIsntValid string

func NewPathDoesntExist(x interface{}) PathDoesntExist {
	return PathDoesntExist(fmt.Sprintf("%v", x))
}

func NewTypeCastIsntValid(x interface{}, cast string) TypeCastIsntValid {
	return TypeCastIsntValid(fmt.Sprintf("%v to %s (it is %v)", x, cast, reflect.TypeOf(x)))
}

func (path PathDoesntExist) Error() string {
	return "Path doesn't exist: " + string(path)
}

func (typ TypeCastIsntValid) Error() string {
	return "Type-cast isn't valid: " + string(typ)
}

var mapType reflect.Type = reflect.TypeOf(map[string]interface{}{})
func convertToMap(i interface{}) (map[string]interface{}, bool) {
	// See if it is a map first
	if m, ok := i.(map[string]interface{}); ok {
		return m, true
	}

	// Now try converting it
	v := reflect.ValueOf(i)
	if v.Type().ConvertibleTo(mapType) {
		return v.Convert(mapType).Interface().(map[string]interface{}), true
	}

	return nil, false
}

var arrayType reflect.Type = reflect.TypeOf([]interface{}{})
func convertToArray(i interface{}) ([]interface{}, bool) {
	// See if it is an array first
	if a, ok := i.([]interface{}); ok {
		return a, true
	}

	// Now try converting it
	v := reflect.ValueOf(i)
	if v.Type().ConvertibleTo(arrayType) {
		return v.Convert(arrayType).Interface().([]interface{}), true
	}

	return nil, false
}

func Get(i I, paths ...interface{}) (interface{}, error) {
	for _, path := range paths {
		if s, ok := path.(string); ok {
			if m, ok := convertToMap(i); ok {
				i = m[s]
				continue
			}
			return nil, NewPathDoesntExist(path)
		}
		if x, ok := path.(int); ok {
			if a, ok := convertToArray(i); ok {
				if x < 0 {
					x = len(a) + x
				}

				if x < 0 || x >= len(a) {
					return nil, NewPathDoesntExist(path)
				} else {
					i = a[x]
					continue
				}
			} else {
				return nil, NewPathDoesntExist(path)
			}
		}
	}

	return i, nil
}

func Set(i I, value interface{}, paths ...interface{}) error {
	for j, path := range paths {
		if s, ok := path.(string); ok {
			if m, ok := convertToMap(i); ok {
				if j < len(paths)-1 {
					if _, ok = m[s]; !ok {
						m[s] = make(map[string]interface{})
					}
					i = m[s]
					continue
				} else {
					m[s] = value
					return nil
				}
			} else {
				return NewPathDoesntExist(path)
			}
		}
		if x, ok := path.(int); ok {
			if a, ok := convertToArray(i); ok {
				if x < 0 {
					x = len(a) + x
				}

				if x < 0 || x >= len(a) {
					return NewPathDoesntExist(path)
				} else {
					if j < len(paths)-1 {
						i = a[x]
						continue
					} else {
						a[x] = value
						return nil
					}
				}
			} else {
				return NewPathDoesntExist(path)
			}
		}
	}
	return nil
}

func String(i I, paths ...interface{}) (string, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return "", err
	}

	if s, ok := i.(string); ok {
		return s, nil
	}

	return "", NewTypeCastIsntValid(i, "string")
}

func Int(i I, paths ...interface{}) (int, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return 0, err
	}

	if s, ok := i.(int); ok {
		return s, nil
	}

	return 0, NewTypeCastIsntValid(i, "int")
}

func Bool(i I, paths ...interface{}) (bool, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return false, err
	}

	if s, ok := i.(bool); ok {
		return s, nil
	}

	return false, NewTypeCastIsntValid(i, "bool")
}

func Float(i I, paths ...interface{}) (float64, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return 0, err
	}

	if s, ok := i.(float64); ok {
		return s, nil
	}

	return 0, NewTypeCastIsntValid(i, "float")
}

func Map(i I, paths ...interface{}) (M, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return nil, err
	}

	if s, ok := convertToMap(i); ok {
		return M(s), nil
	}

	return nil, NewTypeCastIsntValid(i, "map")
}

func Array(i I, paths ...interface{}) (A, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return nil, err
	}

	if s, ok := convertToArray(i); ok {
		return A(s), nil
	}

	return nil, NewTypeCastIsntValid(i, "array")
}

type Q struct {
	i I
	prefix []interface{}
}

func New(i I, prefix_paths...interface{}) *Q {
	return &Q{i, prefix_paths}
}
func (q *Q) Prefix(paths ...interface{}) *Q {
	return New(q.i, append(q.prefix, paths...)...)
}
func (q *Q) I() interface{} {
	return q.i
}
func (q *Q) Q(paths ...interface{}) (*Q, error) {
	if i, err := Get(q.i, append(q.prefix, paths...)...); err != nil {
		return nil, err
	} else {
		return New(i), nil
	}
}
func (q *Q) Get(paths ...interface{}) (interface{}, error) {
	return Get(q.i, append(q.prefix, paths...)...)
}
func (q *Q) String(paths ...interface{}) (string, error) {
	return String(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Int(paths ...interface{}) (int, error) {
	return Int(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Bool(paths ...interface{}) (bool, error) {
	return Bool(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Float(paths ...interface{}) (float64, error) {
	return Float(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Map(paths ...interface{}) (M, error) {
	return Map(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Array(paths ...interface{}) (A, error) {
	return Array(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Set(value interface{}, paths ...interface{}) error {
	return Set(q.i, value, append(q.prefix, paths...)...)
}<commit_msg>Make sure to check for nil<commit_after>package exchequer

import (
	"fmt"
	"reflect"
)

type I interface{}
type A []interface{}
type M map[string]interface{}

type PathDoesntExist string
type TypeCastIsntValid string

func NewPathDoesntExist(x interface{}) PathDoesntExist {
	return PathDoesntExist(fmt.Sprintf("%v", x))
}

func NewTypeCastIsntValid(x interface{}, cast string) TypeCastIsntValid {
	return TypeCastIsntValid(fmt.Sprintf("%v to %s (it is %v)", x, cast, reflect.TypeOf(x)))
}

func (path PathDoesntExist) Error() string {
	return "Path doesn't exist: " + string(path)
}

func (typ TypeCastIsntValid) Error() string {
	return "Type-cast isn't valid: " + string(typ)
}

var mapType reflect.Type = reflect.TypeOf(map[string]interface{}{})
func convertToMap(i interface{}) (map[string]interface{}, bool) {
	// See if it is a map first
	if m, ok := i.(map[string]interface{}); ok {
		return m, true
	} else if i == nil {
		return nil, false
	}

	// Now try converting it
	v := reflect.ValueOf(i)
	if v.Type().ConvertibleTo(mapType) {
		return v.Convert(mapType).Interface().(map[string]interface{}), true
	}

	return nil, false
}

var arrayType reflect.Type = reflect.TypeOf([]interface{}{})
func convertToArray(i interface{}) ([]interface{}, bool) {
	// See if it is an array first
	if a, ok := i.([]interface{}); ok {
		return a, true
	} else if i == nil {
		return nil, false
	}

	// Now try converting it
	v := reflect.ValueOf(i)
	if v.Type().ConvertibleTo(arrayType) {
		return v.Convert(arrayType).Interface().([]interface{}), true
	}

	return nil, false
}

func Get(i I, paths ...interface{}) (interface{}, error) {
	for _, path := range paths {
		if s, ok := path.(string); ok {
			if m, ok := convertToMap(i); ok {
				i = m[s]
				continue
			}
			return nil, NewPathDoesntExist(path)
		}
		if x, ok := path.(int); ok {
			if a, ok := convertToArray(i); ok {
				if x < 0 {
					x = len(a) + x
				}

				if x < 0 || x >= len(a) {
					return nil, NewPathDoesntExist(path)
				} else {
					i = a[x]
					continue
				}
			} else {
				return nil, NewPathDoesntExist(path)
			}
		}
	}

	return i, nil
}

func Set(i I, value interface{}, paths ...interface{}) error {
	for j, path := range paths {
		if s, ok := path.(string); ok {
			if m, ok := convertToMap(i); ok {
				if j < len(paths)-1 {
					if _, ok = m[s]; !ok {
						m[s] = make(map[string]interface{})
					}
					i = m[s]
					continue
				} else {
					m[s] = value
					return nil
				}
			} else {
				return NewPathDoesntExist(path)
			}
		}
		if x, ok := path.(int); ok {
			if a, ok := convertToArray(i); ok {
				if x < 0 {
					x = len(a) + x
				}

				if x < 0 || x >= len(a) {
					return NewPathDoesntExist(path)
				} else {
					if j < len(paths)-1 {
						i = a[x]
						continue
					} else {
						a[x] = value
						return nil
					}
				}
			} else {
				return NewPathDoesntExist(path)
			}
		}
	}
	return nil
}

func String(i I, paths ...interface{}) (string, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return "", err
	}

	if s, ok := i.(string); ok {
		return s, nil
	}

	return "", NewTypeCastIsntValid(i, "string")
}

func Int(i I, paths ...interface{}) (int, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return 0, err
	}

	if s, ok := i.(int); ok {
		return s, nil
	}

	return 0, NewTypeCastIsntValid(i, "int")
}

func Bool(i I, paths ...interface{}) (bool, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return false, err
	}

	if s, ok := i.(bool); ok {
		return s, nil
	}

	return false, NewTypeCastIsntValid(i, "bool")
}

func Float(i I, paths ...interface{}) (float64, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return 0, err
	}

	if s, ok := i.(float64); ok {
		return s, nil
	}

	return 0, NewTypeCastIsntValid(i, "float")
}

func Map(i I, paths ...interface{}) (M, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return nil, err
	}

	if s, ok := convertToMap(i); ok {
		return M(s), nil
	}

	return nil, NewTypeCastIsntValid(i, "map")
}

func Array(i I, paths ...interface{}) (A, error) {
	i, err := Get(i, paths...)
	if err != nil {
		return nil, err
	}

	if s, ok := convertToArray(i); ok {
		return A(s), nil
	}

	return nil, NewTypeCastIsntValid(i, "array")
}

type Q struct {
	i I
	prefix []interface{}
}

func New(i I, prefix_paths...interface{}) *Q {
	return &Q{i, prefix_paths}
}
func (q *Q) Prefix(paths ...interface{}) *Q {
	return New(q.i, append(q.prefix, paths...)...)
}
func (q *Q) I() interface{} {
	return q.i
}
func (q *Q) Q(paths ...interface{}) (*Q, error) {
	if i, err := Get(q.i, append(q.prefix, paths...)...); err != nil {
		return nil, err
	} else {
		return New(i), nil
	}
}
func (q *Q) Get(paths ...interface{}) (interface{}, error) {
	return Get(q.i, append(q.prefix, paths...)...)
}
func (q *Q) String(paths ...interface{}) (string, error) {
	return String(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Int(paths ...interface{}) (int, error) {
	return Int(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Bool(paths ...interface{}) (bool, error) {
	return Bool(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Float(paths ...interface{}) (float64, error) {
	return Float(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Map(paths ...interface{}) (M, error) {
	return Map(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Array(paths ...interface{}) (A, error) {
	return Array(q.i, append(q.prefix, paths...)...)
}
func (q *Q) Set(value interface{}, paths ...interface{}) error {
	return Set(q.i, value, append(q.prefix, paths...)...)
}<|endoftext|>


<commit_before>package expo

import (
	"testing"

	"github.com/jarcoal/httpmock"
)

const token = "ExponentPushToken[xxxxxxxxxxxxxxxxxxxxxx]"

func TestIsExpoPushToken(t *testing.T) {
	if IsExpoPushToken("badToken") {
		t.Errorf("IsExpoPushToken returned unexpected value: got true want false")
	}

	if !IsExpoPushToken(token) {
		t.Errorf("IsExpoPushToken returned unexpected value: got false want true")
	}
}

func TestChunkPushNotifications(t *testing.T) {
	messages := []*PushMessage{
		{To: "token"},
		{To: "token"},
		{To: "token"},
	}

	chunks := ChunkPushNotifications(messages)
	if len(chunks) > 1 {
		t.Errorf("ChunkPushNotifications returned unexpected chunks: chunks length got %v want 1", len(chunks))
	}

	ChunkLimit = 2
	chunks = ChunkPushNotifications(messages)
	if len(chunks) != 2 {
		t.Errorf("ChunkPushNotifications returned unexpected chunks: chunks length got %v want 2", len(chunks))
	}

	ChunkLimit = 1
	chunks = ChunkPushNotifications(messages)
	if len(chunks) != 3 {
		t.Errorf("ChunkPushNotifications returned unexpected chunks: chunks length got %v want 3", len(chunks))
	}
}

func TestSendPushNotification(t *testing.T) {
	httpmock.Activate()
	defer httpmock.DeactivateAndReset()

	message := PushMessage{
		To:    token,
		Title: "Notification title",
		Body:  "Notification content"}

	httpmock.RegisterResponder("POST", baseAPIURL+"/push/send",
		httpmock.NewStringResponder(200, `{"data": [{"status": "ok"}]}`))

	api, _ := message.Send()
	if api.Data[0].Status != "ok" {
		t.Errorf("SendPushNotification returned unexpected response: status got %s want ok", api.Data[0].Status)
	}
}
<commit_msg>add unit test (#3)<commit_after>package expo

import (
	"net/http"
	"testing"

	"github.com/jarcoal/httpmock"
)

const token = "ExponentPushToken[xxxxxxxxxxxxxxxxxxxxxx]"
const status200 = `{"data": [{"status": "ok"}]}`

func TestIsExpoPushToken(t *testing.T) {
	if IsExpoPushToken("badToken") {
		t.Errorf("IsExpoPushToken returned unexpected value: got true want false")
	}

	if !IsExpoPushToken(token) {
		t.Errorf("IsExpoPushToken returned unexpected value: got false want true")
	}
}

func TestChunkPushNotifications(t *testing.T) {
	messages := []*PushMessage{
		{To: "token"},
		{To: "token"},
		{To: "token"},
	}

	chunks := ChunkPushNotifications(messages)
	if len(chunks) > 1 {
		t.Errorf("ChunkPushNotifications returned unexpected chunks: chunks length got %v want 1", len(chunks))
	}

	ChunkLimit = 2
	chunks = ChunkPushNotifications(messages)
	if len(chunks) != 2 {
		t.Errorf("ChunkPushNotifications returned unexpected chunks: chunks length got %v want 2", len(chunks))
	}

	ChunkLimit = 1
	chunks = ChunkPushNotifications(messages)
	if len(chunks) != 3 {
		t.Errorf("ChunkPushNotifications returned unexpected chunks: chunks length got %v want 3", len(chunks))
	}
}

func TestSendPushNotification(t *testing.T) {
	httpmock.Activate()
	defer httpmock.DeactivateAndReset()

	message := PushMessage{
		To:    token,
		Title: "Notification title",
		Body:  "Notification content"}

	httpmock.RegisterResponder("POST", baseAPIURL+"/push/send",
		httpmock.NewStringResponder(200, status200))

	api, _ := message.Send()
	if api.Data[0].Status != "ok" {
		t.Errorf("SendPushNotification returned unexpected response: status got %s want ok", api.Data[0].Status)
	}
}

func TestSendPushNotifications(t *testing.T) {
	httpmock.Activate()
	defer httpmock.DeactivateAndReset()

	m1 := PushMessage{
		To:    token,
		Title: "Notification title",
		Body:  "Notification content"}

	m2 := PushMessage{
		To:    token,
		Title: "Notification title",
		Body:  "Notification content"}

	httpmock.RegisterResponder("POST", baseAPIURL+"/push/send",
		httpmock.NewStringResponder(200, status200))

	api, _ := SendPushNotifications([]*PushMessage{&m1, &m2})
	if api.Data[0].Status != "ok" {
		t.Errorf("SendPushNotifications returned unexpected response: status got %s want ok", api.Data[0].Status)
	}
}

func TestBodyGzip(t *testing.T) {
	httpmock.Activate()
	defer httpmock.DeactivateAndReset()

	message := PushMessage{
		To:    token,
		Title: "Notification title",
		Body:  "Notification content"}

	MaxBodySizeWithoutGzip = 1

	httpmock.RegisterResponder("POST", baseAPIURL+"/push/send",
		func(req *http.Request) (*http.Response, error) {
			if req.ContentLength != 110 {
				t.Errorf("SendPushNotification send unexpected message: ContentLength got %v want 110", req.ContentLength)
			}

			resp, _ := httpmock.NewJsonResponse(200, status200)
			return resp, nil
		})

	SendPushNotification(&message)
}
<|endoftext|>


<commit_before>// Extract data from reflect.Value into a interface
// It's NOT reflect.Value.Interface(). The method only valid on struct but
// a basic type like int64
// Limit by refelct.Value, only convert to int64, float64, string,
// and collections or struct
package pgears

import (
	"fmt"
	"reflect"
	"encoding/json"
)

func SelectExtractor(val reflect.Value) func(reflect.Value)interface{} {
	var v = val
	var typ = v.Type()
	if typ.Kind() == reflect.Ptr {
		var v = reflect.Indirect(v)
		typ = v.Type()
	}
	var ret func(reflect.Value)interface{}
	switch typ.Kind() {
	case reflect.Bool:
		ret = ExtractBool
	case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8:
		ret = ExtractInt
	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16, reflect.Uint8:
		ret = ExtractUint
	case reflect.Float64, reflect.Float32:
		ret = ExtractFloat
	case reflect.String:
		ret = ExtractString
	case reflect.Struct, reflect.Map, reflect.Slice, reflect.Array, reflect.Chan, reflect.Func, reflect.Interface:
		ret = ExtractObject
	default:
		var message = fmt.Sprintf("I don't know how to extract a %v", v)
		panic(message)
	}
	return ret
}

func ExtractField(val reflect.Value, field reflect.StructField) interface{}{
	var itf interface{} = Extract(val)
	if field.Tag.Get("jsonto") != "" {
		itf, _ = json.Marshal(itf)
	}
	return itf
}

func Extract(val reflect.Value) interface{} {
	var v = val
	var typ = v.Type()
	if typ.Kind() == reflect.Ptr {
		if val.IsNil() {
			return nil
		}
		var v = reflect.Indirect(v)
		typ = v.Type()
	}
	var ret interface{}
	switch typ.Kind() {
	case reflect.Bool:
		ret = val.Bool()
	case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8:
		ret = val.Int()
	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16, reflect.Uint8:
		ret = val.Uint()
	case reflect.Float64, reflect.Float32:
		ret = val.Float()
	case reflect.String:
		ret = val.String()
	case reflect.Complex64, reflect.Complex128:
		ret = val.Complex()
	case reflect.Struct, reflect.Map, reflect.Slice, reflect.Array, reflect.Chan, reflect.Func, reflect.Interface:
		ret = val.Interface()
	default:
		var message = fmt.Sprintf("I don't know how to extract a %v", v)
		panic(message)
	}
	return ret
}

func ExtractBool(val reflect.Value) (ret interface{}) {
	ret = val.Bool()
	return ret
}

func ExtractInt(val reflect.Value) (ret interface{}) {
	ret = val.Int()
	return ret
}

func ExtractUint(val reflect.Value) (ret interface{}) {
	ret = val.Uint()
	return ret
}

func ExtractFloat(val reflect.Value) (ret interface{}) {
	ret = val.Float()
	return ret
}

func ExtractString(val reflect.Value) (ret interface{}) {
	ret = val.String()
	return ret
}
// all object can be box to a interface{}
func ExtractObject(val reflect.Value) (ret interface{}) {
	ret = val.Interface()
	return ret
}

func ExtractJsonMap(val reflect.Value) (ret interface{}) {
	ret, _ = json.Marshal(val.Interface())
	return ret
}
<commit_msg>add panic in in extract field function<commit_after>// Extract data from reflect.Value into a interface
// It's NOT reflect.Value.Interface(). The method only valid on struct but
// a basic type like int64
// Limit by refelct.Value, only convert to int64, float64, string,
// and collections or struct
package pgears

import (
	"fmt"
	"reflect"
	"encoding/json"
)

func SelectExtractor(val reflect.Value) func(reflect.Value)interface{} {
	var v = val
	var typ = v.Type()
	if typ.Kind() == reflect.Ptr {
		var v = reflect.Indirect(v)
		typ = v.Type()
	}
	var ret func(reflect.Value)interface{}
	switch typ.Kind() {
	case reflect.Bool:
		ret = ExtractBool
	case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8:
		ret = ExtractInt
	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16, reflect.Uint8:
		ret = ExtractUint
	case reflect.Float64, reflect.Float32:
		ret = ExtractFloat
	case reflect.String:
		ret = ExtractString
	case reflect.Struct, reflect.Map, reflect.Slice, reflect.Array, reflect.Chan, reflect.Func, reflect.Interface:
		ret = ExtractObject
	default:
		var message = fmt.Sprintf("I don't know how to extract a %v", v)
		panic(message)
	}
	return ret
}

func ExtractField(val reflect.Value, field reflect.StructField) interface{}{
	var itf interface{} = Extract(val)
	if field.Tag.Get("jsonto") != "" {
		itf, err := json.Marshal(itf)
		if err != nil {
			panic(err)
		}
	}
	return itf
}

func Extract(val reflect.Value) interface{} {
	var v = val
	var typ = v.Type()
	if typ.Kind() == reflect.Ptr {
		if val.IsNil() {
			return nil
		}
		var v = reflect.Indirect(v)
		typ = v.Type()
	}
	var ret interface{}
	switch typ.Kind() {
	case reflect.Bool:
		ret = val.Bool()
	case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8:
		ret = val.Int()
	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16, reflect.Uint8:
		ret = val.Uint()
	case reflect.Float64, reflect.Float32:
		ret = val.Float()
	case reflect.String:
		ret = val.String()
	case reflect.Complex64, reflect.Complex128:
		ret = val.Complex()
	case reflect.Struct, reflect.Map, reflect.Slice, reflect.Array, reflect.Chan, reflect.Func, reflect.Interface:
		ret = val.Interface()
	default:
		var message = fmt.Sprintf("I don't know how to extract a %v", v)
		panic(message)
	}
	return ret
}

func ExtractBool(val reflect.Value) (ret interface{}) {
	ret = val.Bool()
	return ret
}

func ExtractInt(val reflect.Value) (ret interface{}) {
	ret = val.Int()
	return ret
}

func ExtractUint(val reflect.Value) (ret interface{}) {
	ret = val.Uint()
	return ret
}

func ExtractFloat(val reflect.Value) (ret interface{}) {
	ret = val.Float()
	return ret
}

func ExtractString(val reflect.Value) (ret interface{}) {
	ret = val.String()
	return ret
}
// all object can be box to a interface{}
func ExtractObject(val reflect.Value) (ret interface{}) {
	ret = val.Interface()
	return ret
}

func ExtractJsonMap(val reflect.Value) (ret interface{}) {
	ret, _ = json.Marshal(val.Interface())
	return ret
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	rw "github.com/whyrusleeping/gx-go/rewrite"
	gx "github.com/whyrusleeping/gx/gxutil"
	. "github.com/whyrusleeping/stump"
)

func doUpdate(dir, oldimp, newimp string) error {
	rwf := func(in string) string {
		if in == oldimp {
			return newimp
		}

		if strings.HasPrefix(in, oldimp+"/") {
			return strings.Replace(in, oldimp, newimp, 1)
		}

		return in
	}

	filter := func(in string) bool {
		return strings.HasSuffix(in, ".go") && !strings.HasPrefix(in, "vendor")
	}

	return rw.RewriteImports(dir, rwf, filter)
}

func pathIsNotStdlib(path string) bool {
	first := strings.Split(path, "/")[0]

	if len(strings.Split(first, ".")) > 1 {
		return true
	}
	return false
}

type Importer struct {
	pkgs    map[string]*gx.Dependency
	gopath  string
	pm      *gx.PM
	rewrite bool
	yesall  bool
	preMap  map[string]string

	bctx build.Context
}

func NewImporter(rw bool, gopath string, premap map[string]string) (*Importer, error) {
	cfg, err := gx.LoadConfig()
	if err != nil {
		return nil, err
	}

	pm, err := gx.NewPM(cfg)
	if err != nil {
		return nil, err
	}

	if premap == nil {
		premap = make(map[string]string)
	}

	bctx := build.Default
	bctx.GOPATH = gopath

	return &Importer{
		pkgs:    make(map[string]*gx.Dependency),
		gopath:  gopath,
		pm:      pm,
		rewrite: rw,
		preMap:  premap,
		bctx:    bctx,
	}, nil
}

func getGoPath() (string, error) {
	gopath := os.Getenv("GOPATH")
	if gopath == "" {
		return "", errors.New("gopath not set")
	}
	return gopath, nil
}

// this function is an attempt to keep subdirectories of a package as part of
// the same logical gx package. It has a special case for golang.org/x/ packages
func getBaseDVCS(path string) string {
	parts := strings.Split(path, "/")
	depth := 3
	/*
		if parts[0] == "golang.org" && parts[1] == "x" {
			depth = 4
		}
	*/

	if len(parts) > depth {
		return strings.Join(parts[:3], "/")
	}
	return path
}

func (i *Importer) GxPublishGoPackage(imppath string) (*gx.Dependency, error) {
	imppath = getBaseDVCS(imppath)
	if d, ok := i.pkgs[imppath]; ok {
		return d, nil
	}

	if hash, ok := i.preMap[imppath]; ok {
		pkg, err := i.pm.GetPackage(hash)
		if err != nil {
			return nil, err
		}

		dep := &gx.Dependency{
			Hash:    hash,
			Name:    pkg.Name,
			Version: pkg.Version,
		}
		i.pkgs[imppath] = dep
		return dep, nil
	}

	// make sure its local
	err := i.GoGet(imppath)
	if err != nil {
		if !strings.Contains(err.Error(), "no buildable Go source files") {
			return nil, err
		}
	}

	pkgpath := path.Join(i.gopath, "src", imppath)
	pkgFilePath := path.Join(pkgpath, gx.PkgFileName)
	pkg, err := LoadPackageFile(pkgFilePath)
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}

		// init as gx package
		parts := strings.Split(imppath, "/")
		pkgname := parts[len(parts)-1]
		if !i.yesall {
			p := fmt.Sprintf("enter name for import '%s'", imppath)
			nname, err := prompt(p, pkgname)
			if err != nil {
				return nil, err
			}

			pkgname = nname
		}

		err = i.pm.InitPkg(pkgpath, pkgname, "go", nil)
		if err != nil {
			return nil, err
		}

		pkg, err = LoadPackageFile(pkgFilePath)
		if err != nil {
			return nil, err
		}
	}

	// wipe out existing dependencies
	pkg.Dependencies = nil

	// recurse!
	depsToVendor, err := i.depsToVendorForPackage(imppath)
	if err != nil {
		return nil, err
	}

	for n, child := range depsToVendor {
		Log("- processing dep %s for %s [%d / %d]", child, imppath, n+1, len(depsToVendor))
		if strings.HasPrefix(child, imppath) {
			continue
		}
		childdep, err := i.GxPublishGoPackage(child)
		if err != nil {
			return nil, err
		}

		pkg.Dependencies = append(pkg.Dependencies, childdep)
	}

	err = gx.SavePackageFile(pkg, pkgFilePath)
	if err != nil {
		return nil, err
	}

	fullpkgpath, err := filepath.Abs(pkgpath)
	if err != nil {
		return nil, err
	}

	err = i.rewriteImports(fullpkgpath)
	if err != nil {
		return nil, err
	}

	err = writeGxIgnore(pkgpath, []string{"Godeps/*"})
	if err != nil {
		return nil, err
	}

	hash, err := i.pm.PublishPackage(pkgpath, &pkg.PackageBase)
	if err != nil {
		return nil, err
	}

	Log("published %s as %s", imppath, hash)

	dep := &gx.Dependency{
		Hash:    hash,
		Name:    pkg.Name,
		Version: pkg.Version,
	}
	i.pkgs[imppath] = dep
	return dep, nil
}

func (i *Importer) depsToVendorForPackage(path string) ([]string, error) {
	rdeps := make(map[string]struct{})

	gopkg, err := i.bctx.Import(path, "", 0)
	if err != nil {
		_, ok := err.(*build.NoGoError)
		if !ok {
			return nil, err
		}
		// if theres no go code here, there still might be some in lower directories
	} else {
		imps := append(gopkg.Imports, gopkg.TestImports...)
		// if the package existed and has go code in it
		gdeps := getBaseDVCS(path) + "/Godeps/_workspace/src/"
		for _, child := range imps {
			if strings.HasPrefix(child, gdeps) {
				child = child[len(gdeps):]
			}

			child = getBaseDVCS(child)
			if pathIsNotStdlib(child) && !strings.HasPrefix(child, path) {
				rdeps[child] = struct{}{}
			}
		}
	}

	dirents, err := ioutil.ReadDir(filepath.Join(i.gopath, "src", path))
	if err != nil {
		return nil, err
	}

	for _, e := range dirents {
		if !e.IsDir() || skipDir(e.Name()) {
			continue
		}

		out, err := i.depsToVendorForPackage(filepath.Join(path, e.Name()))
		if err != nil {
			return nil, err
		}

		for _, o := range out {
			rdeps[o] = struct{}{}
		}
	}

	var depsToVendor []string
	for d, _ := range rdeps {
		depsToVendor = append(depsToVendor, d)
	}

	return depsToVendor, nil
}

func skipDir(name string) bool {
	switch name {
	case "Godeps", "vendor", ".git":
		return true
	default:
		return false
	}
}

func (i *Importer) rewriteImports(pkgpath string) error {

	filter := func(p string) bool {
		return !strings.HasPrefix(p, "vendor") &&
			!strings.HasPrefix(p, ".git") &&
			strings.HasSuffix(p, ".go") &&
			!strings.HasPrefix(p, "Godeps")
	}

	base := pkgpath[len(i.gopath)+5:]
	gdepath := base + "/Godeps/_workspace/src/"
	rwf := func(in string) string {
		if strings.HasPrefix(in, gdepath) {
			in = in[len(gdepath):]
		}

		if !i.rewrite {
			// if rewrite not specified, just fixup godeps paths
			return in
		}

		dep, ok := i.pkgs[in]
		if ok {
			return "gx/" + dep.Hash + "/" + dep.Name
		}

		parts := strings.Split(in, "/")
		if len(parts) > 3 {
			obase := strings.Join(parts[:3], "/")
			dep, bok := i.pkgs[obase]
			if !bok {
				return in
			}

			return strings.Replace(in, obase, "gx/"+dep.Hash+"/"+dep.Name, 1)
		}

		return in
	}

	return rw.RewriteImports(pkgpath, rwf, filter)
}

// TODO: take an option to grab packages from local GOPATH
func (imp *Importer) GoGet(path string) error {
	cmd := exec.Command("go", "get", path)
	env := os.Environ()
	for i, e := range env {
		if strings.HasPrefix(e, "GOPATH=") {
			env[i] = "GOPATH=" + imp.gopath
		}
	}
	cmd.Env = env
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("go get failed: %s - %s", string(out), err)
	}
	return nil
}

func writeGxIgnore(dir string, ignore []string) error {
	return ioutil.WriteFile(filepath.Join(dir, ".gxignore"), []byte(strings.Join(ignore, "\n")), 0644)
}
<commit_msg>dont use GetPackage anymore<commit_after>package main

import (
	"errors"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	rw "github.com/whyrusleeping/gx-go/rewrite"
	gx "github.com/whyrusleeping/gx/gxutil"
	. "github.com/whyrusleeping/stump"
)

func doUpdate(dir, oldimp, newimp string) error {
	rwf := func(in string) string {
		if in == oldimp {
			return newimp
		}

		if strings.HasPrefix(in, oldimp+"/") {
			return strings.Replace(in, oldimp, newimp, 1)
		}

		return in
	}

	filter := func(in string) bool {
		return strings.HasSuffix(in, ".go") && !strings.HasPrefix(in, "vendor")
	}

	return rw.RewriteImports(dir, rwf, filter)
}

func pathIsNotStdlib(path string) bool {
	first := strings.Split(path, "/")[0]

	if len(strings.Split(first, ".")) > 1 {
		return true
	}
	return false
}

type Importer struct {
	pkgs    map[string]*gx.Dependency
	gopath  string
	pm      *gx.PM
	rewrite bool
	yesall  bool
	preMap  map[string]string

	bctx build.Context
}

func NewImporter(rw bool, gopath string, premap map[string]string) (*Importer, error) {
	cfg, err := gx.LoadConfig()
	if err != nil {
		return nil, err
	}

	pm, err := gx.NewPM(cfg)
	if err != nil {
		return nil, err
	}

	if premap == nil {
		premap = make(map[string]string)
	}

	bctx := build.Default
	bctx.GOPATH = gopath

	return &Importer{
		pkgs:    make(map[string]*gx.Dependency),
		gopath:  gopath,
		pm:      pm,
		rewrite: rw,
		preMap:  premap,
		bctx:    bctx,
	}, nil
}

func getGoPath() (string, error) {
	gopath := os.Getenv("GOPATH")
	if gopath == "" {
		return "", errors.New("gopath not set")
	}
	return gopath, nil
}

// this function is an attempt to keep subdirectories of a package as part of
// the same logical gx package. It has a special case for golang.org/x/ packages
func getBaseDVCS(path string) string {
	parts := strings.Split(path, "/")
	depth := 3
	/*
		if parts[0] == "golang.org" && parts[1] == "x" {
			depth = 4
		}
	*/

	if len(parts) > depth {
		return strings.Join(parts[:3], "/")
	}
	return path
}

func (i *Importer) GxPublishGoPackage(imppath string) (*gx.Dependency, error) {
	imppath = getBaseDVCS(imppath)
	if d, ok := i.pkgs[imppath]; ok {
		return d, nil
	}

	if hash, ok := i.preMap[imppath]; ok {
		pkg, err := i.pm.GetPackageTo(hash, filepath.Join(vendorDir, hash))
		if err != nil {
			return nil, err
		}

		dep := &gx.Dependency{
			Hash:    hash,
			Name:    pkg.Name,
			Version: pkg.Version,
		}
		i.pkgs[imppath] = dep
		return dep, nil
	}

	// make sure its local
	err := i.GoGet(imppath)
	if err != nil {
		if !strings.Contains(err.Error(), "no buildable Go source files") {
			return nil, err
		}
	}

	pkgpath := path.Join(i.gopath, "src", imppath)
	pkgFilePath := path.Join(pkgpath, gx.PkgFileName)
	pkg, err := LoadPackageFile(pkgFilePath)
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}

		// init as gx package
		parts := strings.Split(imppath, "/")
		pkgname := parts[len(parts)-1]
		if !i.yesall {
			p := fmt.Sprintf("enter name for import '%s'", imppath)
			nname, err := prompt(p, pkgname)
			if err != nil {
				return nil, err
			}

			pkgname = nname
		}

		err = i.pm.InitPkg(pkgpath, pkgname, "go", nil)
		if err != nil {
			return nil, err
		}

		pkg, err = LoadPackageFile(pkgFilePath)
		if err != nil {
			return nil, err
		}
	}

	// wipe out existing dependencies
	pkg.Dependencies = nil

	// recurse!
	depsToVendor, err := i.depsToVendorForPackage(imppath)
	if err != nil {
		return nil, err
	}

	for n, child := range depsToVendor {
		Log("- processing dep %s for %s [%d / %d]", child, imppath, n+1, len(depsToVendor))
		if strings.HasPrefix(child, imppath) {
			continue
		}
		childdep, err := i.GxPublishGoPackage(child)
		if err != nil {
			return nil, err
		}

		pkg.Dependencies = append(pkg.Dependencies, childdep)
	}

	err = gx.SavePackageFile(pkg, pkgFilePath)
	if err != nil {
		return nil, err
	}

	fullpkgpath, err := filepath.Abs(pkgpath)
	if err != nil {
		return nil, err
	}

	err = i.rewriteImports(fullpkgpath)
	if err != nil {
		return nil, err
	}

	err = writeGxIgnore(pkgpath, []string{"Godeps/*"})
	if err != nil {
		return nil, err
	}

	hash, err := i.pm.PublishPackage(pkgpath, &pkg.PackageBase)
	if err != nil {
		return nil, err
	}

	Log("published %s as %s", imppath, hash)

	dep := &gx.Dependency{
		Hash:    hash,
		Name:    pkg.Name,
		Version: pkg.Version,
	}
	i.pkgs[imppath] = dep
	return dep, nil
}

func (i *Importer) depsToVendorForPackage(path string) ([]string, error) {
	rdeps := make(map[string]struct{})

	gopkg, err := i.bctx.Import(path, "", 0)
	if err != nil {
		_, ok := err.(*build.NoGoError)
		if !ok {
			return nil, err
		}
		// if theres no go code here, there still might be some in lower directories
	} else {
		imps := append(gopkg.Imports, gopkg.TestImports...)
		// if the package existed and has go code in it
		gdeps := getBaseDVCS(path) + "/Godeps/_workspace/src/"
		for _, child := range imps {
			if strings.HasPrefix(child, gdeps) {
				child = child[len(gdeps):]
			}

			child = getBaseDVCS(child)
			if pathIsNotStdlib(child) && !strings.HasPrefix(child, path) {
				rdeps[child] = struct{}{}
			}
		}
	}

	dirents, err := ioutil.ReadDir(filepath.Join(i.gopath, "src", path))
	if err != nil {
		return nil, err
	}

	for _, e := range dirents {
		if !e.IsDir() || skipDir(e.Name()) {
			continue
		}

		out, err := i.depsToVendorForPackage(filepath.Join(path, e.Name()))
		if err != nil {
			return nil, err
		}

		for _, o := range out {
			rdeps[o] = struct{}{}
		}
	}

	var depsToVendor []string
	for d, _ := range rdeps {
		depsToVendor = append(depsToVendor, d)
	}

	return depsToVendor, nil
}

func skipDir(name string) bool {
	switch name {
	case "Godeps", "vendor", ".git":
		return true
	default:
		return false
	}
}

func (i *Importer) rewriteImports(pkgpath string) error {

	filter := func(p string) bool {
		return !strings.HasPrefix(p, "vendor") &&
			!strings.HasPrefix(p, ".git") &&
			strings.HasSuffix(p, ".go") &&
			!strings.HasPrefix(p, "Godeps")
	}

	base := pkgpath[len(i.gopath)+5:]
	gdepath := base + "/Godeps/_workspace/src/"
	rwf := func(in string) string {
		if strings.HasPrefix(in, gdepath) {
			in = in[len(gdepath):]
		}

		if !i.rewrite {
			// if rewrite not specified, just fixup godeps paths
			return in
		}

		dep, ok := i.pkgs[in]
		if ok {
			return "gx/" + dep.Hash + "/" + dep.Name
		}

		parts := strings.Split(in, "/")
		if len(parts) > 3 {
			obase := strings.Join(parts[:3], "/")
			dep, bok := i.pkgs[obase]
			if !bok {
				return in
			}

			return strings.Replace(in, obase, "gx/"+dep.Hash+"/"+dep.Name, 1)
		}

		return in
	}

	return rw.RewriteImports(pkgpath, rwf, filter)
}

// TODO: take an option to grab packages from local GOPATH
func (imp *Importer) GoGet(path string) error {
	cmd := exec.Command("go", "get", path)
	env := os.Environ()
	for i, e := range env {
		if strings.HasPrefix(e, "GOPATH=") {
			env[i] = "GOPATH=" + imp.gopath
		}
	}
	cmd.Env = env
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("go get failed: %s - %s", string(out), err)
	}
	return nil
}

func writeGxIgnore(dir string, ignore []string) error {
	return ioutil.WriteFile(filepath.Join(dir, ".gxignore"), []byte(strings.Join(ignore, "\n")), 0644)
}
<|endoftext|>


<commit_before>package yunti_core

import "net"
import "golang.org/x/net/websocket"
import "net/url"

type leaptransp_stage int

const (
    leaptransp_stage_NULL leaptransp_stage = iota
    leaptransp_stage_initialized
    leaptransp_stage_connected
    leaptransp_stage_serverauthdone
    leaptransp_stage_clientauthdone
    leaptransp_stage_syncconnest
    leaptransp_stage_asyncconnest
    leaptransp_stage_failed
)

type leaptransp struct{
  Leaptype string
  Stage leaptransp_stage
  Sync bool
  conn interface
}

func leaptransp_websocket_Connect_genorigin(wsaddr string)string{
  oriu,_:=url.Parse(wsaddr)
  switch strings.ToUpper(oriu.Scheme){
  case "WSS":
  oriu.Scheme="https"
  case "WS":
  oriu.Scheme="http"
}
origin=oriu.String()
return origin
}

func (lt *leaptransp)Leaptransp_websocket_Connect(wsaddr string){

origin:=leaptransp_websocket_Connect_genorigin(wsaddr)

wsc,err:=websocket.Dial(wsaddr, "", origin)

if err != nil {
  log.Fatal(err)
}

}
<commit_msg>finialize websocket connect<commit_after>package yunti_core

import "net"
import "golang.org/x/net/websocket"
import "net/url"

type leaptransp_stage int

const (
    leaptransp_stage_NULL leaptransp_stage = iota
    leaptransp_stage_initialized
    leaptransp_stage_connected
    leaptransp_stage_serverauthdone
    leaptransp_stage_clientauthdone
    leaptransp_stage_syncconnest
    leaptransp_stage_asyncconnest
    leaptransp_stage_failed
)

type leaptransp struct{
  Leaptype string
  Stage leaptransp_stage
  Sync bool
  conn interface
}

func leaptransp_websocket_Connect_genorigin(wsaddr string)string{
  oriu,_:=url.Parse(wsaddr)
  switch strings.ToUpper(oriu.Scheme){
  case "WSS":
  oriu.Scheme="https"
  case "WS":
  oriu.Scheme="http"
}
origin=oriu.String()
return origin
}

func (lt *leaptransp)Leaptransp_websocket_Connect(wsaddr string)int{

origin:=leaptransp_websocket_Connect_genorigin(wsaddr)

wsc,err:=websocket.Dial(wsaddr, "", origin)

if err != nil {
  return -1
}

lt.conn=&wsc

}
<|endoftext|>


<commit_before><commit_msg>use utf8mb4<commit_after><|endoftext|>


<commit_before><commit_msg>add minimal golang example<commit_after><|endoftext|>


<commit_before><commit_msg>Move endpoint to accept only JSON<commit_after><|endoftext|>


<commit_before>package metadata

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/rancher/external-dns/config"
	"github.com/rancher/external-dns/utils"
	"github.com/rancher/go-rancher-metadata/metadata"
	"time"
)

const (
	metadataUrl = "http://rancher-metadata/2015-12-19"
)

type MetadataClient struct {
	MetadataClient  *metadata.Client
	EnvironmentName string
	EnvironmentUUID string
}

func getEnvironment(m *metadata.Client) (string, string, error) {
	timeout := 30 * time.Second
	var err error
	var stack metadata.Stack
	for i := 1 * time.Second; i < timeout; i *= time.Duration(2) {
		stack, err = m.GetSelfStack()
		if err != nil {
			logrus.Errorf("Error reading stack info: %v...will retry", err)
			time.Sleep(i)
		} else {
			return stack.EnvironmentName, stack.EnvironmentUUID, nil
		}
	}
	return "", "", fmt.Errorf("Error reading stack info: %v", err)
}

func NewMetadataClient() (*MetadataClient, error) {
	m, err := metadata.NewClientAndWait(metadataUrl)
	if err != nil {
		logrus.Fatalf("Failed to configure rancher-metadata: %v", err)
	}

	envName, envUUID, err := getEnvironment(m)
	if err != nil {
		logrus.Fatalf("Error reading stack info: %v", err)
	}

	return &MetadataClient{
		MetadataClient:  m,
		EnvironmentName: envName,
		EnvironmentUUID: envUUID,
	}, nil
}

func (m *MetadataClient) GetVersion() (string, error) {
	return m.MetadataClient.GetVersion()
}

func (m *MetadataClient) GetMetadataDnsRecords() (map[string]utils.DnsRecord, error) {
	dnsEntries := make(map[string]utils.DnsRecord)
	err := m.getContainersDnsRecords(dnsEntries, "", "")
	if err != nil {
		return dnsEntries, err
	}
	return dnsEntries, nil
}

func (m *MetadataClient) getContainersDnsRecords(dnsEntries map[string]utils.DnsRecord, serviceName string, stackName string) error {
	containers, err := m.MetadataClient.GetContainers()
	if err != nil {
		return err
	}

	ourFqdns := make(map[string]struct{})
	for _, container := range containers {
		if len(container.ServiceName) == 0 || len(container.Ports) == 0 || !containerStateOK(container) {
			continue
		}

		if len(serviceName) != 0 {
			if serviceName != container.ServiceName {
				continue
			}
			if stackName != container.StackName {
				continue
			}
		}

		hostUUID := container.HostUUID
		if len(hostUUID) == 0 {
			logrus.Debugf("Container's %v host_uuid is empty", container.Name)
			continue
		}
		host, err := m.MetadataClient.GetHost(hostUUID)
		if err != nil {
			logrus.Infof("%v", err)
			continue
		}

		ip, ok := host.Labels["io.rancher.host.external_dns_ip"]

		if !ok || ip == "" {
			ip = host.AgentIP
		}

		fqdn := utils.FqdnFromTemplate(config.NameTemplate, container.ServiceName, container.StackName,
			m.EnvironmentName, config.RootDomainName)
		records := []string{ip}
		dnsEntry := utils.DnsRecord{fqdn, records, "A", config.TTL}

		addToDnsEntries(dnsEntry, dnsEntries)
		ourFqdns[fqdn] = struct{}{}
	}

	if len(ourFqdns) > 0 {
		fqdn := utils.StateFqdn(m.EnvironmentUUID, config.RootDomainName)
		stateRec := utils.StateRecord(fqdn, config.TTL, ourFqdns)
		addToDnsEntries(stateRec, dnsEntries)
	}

	return nil
}

func addToDnsEntries(dnsEntry utils.DnsRecord, dnsEntries map[string]utils.DnsRecord) {
	var records []string
	if _, ok := dnsEntries[dnsEntry.Fqdn]; !ok {
		records = dnsEntry.Records
	} else {
		records = dnsEntries[dnsEntry.Fqdn].Records
		records = append(records, dnsEntry.Records...)
	}
	dnsEntry = utils.DnsRecord{dnsEntry.Fqdn, records, dnsEntry.Type, dnsEntry.TTL}
	dnsEntries[dnsEntry.Fqdn] = dnsEntry
}

func containerStateOK(container metadata.Container) bool {
	switch container.State {
	case "running":
	default:
		return false
	}

	switch container.HealthState {
	case "healthy":
	case "updating-healthy":
	case "":
	default:
		return false
	}

	return true
}
<commit_msg>Don't create DNS records for system services (#50)<commit_after>package metadata

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/rancher/external-dns/config"
	"github.com/rancher/external-dns/utils"
	"github.com/rancher/go-rancher-metadata/metadata"
	"time"
)

const (
	metadataUrl = "http://rancher-metadata/2015-12-19"
)

type MetadataClient struct {
	MetadataClient  *metadata.Client
	EnvironmentName string
	EnvironmentUUID string
}

func getEnvironment(m *metadata.Client) (string, string, error) {
	timeout := 30 * time.Second
	var err error
	var stack metadata.Stack
	for i := 1 * time.Second; i < timeout; i *= time.Duration(2) {
		stack, err = m.GetSelfStack()
		if err != nil {
			logrus.Errorf("Error reading stack info: %v...will retry", err)
			time.Sleep(i)
		} else {
			return stack.EnvironmentName, stack.EnvironmentUUID, nil
		}
	}
	return "", "", fmt.Errorf("Error reading stack info: %v", err)
}

func NewMetadataClient() (*MetadataClient, error) {
	m, err := metadata.NewClientAndWait(metadataUrl)
	if err != nil {
		logrus.Fatalf("Failed to configure rancher-metadata: %v", err)
	}

	envName, envUUID, err := getEnvironment(m)
	if err != nil {
		logrus.Fatalf("Error reading stack info: %v", err)
	}

	return &MetadataClient{
		MetadataClient:  m,
		EnvironmentName: envName,
		EnvironmentUUID: envUUID,
	}, nil
}

func (m *MetadataClient) GetVersion() (string, error) {
	return m.MetadataClient.GetVersion()
}

func (m *MetadataClient) GetMetadataDnsRecords() (map[string]utils.DnsRecord, error) {
	dnsEntries := make(map[string]utils.DnsRecord)
	err := m.getContainersDnsRecords(dnsEntries)
	if err != nil {
		return dnsEntries, err
	}
	return dnsEntries, nil
}

func (m *MetadataClient) getContainersDnsRecords(dnsEntries map[string]utils.DnsRecord) error {
	services, err := m.MetadataClient.GetServices()
	if err != nil {
		return err
	}

	ourFqdns := make(map[string]struct{})
	hostMeta := make(map[string]metadata.Host)
	for _, service := range services {
		if service.Kind != "service" && service.Kind != "loadBalancerService" {
			continue
		}

		for _, container := range service.Containers {
			if len(container.Ports) == 0 || !containerStateOK(container) {
				continue
			}

			hostUUID := container.HostUUID
			if len(hostUUID) == 0 {
				logrus.Debugf("Container's %v host_uuid is empty", container.Name)
				continue
			}

			var host metadata.Host
			if _, ok := hostMeta[hostUUID]; ok {
				host = hostMeta[hostUUID]
			} else {
				host, err := m.MetadataClient.GetHost(hostUUID)
				if err != nil {
					logrus.Warnf("Failed to get host metadata: %v", err)
					continue
				}
				hostMeta[hostUUID] = host
			}

			ip, ok := host.Labels["io.rancher.host.external_dns_ip"]
			if !ok || ip == "" {
				ip = host.AgentIP
			}

			fqdn := utils.FqdnFromTemplate(config.NameTemplate, container.ServiceName, container.StackName,
				m.EnvironmentName, config.RootDomainName)
			records := []string{ip}
			dnsEntry := utils.DnsRecord{fqdn, records, "A", config.TTL}

			addToDnsEntries(dnsEntry, dnsEntries)
			ourFqdns[fqdn] = struct{}{}
		}
	}

	if len(ourFqdns) > 0 {
		fqdn := utils.StateFqdn(m.EnvironmentUUID, config.RootDomainName)
		stateRec := utils.StateRecord(fqdn, config.TTL, ourFqdns)
		addToDnsEntries(stateRec, dnsEntries)
	}

	return nil
}

func addToDnsEntries(dnsEntry utils.DnsRecord, dnsEntries map[string]utils.DnsRecord) {
	var records []string
	if _, ok := dnsEntries[dnsEntry.Fqdn]; !ok {
		records = dnsEntry.Records
	} else {
		records = dnsEntries[dnsEntry.Fqdn].Records
		records = append(records, dnsEntry.Records...)
	}
	dnsEntry = utils.DnsRecord{dnsEntry.Fqdn, records, dnsEntry.Type, dnsEntry.TTL}
	dnsEntries[dnsEntry.Fqdn] = dnsEntry
}

func containerStateOK(container metadata.Container) bool {
	switch container.State {
	case "running":
	default:
		return false
	}

	switch container.HealthState {
	case "healthy":
	case "updating-healthy":
	case "":
	default:
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package levant

import (
	"fmt"
	"strings"
	"time"

	nomad "github.com/hashicorp/nomad/api"
	nomadStructs "github.com/hashicorp/nomad/nomad/structs"
	"github.com/jrasell/levant/logging"
)

type nomadClient struct {
	nomad *nomad.Client
}

// NomadClient is an interface to the Nomad API and deployment functions.
type NomadClient interface {
	// Deploy triggers a register of the job resulting in a Nomad deployment which
	// is monitored to determine the eventual state.
	Deploy(*nomad.Job, int, bool) bool
}

// NewNomadClient is used to create a new client to interact with Nomad.
func NewNomadClient(addr string) (NomadClient, error) {
	config := nomad.DefaultConfig()

	if addr != "" {
		config.Address = addr
	}

	c, err := nomad.NewClient(config)
	if err != nil {
		return nil, err
	}

	return &nomadClient{nomad: c}, nil
}

// Deploy triggers a register of the job resulting in a Nomad deployment which
// is monitored to determine the eventual state.
func (c *nomadClient) Deploy(job *nomad.Job, autoPromote int, forceCount bool) (success bool) {

	// Validate the job to check it is syntactically correct.
	if _, _, err := c.nomad.Jobs().Validate(job, nil); err != nil {
		logging.Error("levant/deploy: job validation failed: %v", err)
		return
	}

	// If job.Type isn't set we can't continue
	if job.Type == nil {
		logging.Error("levant/deploy: Nomad job `type` is not set; should be set to `%s`, `%s` or `%s`",
			nomadStructs.JobTypeBatch, nomadStructs.JobTypeSystem, nomadStructs.JobTypeService)
		return
	}

	if !forceCount {
		logging.Debug("levant/deploy: running dynamic job count updater for job %s", *job.Name)
		if err := c.dynamicGroupCountUpdater(job); err != nil {
			return
		}
	}

	// Check that the job has at least 1 TaskGroup with count > 0 (GH-16) if the
	// job is not a System job. Systems jobs do not define counts so cannot be
	// checked.
	if *job.Type != nomadStructs.JobTypeSystem {
		tgCount := 0
		for _, group := range job.TaskGroups {
			tgCount += *group.Count
		}
		if tgCount == 0 {
			logging.Error("levant/deploy: all TaskGroups have a count of 0, nothing to do")
			return
		}
	}

	logging.Info("levant/deploy: triggering a deployment of job %s", *job.Name)

	eval, _, err := c.nomad.Jobs().Register(job, nil)
	if err != nil {
		logging.Error("levant/deploy: unable to register job %s with Nomad: %v", *job.Name, err)
		return
	}

	// GH-50: batch job types do not return an evaluation upon registration.
	if eval.EvalID == "" && *job.Type == nomadStructs.JobTypeBatch {
		return c.checkBatchJob(job.Name)
	}

	// Trigger the evaluationInspector to identify any potential errors in the
	// Nomad evaluation run. As far as I can tell from testing; a single alloc
	// failure in an evaluation means no allocs will be placed so we exit here.
	err = c.evaluationInspector(&eval.EvalID)
	if err != nil {
		logging.Error("levant/deploy: %v", err)
		return
	}

	switch *job.Type {
	case nomadStructs.JobTypeService:
		logging.Info("levant/deploy: beginning deployment watcher for job %s", *job.Name)

		// Get the deploymentID from the evaluationID so that we can watch the
		// deployment for end status.
		depID, err := c.getDeploymentID(eval.EvalID)
		if err != nil {
			logging.Error("levant/deploy: unable to get info of evaluation %s: %v", eval.EvalID, err)
			return
		}

		// Get the success of the deployment.
		success = c.deploymentWatcher(depID, autoPromote)

		// If the deployment has not been successful; check whether the job is
		// configured to auto-revert so that this can be tracked.
		if !success {
			dep, _, err := c.nomad.Deployments().Info(depID, nil)
			if err != nil {
				logging.Error("levant/deploy: unable to query deployment %s for auto-revert check: %v",
					dep.ID, err)
				break
			}
			c.checkAutoRevert(dep)
		}

	default:
		logging.Debug("levant/deploy: job type %s does not support Nomad deployment model", *job.Type)
		success = true
	}

	return
}

func (c *nomadClient) evaluationInspector(evalID *string) error {

	for {
		evalInfo, _, err := c.nomad.Evaluations().Info(*evalID, nil)
		if err != nil {
			return err
		}

		switch evalInfo.Status {
		case nomadStructs.EvalStatusComplete, nomadStructs.EvalStatusFailed, nomadStructs.EvalStatusCancelled:
			if len(evalInfo.FailedTGAllocs) == 0 {
				logging.Info("levant/deploy: evaluation %s finished successfully", *evalID)
				return nil
			}

			for group, metrics := range evalInfo.FailedTGAllocs {

				// Check if any nodes have been exhausted of resources and therfore are
				// unable to place allocs.
				if metrics.NodesExhausted > 0 {
					var exhausted, dimension []string
					for e := range metrics.ClassExhausted {
						exhausted = append(exhausted, e)
					}
					for d := range metrics.DimensionExhausted {
						dimension = append(dimension, d)
					}
					logging.Error("levant/deploy: task group %s failed to place allocs, failed on %v and exhausted %v",
						group, exhausted, dimension)
				}

				// Check if any node classes were filtered causing alloc placement
				// failures.
				if len(metrics.ClassFiltered) > 0 {
					for f := range metrics.ClassFiltered {
						logging.Error("levant/deploy: task group %s failed to place %v allocs as class \"%s\" was filtered",
							group, len(metrics.ClassFiltered), f)
					}
				}

				// Check if any node constraints were filtered causing alloc placement
				// failures.
				if len(metrics.ConstraintFiltered) > 0 {
					for cf := range metrics.ConstraintFiltered {
						logging.Error("levant/deploy: task group %s failed to place %v allocs as constraint \"%s\" was filtered",
							group, len(metrics.ConstraintFiltered), cf)
					}
				}
			}

			return fmt.Errorf("evaluation %v finished with status %s but failed to place allocations",
				*evalID, evalInfo.Status)

		default:
			time.Sleep(1 * time.Second)
			continue
		}
	}
}

func (c *nomadClient) deploymentWatcher(depID string, autoPromote int) (success bool) {

	var canaryChan chan interface{}
	deploymentChan := make(chan interface{})

	t := time.Now()
	wt := time.Duration(5 * time.Second)

	// Setup the canaryChan and launch the autoPromote go routine if autoPromote
	// has been enabled.
	if autoPromote > 0 {
		canaryChan = make(chan interface{})
		go c.canaryAutoPromote(depID, autoPromote, canaryChan, deploymentChan)
	}

	q := &nomad.QueryOptions{WaitIndex: 1, AllowStale: true, WaitTime: wt}

	for {

		dep, meta, err := c.nomad.Deployments().Info(depID, q)
		logging.Debug("levant/deploy: deployment %v running for %.2fs", depID, time.Since(t).Seconds())

		// Listen for the deploymentChan closing which indicates Levant should exit
		// the deployment watcher.
		select {
		case <-deploymentChan:
			return false
		default:
			break
		}

		if err != nil {
			logging.Error("levant/deploy: unable to get info of deployment %s: %v", depID, err)
			return
		}

		if meta.LastIndex <= q.WaitIndex {
			continue
		}

		q.WaitIndex = meta.LastIndex

		cont, err := c.checkDeploymentStatus(dep, canaryChan)
		if err != nil {
			return false
		}

		if cont {
			continue
		} else {
			return true
		}
	}
}

func (c *nomadClient) checkDeploymentStatus(dep *nomad.Deployment, shutdownChan chan interface{}) (bool, error) {

	switch dep.Status {
	case nomadStructs.DeploymentStatusSuccessful:
		logging.Info("levant/deploy: deployment %v has completed successfully", dep.ID)
		return false, nil
	case nomadStructs.DeploymentStatusRunning:
		return true, nil
	default:
		if shutdownChan != nil {
			logging.Debug("levant/deploy: deployment %v meaning canary auto promote will shutdown", dep.Status)
			close(shutdownChan)
		}

		logging.Error("levant/deploy: deployment %v has status %s", dep.ID, dep.Status)

		// Launch the failure inspector.
		c.checkFailedDeployment(&dep.ID)

		return false, fmt.Errorf("deployment failed")
	}
}

// canaryAutoPromote handles Levant's canary-auto-promote functionality.
func (c *nomadClient) canaryAutoPromote(depID string, waitTime int, shutdownChan, deploymentChan chan interface{}) {

	// Setup the AutoPromote timer.
	autoPromote := time.After(time.Duration(waitTime) * time.Second)

	for {
		select {
		case <-autoPromote:
			logging.Info("levant/deploy: auto-promote period %vs has been reached for deployment %s",
				waitTime, depID)

			// Check the deployment is healthy before promoting.
			if healthy := c.checkCanaryDeploymentHealth(depID); !healthy {
				logging.Error("levant/deploy: the canary deployment %s has unhealthy allocations, unable to promote", depID)
				close(deploymentChan)
				return
			}

			logging.Info("levant/deploy: triggering auto promote of deployment %s", depID)

			// Promote the deployment.
			_, _, err := c.nomad.Deployments().PromoteAll(depID, nil)
			if err != nil {
				logging.Error("levant/deploy: unable to promote deployment %s: %v", depID, err)
				close(deploymentChan)
				return
			}

		case <-shutdownChan:
			logging.Info("levant/deploy: canary auto promote has been shutdown")
			return
		}
	}
}

// checkCanaryDeploymentHealth is used to check the health status of each
// task-group within a canary deployment.
func (c *nomadClient) checkCanaryDeploymentHealth(depID string) (healthy bool) {

	var unhealthy int

	dep, _, err := c.nomad.Deployments().Info(depID, &nomad.QueryOptions{AllowStale: true})
	if err != nil {
		logging.Error("levant/deploy: unable to query deployment %s for health: %v", depID, err)
		return
	}

	// Itertate each task in the deployment to determine is health status. If an
	// unhealthy task is found, incrament the unhealthy counter.
	for taskName, taskInfo := range dep.TaskGroups {
		if taskInfo.DesiredCanaries != taskInfo.HealthyAllocs {
			logging.Error("levant/deploy: task %s has unhealthy allocations in deployment %s", taskName, depID)
			unhealthy++
		}
	}

	// If zero unhealthy tasks were found, continue with the auto promotion.
	if unhealthy == 0 {
		logging.Debug("levant/deploy: deployment %s has 0 unhealthy allocations", depID)
		healthy = true
	}

	return
}

// getDeploymentID finds the Nomad deploymentID associated to a Nomad
// evaluationID. This is only needed as sometimes Nomad initially returns eval
// info with an empty deploymentID; and a retry is required in order to get the
// updated response from Nomad.
func (c *nomadClient) getDeploymentID(evalID string) (depID string, err error) {

	var evalInfo *nomad.Evaluation

	for {
		if evalInfo, _, err = c.nomad.Evaluations().Info(evalID, nil); err != nil {
			return
		}

		if evalInfo.DeploymentID == "" {
			logging.Debug("levant/deploy: Nomad returned an empty deployment for evaluation %v; retrying", evalID)
			time.Sleep(2 * time.Second)
			continue
		} else {
			break
		}
	}

	return evalInfo.DeploymentID, nil
}

// dynamicGroupCountUpdater takes the templated and rendered job and updates the
// group counts based on the currently deployed job; if its running.
func (c *nomadClient) dynamicGroupCountUpdater(job *nomad.Job) error {

	// Gather information about the current state, if any, of the job on the
	// Nomad cluster.
	rJob, _, err := c.nomad.Jobs().Info(*job.Name, &nomad.QueryOptions{})

	// This is a hack due to GH-1849; we check the error string for 404 which
	// indicates the job is not running, not that there was an error in the API
	// call.
	if err != nil && strings.Contains(err.Error(), "404") {
		logging.Info("levant/deploy: job %s not running, using template file group counts", *job.Name)
		return nil
	} else if err != nil {
		logging.Error("levant/deploy: unable to perform job evaluation: %v", err)
		return err
	}

	// Iterate the templated job and the Nomad returned job and update group count
	// based on matches.
	for _, rGroup := range rJob.TaskGroups {
		for _, group := range job.TaskGroups {
			if *rGroup.Name == *group.Name {
				logging.Info("levant/deploy: using dynamic count %v for job %s and group %s",
					*rGroup.Count, *job.Name, *group.Name)
				group.Count = rGroup.Count
			}
		}
	}
	return nil
}
<commit_msg>levant/deploy: skip health checks for task groups without canaries<commit_after>package levant

import (
	"fmt"
	"strings"
	"time"

	nomad "github.com/hashicorp/nomad/api"
	nomadStructs "github.com/hashicorp/nomad/nomad/structs"
	"github.com/jrasell/levant/logging"
)

type nomadClient struct {
	nomad *nomad.Client
}

// NomadClient is an interface to the Nomad API and deployment functions.
type NomadClient interface {
	// Deploy triggers a register of the job resulting in a Nomad deployment which
	// is monitored to determine the eventual state.
	Deploy(*nomad.Job, int, bool) bool
}

// NewNomadClient is used to create a new client to interact with Nomad.
func NewNomadClient(addr string) (NomadClient, error) {
	config := nomad.DefaultConfig()

	if addr != "" {
		config.Address = addr
	}

	c, err := nomad.NewClient(config)
	if err != nil {
		return nil, err
	}

	return &nomadClient{nomad: c}, nil
}

// Deploy triggers a register of the job resulting in a Nomad deployment which
// is monitored to determine the eventual state.
func (c *nomadClient) Deploy(job *nomad.Job, autoPromote int, forceCount bool) (success bool) {

	// Validate the job to check it is syntactically correct.
	if _, _, err := c.nomad.Jobs().Validate(job, nil); err != nil {
		logging.Error("levant/deploy: job validation failed: %v", err)
		return
	}

	// If job.Type isn't set we can't continue
	if job.Type == nil {
		logging.Error("levant/deploy: Nomad job `type` is not set; should be set to `%s`, `%s` or `%s`",
			nomadStructs.JobTypeBatch, nomadStructs.JobTypeSystem, nomadStructs.JobTypeService)
		return
	}

	if !forceCount {
		logging.Debug("levant/deploy: running dynamic job count updater for job %s", *job.Name)
		if err := c.dynamicGroupCountUpdater(job); err != nil {
			return
		}
	}

	// Check that the job has at least 1 TaskGroup with count > 0 (GH-16) if the
	// job is not a System job. Systems jobs do not define counts so cannot be
	// checked.
	if *job.Type != nomadStructs.JobTypeSystem {
		tgCount := 0
		for _, group := range job.TaskGroups {
			tgCount += *group.Count
		}
		if tgCount == 0 {
			logging.Error("levant/deploy: all TaskGroups have a count of 0, nothing to do")
			return
		}
	}

	logging.Info("levant/deploy: triggering a deployment of job %s", *job.Name)

	eval, _, err := c.nomad.Jobs().Register(job, nil)
	if err != nil {
		logging.Error("levant/deploy: unable to register job %s with Nomad: %v", *job.Name, err)
		return
	}

	// GH-50: batch job types do not return an evaluation upon registration.
	if eval.EvalID == "" && *job.Type == nomadStructs.JobTypeBatch {
		return c.checkBatchJob(job.Name)
	}

	// Trigger the evaluationInspector to identify any potential errors in the
	// Nomad evaluation run. As far as I can tell from testing; a single alloc
	// failure in an evaluation means no allocs will be placed so we exit here.
	err = c.evaluationInspector(&eval.EvalID)
	if err != nil {
		logging.Error("levant/deploy: %v", err)
		return
	}

	switch *job.Type {
	case nomadStructs.JobTypeService:
		logging.Info("levant/deploy: beginning deployment watcher for job %s", *job.Name)

		// Get the deploymentID from the evaluationID so that we can watch the
		// deployment for end status.
		depID, err := c.getDeploymentID(eval.EvalID)
		if err != nil {
			logging.Error("levant/deploy: unable to get info of evaluation %s: %v", eval.EvalID, err)
			return
		}

		// Get the success of the deployment.
		success = c.deploymentWatcher(depID, autoPromote)

		// If the deployment has not been successful; check whether the job is
		// configured to auto-revert so that this can be tracked.
		if !success {
			dep, _, err := c.nomad.Deployments().Info(depID, nil)
			if err != nil {
				logging.Error("levant/deploy: unable to query deployment %s for auto-revert check: %v",
					dep.ID, err)
				break
			}
			c.checkAutoRevert(dep)
		}

	default:
		logging.Debug("levant/deploy: job type %s does not support Nomad deployment model", *job.Type)
		success = true
	}

	return
}

func (c *nomadClient) evaluationInspector(evalID *string) error {

	for {
		evalInfo, _, err := c.nomad.Evaluations().Info(*evalID, nil)
		if err != nil {
			return err
		}

		switch evalInfo.Status {
		case nomadStructs.EvalStatusComplete, nomadStructs.EvalStatusFailed, nomadStructs.EvalStatusCancelled:
			if len(evalInfo.FailedTGAllocs) == 0 {
				logging.Info("levant/deploy: evaluation %s finished successfully", *evalID)
				return nil
			}

			for group, metrics := range evalInfo.FailedTGAllocs {

				// Check if any nodes have been exhausted of resources and therfore are
				// unable to place allocs.
				if metrics.NodesExhausted > 0 {
					var exhausted, dimension []string
					for e := range metrics.ClassExhausted {
						exhausted = append(exhausted, e)
					}
					for d := range metrics.DimensionExhausted {
						dimension = append(dimension, d)
					}
					logging.Error("levant/deploy: task group %s failed to place allocs, failed on %v and exhausted %v",
						group, exhausted, dimension)
				}

				// Check if any node classes were filtered causing alloc placement
				// failures.
				if len(metrics.ClassFiltered) > 0 {
					for f := range metrics.ClassFiltered {
						logging.Error("levant/deploy: task group %s failed to place %v allocs as class \"%s\" was filtered",
							group, len(metrics.ClassFiltered), f)
					}
				}

				// Check if any node constraints were filtered causing alloc placement
				// failures.
				if len(metrics.ConstraintFiltered) > 0 {
					for cf := range metrics.ConstraintFiltered {
						logging.Error("levant/deploy: task group %s failed to place %v allocs as constraint \"%s\" was filtered",
							group, len(metrics.ConstraintFiltered), cf)
					}
				}
			}

			return fmt.Errorf("evaluation %v finished with status %s but failed to place allocations",
				*evalID, evalInfo.Status)

		default:
			time.Sleep(1 * time.Second)
			continue
		}
	}
}

func (c *nomadClient) deploymentWatcher(depID string, autoPromote int) (success bool) {

	var canaryChan chan interface{}
	deploymentChan := make(chan interface{})

	t := time.Now()
	wt := time.Duration(5 * time.Second)

	// Setup the canaryChan and launch the autoPromote go routine if autoPromote
	// has been enabled.
	if autoPromote > 0 {
		canaryChan = make(chan interface{})
		go c.canaryAutoPromote(depID, autoPromote, canaryChan, deploymentChan)
	}

	q := &nomad.QueryOptions{WaitIndex: 1, AllowStale: true, WaitTime: wt}

	for {

		dep, meta, err := c.nomad.Deployments().Info(depID, q)
		logging.Debug("levant/deploy: deployment %v running for %.2fs", depID, time.Since(t).Seconds())

		// Listen for the deploymentChan closing which indicates Levant should exit
		// the deployment watcher.
		select {
		case <-deploymentChan:
			return false
		default:
			break
		}

		if err != nil {
			logging.Error("levant/deploy: unable to get info of deployment %s: %v", depID, err)
			return
		}

		if meta.LastIndex <= q.WaitIndex {
			continue
		}

		q.WaitIndex = meta.LastIndex

		cont, err := c.checkDeploymentStatus(dep, canaryChan)
		if err != nil {
			return false
		}

		if cont {
			continue
		} else {
			return true
		}
	}
}

func (c *nomadClient) checkDeploymentStatus(dep *nomad.Deployment, shutdownChan chan interface{}) (bool, error) {

	switch dep.Status {
	case nomadStructs.DeploymentStatusSuccessful:
		logging.Info("levant/deploy: deployment %v has completed successfully", dep.ID)
		return false, nil
	case nomadStructs.DeploymentStatusRunning:
		return true, nil
	default:
		if shutdownChan != nil {
			logging.Debug("levant/deploy: deployment %v meaning canary auto promote will shutdown", dep.Status)
			close(shutdownChan)
		}

		logging.Error("levant/deploy: deployment %v has status %s", dep.ID, dep.Status)

		// Launch the failure inspector.
		c.checkFailedDeployment(&dep.ID)

		return false, fmt.Errorf("deployment failed")
	}
}

// canaryAutoPromote handles Levant's canary-auto-promote functionality.
func (c *nomadClient) canaryAutoPromote(depID string, waitTime int, shutdownChan, deploymentChan chan interface{}) {

	// Setup the AutoPromote timer.
	autoPromote := time.After(time.Duration(waitTime) * time.Second)

	for {
		select {
		case <-autoPromote:
			logging.Info("levant/deploy: auto-promote period %vs has been reached for deployment %s",
				waitTime, depID)

			// Check the deployment is healthy before promoting.
			if healthy := c.checkCanaryDeploymentHealth(depID); !healthy {
				logging.Error("levant/deploy: the canary deployment %s has unhealthy allocations, unable to promote", depID)
				close(deploymentChan)
				return
			}

			logging.Info("levant/deploy: triggering auto promote of deployment %s", depID)

			// Promote the deployment.
			_, _, err := c.nomad.Deployments().PromoteAll(depID, nil)
			if err != nil {
				logging.Error("levant/deploy: unable to promote deployment %s: %v", depID, err)
				close(deploymentChan)
				return
			}

		case <-shutdownChan:
			logging.Info("levant/deploy: canary auto promote has been shutdown")
			return
		}
	}
}

// checkCanaryDeploymentHealth is used to check the health status of each
// task-group within a canary deployment.
func (c *nomadClient) checkCanaryDeploymentHealth(depID string) (healthy bool) {

	var unhealthy int

	dep, _, err := c.nomad.Deployments().Info(depID, &nomad.QueryOptions{AllowStale: true})
	if err != nil {
		logging.Error("levant/deploy: unable to query deployment %s for health: %v", depID, err)
		return
	}

	// Itertate each task in the deployment to determine is health status. If an
	// unhealthy task is found, incrament the unhealthy counter.
	for taskName, taskInfo := range dep.TaskGroups {
		// skip any task groups which are not configured for canary deployments
		if taskInfo.DesiredCanaries == 0 {
			logging.Debug("levant/deploy: task %s has no desired canaries, skipping health checks in deployment %s", taskName, depID)
			continue
		}

		if taskInfo.DesiredCanaries != taskInfo.HealthyAllocs {
			logging.Error("levant/deploy: task %s has unhealthy allocations in deployment %s", taskName, depID)
			unhealthy++
		}
	}

	// If zero unhealthy tasks were found, continue with the auto promotion.
	if unhealthy == 0 {
		logging.Debug("levant/deploy: deployment %s has 0 unhealthy allocations", depID)
		healthy = true
	}

	return
}

// getDeploymentID finds the Nomad deploymentID associated to a Nomad
// evaluationID. This is only needed as sometimes Nomad initially returns eval
// info with an empty deploymentID; and a retry is required in order to get the
// updated response from Nomad.
func (c *nomadClient) getDeploymentID(evalID string) (depID string, err error) {

	var evalInfo *nomad.Evaluation

	for {
		if evalInfo, _, err = c.nomad.Evaluations().Info(evalID, nil); err != nil {
			return
		}

		if evalInfo.DeploymentID == "" {
			logging.Debug("levant/deploy: Nomad returned an empty deployment for evaluation %v; retrying", evalID)
			time.Sleep(2 * time.Second)
			continue
		} else {
			break
		}
	}

	return evalInfo.DeploymentID, nil
}

// dynamicGroupCountUpdater takes the templated and rendered job and updates the
// group counts based on the currently deployed job; if its running.
func (c *nomadClient) dynamicGroupCountUpdater(job *nomad.Job) error {

	// Gather information about the current state, if any, of the job on the
	// Nomad cluster.
	rJob, _, err := c.nomad.Jobs().Info(*job.Name, &nomad.QueryOptions{})

	// This is a hack due to GH-1849; we check the error string for 404 which
	// indicates the job is not running, not that there was an error in the API
	// call.
	if err != nil && strings.Contains(err.Error(), "404") {
		logging.Info("levant/deploy: job %s not running, using template file group counts", *job.Name)
		return nil
	} else if err != nil {
		logging.Error("levant/deploy: unable to perform job evaluation: %v", err)
		return err
	}

	// Iterate the templated job and the Nomad returned job and update group count
	// based on matches.
	for _, rGroup := range rJob.TaskGroups {
		for _, group := range job.TaskGroups {
			if *rGroup.Name == *group.Name {
				logging.Info("levant/deploy: using dynamic count %v for job %s and group %s",
					*rGroup.Count, *job.Name, *group.Name)
				group.Count = rGroup.Count
			}
		}
	}
	return nil
}
<|endoftext|>