<commit_before>package pwr

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"

	"github.com/itchio/wharf/wire"
)

func swap(a []int, i, j int) { a[i], a[j] = a[j], a[i] }

// Ternary-Split Quicksort, cf. http://www.larsson.dogma.net/ssrev-tr.pdf
func split(I, V []int, start, length, h int) {
	var i, j, k, x, jj, kk int

	if length < 16 {
		for k = start; k < start+length; k += j {
			j = 1
			x = V[I[k]+h]
			for i = 1; k+i < start+length; i++ {
				if V[I[k+i]+h] < x {
					x = V[I[k+i]+h]
					j = 0
				}
				if V[I[k+i]+h] == x {
					swap(I, k+i, k+j)
					j++
				}
			}
			for i = 0; i < j; i++ {
				V[I[k+i]] = k + j - 1
			}
			if j == 1 {
				I[k] = -1
			}
		}
		return
	}

	x = V[I[start+length/2]+h]
	jj = 0
	kk = 0
	for i = start; i < start+length; i++ {
		if V[I[i]+h] < x {
			jj++
		}
		if V[I[i]+h] == x {
			kk++
		}
	}
	jj += start
	kk += jj

	i = start
	j = 0
	k = 0
	for i < jj {
		if V[I[i]+h] < x {
			i++
		} else if V[I[i]+h] == x {
			swap(I, i, jj+j)
			j++
		} else {
			swap(I, i, kk+k)
			k++
		}
	}

	for jj+j < kk {
		if V[I[jj+j]+h] == x {
			j++
		} else {
			swap(I, jj+j, kk+k)
			k++
		}
	}

	if jj > start {
		split(I, V, start, jj-start, h)
	}

	for i = 0; i < kk-jj; i++ {
		V[I[jj+i]] = kk - 1
	}
	if jj == kk-1 {
		I[jj] = -1
	}

	if start+length > kk {
		split(I, V, kk, start+length-kk, h)
	}
}

// Faster Suffix Sorting, see: http://www.larsson.dogma.net/ssrev-tr.pdf
// Output `I` is a sorted suffix array.
func qsufsort(obuf []byte) []int {
	var buckets [256]int
	var i, h int
	I := make([]int, len(obuf)+1)
	V := make([]int, len(obuf)+1)

	for _, c := range obuf {
		buckets[c]++
	}
	for i = 1; i < 256; i++ {
		buckets[i] += buckets[i-1]
	}
	copy(buckets[1:], buckets[:])
	buckets[0] = 0

	for i, c := range obuf {
		buckets[c]++
		I[buckets[c]] = i
	}

	I[0] = len(obuf)
	for i, c := range obuf {
		V[i] = buckets[c]
	}

	V[len(obuf)] = 0
	for i = 1; i < 256; i++ {
		if buckets[i] == buckets[i-1]+1 {
			I[buckets[i]] = -1
		}
	}
	I[0] = -1

	for h = 1; I[0] != -(len(obuf) + 1); h += h {
		var n int
		for i = 0; i < len(obuf)+1; {
			if I[i] < 0 {
				n -= I[i]
				i -= I[i]
			} else {
				if n != 0 {
					I[i-n] = -n
				}
				n = V[I[i]] + 1 - i
				split(I, V, i, n, h)
				i += n
				n = 0
			}
		}
		if n != 0 {
			I[i-n] = -n
		}
	}

	for i = 0; i < len(obuf)+1; i++ {
		I[V[i]] = i
	}
	return I
}

// Returns the number of bytes common to a and b
func matchlen(a, b []byte) (i int) {
	for i < len(a) && i < len(b) && a[i] == b[i] {
		i++
	}
	return i
}

func search(I []int, obuf, nbuf []byte, st, en int) (pos, n int) {
	if en-st < 2 {
		x := matchlen(obuf[I[st]:], nbuf)
		y := matchlen(obuf[I[en]:], nbuf)

		if x > y {
			return I[st], x
		}
		return I[en], y
	}

	x := st + (en-st)/2
	if bytes.Compare(obuf[I[x]:], nbuf) < 0 {
		return search(I, obuf, nbuf, x, en)
	}
	return search(I, obuf, nbuf, st, x)
}

// BSDiff computes the difference between old and new, according to the bsdiff
// algorithm, and writes the result to patch.
func BSDiff(old, new io.Reader, patch wire.WriteContext) error {
	obuf, err := ioutil.ReadAll(old)
	if err != nil {
		return err
	}

	nbuf, err := ioutil.ReadAll(new)
	if err != nil {
		return err
	}

	var lenf int
	I := qsufsort(obuf)
	db := make([]byte, len(nbuf))
	eb := make([]byte, len(nbuf))
	var dblen, eblen int

	bsdc := &BsdiffControl{}

	// Compute the differences, writing ctrl as we go
	var scan, pos, length int
	var lastscan, lastpos, lastoffset int
	for scan < len(nbuf) {
		var oldscore int
		scan += length
		for scsc := scan; scan < len(nbuf); scan++ {
			pos, length = search(I, obuf, nbuf[scan:], 0, len(obuf))

			for ; scsc < scan+length; scsc++ {
				if scsc+lastoffset < len(obuf) &&
					obuf[scsc+lastoffset] == nbuf[scsc] {
					oldscore++
				}
			}

			if (length == oldscore && length != 0) || length > oldscore+8 {
				break
			}

			if scan+lastoffset < len(obuf) && obuf[scan+lastoffset] == nbuf[scan] {
				oldscore--
			}
		}

		if length != oldscore || scan == len(nbuf) {
			var s, Sf int
			lenf = 0
			for i := 0; lastscan+i < scan && lastpos+i < len(obuf); {
				if obuf[lastpos+i] == nbuf[lastscan+i] {
					s++
				}
				i++
				if s*2-i > Sf*2-lenf {
					Sf = s
					lenf = i
				}
			}

			lenb := 0
			if scan < len(nbuf) {
				var s, Sb int
				for i := 1; (scan >= lastscan+i) && (pos >= i); i++ {
					if obuf[pos-i] == nbuf[scan-i] {
						s++
					}
					if s*2-i > Sb*2-lenb {
						Sb = s
						lenb = i
					}
				}
			}

			if lastscan+lenf > scan-lenb {
				overlap := (lastscan + lenf) - (scan - lenb)
				s := 0
				Ss := 0
				lens := 0
				for i := 0; i < overlap; i++ {
					if nbuf[lastscan+lenf-overlap+i] == obuf[lastpos+lenf-overlap+i] {
						s++
					}
					if nbuf[scan-lenb+i] == obuf[pos-lenb+i] {
						s--
					}
					if s > Ss {
						Ss = s
						lens = i + 1
					}
				}

				lenf += lens - overlap
				lenb -= lens
			}

			for i := 0; i < lenf; i++ {
				db[dblen+i] = nbuf[lastscan+i] - obuf[lastpos+i]
			}
			for i := 0; i < (scan-lenb)-(lastscan+lenf); i++ {
				eb[eblen+i] = nbuf[lastscan+lenf+i]
			}

			dblen += lenf
			eblen += (scan - lenb) - (lastscan + lenf)

			bsdc.Add = int64(lenf)
			bsdc.Copy = int64((scan - lenb) - (lastscan + lenf))
			bsdc.Seek = int64((pos - lenb) - (lastpos + lenf))

			err := patch.WriteMessage(bsdc)
			if err != nil {
				return err
			}

			lastscan = scan - lenb
			lastpos = pos - lenb
			lastoffset = pos - scan
		}
	}

	// Write sentinel control message
	bsdc.Reset()
	bsdc.Add = -1
	err = patch.WriteMessage(bsdc)
	if err != nil {
		return err
	}

	// Write diff data
	bsdd := &BsdiffDiff{
		Data: db[:dblen],
	}

	err = patch.WriteMessage(bsdd)
	if err != nil {
		return err
	}

	// Write extra data
	bsed := &BsdiffExtra{
		Data: eb[:eblen],
	}

	err = patch.WriteMessage(bsed)
	if err != nil {
		return err
	}

	return nil
}

// ErrCorrupt indicates that a patch is corrupted, most often that it would produce a longer file
// than specified
var ErrCorrupt = errors.New("corrupt patch")

// BSPatch applies patch to old, according to the bspatch algorithm,
// and writes the result to new.
func BSPatch(old io.Reader, new io.Writer, newSize int64, patch wire.ReadContext) error {
	var ctrlOps []BsdiffControl

	ctrlOp := &BsdiffControl{}

	for {
		ctrlOp.Reset()
		err := patch.ReadMessage(ctrlOp)
		if err != nil {
			return err
		}

		if ctrlOp.Add == -1 {
			break
		}

		ctrlOps = append(ctrlOps, BsdiffControl{
			Add:  ctrlOp.Add,
			Seek: ctrlOp.Seek,
			Copy: ctrlOp.Copy,
		})
	}

	diffMessage := &BsdiffDiff{}
	err := patch.ReadMessage(diffMessage)
	if err != nil {
		return err
	}
	diff := diffMessage.Data

	var diffOffset int64

	extraMessage := &BsdiffExtra{}
	err = patch.ReadMessage(extraMessage)
	if err != nil {
		return err
	}
	extra := extraMessage.Data

	var extraOffset int64

	obuf, err := ioutil.ReadAll(old)
	if err != nil {
		return err
	}

	nbuf := make([]byte, newSize)

	var oldpos, newpos int64

	for _, ctrl := range ctrlOps {
		// Sanity-check
		if newpos+ctrl.Add > newSize {
			return ErrCorrupt
		}

		// Read diff string
		copy(nbuf[newpos:newpos+ctrl.Add], diff[diffOffset:diffOffset+ctrl.Add])
		diffOffset += ctrl.Add

		// Add old data to diff string
		for i := int64(0); i < ctrl.Add; i++ {
			nbuf[newpos+i] += obuf[oldpos+i]
		}

		// Adjust pointers
		newpos += ctrl.Add
		oldpos += ctrl.Add

		// Sanity-check
		if newpos+ctrl.Copy > newSize {
			return ErrCorrupt
		}

		// Read extra string
		copy(nbuf[newpos:newpos+ctrl.Copy], extra[extraOffset:extraOffset+ctrl.Copy])
		extraOffset += ctrl.Copy

		// Adjust pointers
		newpos += ctrl.Copy
		oldpos += ctrl.Seek
	}

	// Write the new file
	for len(nbuf) > 0 {
		n, err := new.Write(nbuf)
		if err != nil {
			return err
		}
		nbuf = nbuf[n:]
	}

	return nil
}
<commit_msg>typing<commit_after>package pwr

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"

	"github.com/itchio/wharf/wire"
)

func swap(a []int, i, j int) { a[i], a[j] = a[j], a[i] }

// Ternary-Split Quicksort, cf. http://www.larsson.dogma.net/ssrev-tr.pdf
func split(I, V []int, start, length, h int) {
	var i, j, k, x, jj, kk int

	if length < 16 {
		for k = start; k < start+length; k += j {
			j = 1
			x = V[I[k]+h]
			for i = 1; k+i < start+length; i++ {
				if V[I[k+i]+h] < x {
					x = V[I[k+i]+h]
					j = 0
				}
				if V[I[k+i]+h] == x {
					swap(I, k+i, k+j)
					j++
				}
			}
			for i = 0; i < j; i++ {
				V[I[k+i]] = k + j - 1
			}
			if j == 1 {
				I[k] = -1
			}
		}
		return
	}

	x = V[I[start+length/2]+h]
	jj = 0
	kk = 0
	for i = start; i < start+length; i++ {
		if V[I[i]+h] < x {
			jj++
		}
		if V[I[i]+h] == x {
			kk++
		}
	}
	jj += start
	kk += jj

	i = start
	j = 0
	k = 0
	for i < jj {
		if V[I[i]+h] < x {
			i++
		} else if V[I[i]+h] == x {
			swap(I, i, jj+j)
			j++
		} else {
			swap(I, i, kk+k)
			k++
		}
	}

	for jj+j < kk {
		if V[I[jj+j]+h] == x {
			j++
		} else {
			swap(I, jj+j, kk+k)
			k++
		}
	}

	if jj > start {
		split(I, V, start, jj-start, h)
	}

	for i = 0; i < kk-jj; i++ {
		V[I[jj+i]] = kk - 1
	}
	if jj == kk-1 {
		I[jj] = -1
	}

	if start+length > kk {
		split(I, V, kk, start+length-kk, h)
	}
}

// Faster Suffix Sorting, see: http://www.larsson.dogma.net/ssrev-tr.pdf
// Output `I` is a sorted suffix array.
func qsufsort(obuf []byte) []int {
	var buckets [256]int
	var i, h int
	I := make([]int, len(obuf)+1)
	V := make([]int, len(obuf)+1)

	for _, c := range obuf {
		buckets[c]++
	}
	for i = 1; i < 256; i++ {
		buckets[i] += buckets[i-1]
	}
	copy(buckets[1:], buckets[:])
	buckets[0] = 0

	for i, c := range obuf {
		buckets[c]++
		I[buckets[c]] = i
	}

	I[0] = len(obuf)
	for i, c := range obuf {
		V[i] = buckets[c]
	}

	V[len(obuf)] = 0
	for i = 1; i < 256; i++ {
		if buckets[i] == buckets[i-1]+1 {
			I[buckets[i]] = -1
		}
	}
	I[0] = -1

	for h = 1; I[0] != -(len(obuf) + 1); h += h {
		var n int
		for i = 0; i < len(obuf)+1; {
			if I[i] < 0 {
				n -= I[i]
				i -= I[i]
			} else {
				if n != 0 {
					I[i-n] = -n
				}
				n = V[I[i]] + 1 - i
				split(I, V, i, n, h)
				i += n
				n = 0
			}
		}
		if n != 0 {
			I[i-n] = -n
		}
	}

	for i = 0; i < len(obuf)+1; i++ {
		I[V[i]] = i
	}
	return I
}

// Returns the number of bytes common to a and b
func matchlen(a, b []byte) (i int) {
	for i < len(a) && i < len(b) && a[i] == b[i] {
		i++
	}
	return i
}

func search(I []int, obuf, nbuf []byte, st, en int) (pos, n int) {
	if en-st < 2 {
		x := matchlen(obuf[I[st]:], nbuf)
		y := matchlen(obuf[I[en]:], nbuf)

		if x > y {
			return I[st], x
		}
		return I[en], y
	}

	x := st + (en-st)/2
	if bytes.Compare(obuf[I[x]:], nbuf) < 0 {
		return search(I, obuf, nbuf, x, en)
	}
	return search(I, obuf, nbuf, st, x)
}

// BSDiff computes the difference between old and new, according to the bsdiff
// algorithm, and writes the result to patch.
func BSDiff(old, new io.Reader, patch *wire.WriteContext) error {
	obuf, err := ioutil.ReadAll(old)
	if err != nil {
		return err
	}

	nbuf, err := ioutil.ReadAll(new)
	if err != nil {
		return err
	}

	var lenf int
	I := qsufsort(obuf)
	db := make([]byte, len(nbuf))
	eb := make([]byte, len(nbuf))
	var dblen, eblen int

	bsdc := &BsdiffControl{}

	// Compute the differences, writing ctrl as we go
	var scan, pos, length int
	var lastscan, lastpos, lastoffset int
	for scan < len(nbuf) {
		var oldscore int
		scan += length
		for scsc := scan; scan < len(nbuf); scan++ {
			pos, length = search(I, obuf, nbuf[scan:], 0, len(obuf))

			for ; scsc < scan+length; scsc++ {
				if scsc+lastoffset < len(obuf) &&
					obuf[scsc+lastoffset] == nbuf[scsc] {
					oldscore++
				}
			}

			if (length == oldscore && length != 0) || length > oldscore+8 {
				break
			}

			if scan+lastoffset < len(obuf) && obuf[scan+lastoffset] == nbuf[scan] {
				oldscore--
			}
		}

		if length != oldscore || scan == len(nbuf) {
			var s, Sf int
			lenf = 0
			for i := 0; lastscan+i < scan && lastpos+i < len(obuf); {
				if obuf[lastpos+i] == nbuf[lastscan+i] {
					s++
				}
				i++
				if s*2-i > Sf*2-lenf {
					Sf = s
					lenf = i
				}
			}

			lenb := 0
			if scan < len(nbuf) {
				var s, Sb int
				for i := 1; (scan >= lastscan+i) && (pos >= i); i++ {
					if obuf[pos-i] == nbuf[scan-i] {
						s++
					}
					if s*2-i > Sb*2-lenb {
						Sb = s
						lenb = i
					}
				}
			}

			if lastscan+lenf > scan-lenb {
				overlap := (lastscan + lenf) - (scan - lenb)
				s := 0
				Ss := 0
				lens := 0
				for i := 0; i < overlap; i++ {
					if nbuf[lastscan+lenf-overlap+i] == obuf[lastpos+lenf-overlap+i] {
						s++
					}
					if nbuf[scan-lenb+i] == obuf[pos-lenb+i] {
						s--
					}
					if s > Ss {
						Ss = s
						lens = i + 1
					}
				}

				lenf += lens - overlap
				lenb -= lens
			}

			for i := 0; i < lenf; i++ {
				db[dblen+i] = nbuf[lastscan+i] - obuf[lastpos+i]
			}
			for i := 0; i < (scan-lenb)-(lastscan+lenf); i++ {
				eb[eblen+i] = nbuf[lastscan+lenf+i]
			}

			dblen += lenf
			eblen += (scan - lenb) - (lastscan + lenf)

			bsdc.Add = int64(lenf)
			bsdc.Copy = int64((scan - lenb) - (lastscan + lenf))
			bsdc.Seek = int64((pos - lenb) - (lastpos + lenf))

			err := patch.WriteMessage(bsdc)
			if err != nil {
				return err
			}

			lastscan = scan - lenb
			lastpos = pos - lenb
			lastoffset = pos - scan
		}
	}

	// Write sentinel control message
	bsdc.Reset()
	bsdc.Add = -1
	err = patch.WriteMessage(bsdc)
	if err != nil {
		return err
	}

	// Write diff data
	bsdd := &BsdiffDiff{
		Data: db[:dblen],
	}

	err = patch.WriteMessage(bsdd)
	if err != nil {
		return err
	}

	// Write extra data
	bsed := &BsdiffExtra{
		Data: eb[:eblen],
	}

	err = patch.WriteMessage(bsed)
	if err != nil {
		return err
	}

	return nil
}

// ErrCorrupt indicates that a patch is corrupted, most often that it would produce a longer file
// than specified
var ErrCorrupt = errors.New("corrupt patch")

// BSPatch applies patch to old, according to the bspatch algorithm,
// and writes the result to new.
func BSPatch(old io.Reader, new io.Writer, newSize int64, patch *wire.ReadContext) error {
	var ctrlOps []BsdiffControl

	ctrlOp := &BsdiffControl{}

	for {
		ctrlOp.Reset()
		err := patch.ReadMessage(ctrlOp)
		if err != nil {
			return err
		}

		if ctrlOp.Add == -1 {
			break
		}

		ctrlOps = append(ctrlOps, BsdiffControl{
			Add:  ctrlOp.Add,
			Seek: ctrlOp.Seek,
			Copy: ctrlOp.Copy,
		})
	}

	diffMessage := &BsdiffDiff{}
	err := patch.ReadMessage(diffMessage)
	if err != nil {
		return err
	}
	diff := diffMessage.Data

	var diffOffset int64

	extraMessage := &BsdiffExtra{}
	err = patch.ReadMessage(extraMessage)
	if err != nil {
		return err
	}
	extra := extraMessage.Data

	var extraOffset int64

	obuf, err := ioutil.ReadAll(old)
	if err != nil {
		return err
	}

	nbuf := make([]byte, newSize)

	var oldpos, newpos int64

	for _, ctrl := range ctrlOps {
		// Sanity-check
		if newpos+ctrl.Add > newSize {
			return ErrCorrupt
		}

		// Read diff string
		copy(nbuf[newpos:newpos+ctrl.Add], diff[diffOffset:diffOffset+ctrl.Add])
		diffOffset += ctrl.Add

		// Add old data to diff string
		for i := int64(0); i < ctrl.Add; i++ {
			nbuf[newpos+i] += obuf[oldpos+i]
		}

		// Adjust pointers
		newpos += ctrl.Add
		oldpos += ctrl.Add

		// Sanity-check
		if newpos+ctrl.Copy > newSize {
			return ErrCorrupt
		}

		// Read extra string
		copy(nbuf[newpos:newpos+ctrl.Copy], extra[extraOffset:extraOffset+ctrl.Copy])
		extraOffset += ctrl.Copy

		// Adjust pointers
		newpos += ctrl.Copy
		oldpos += ctrl.Seek
	}

	// Write the new file
	for len(nbuf) > 0 {
		n, err := new.Write(nbuf)
		if err != nil {
			return err
		}
		nbuf = nbuf[n:]
	}

	return nil
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"
)

func sprint(x interface{}) StringType {
	return StringType(fmt.Sprint(x))
}

type dumpable interface {
	dump() Value
}

// Dump dumps a value into a string type value.
var Dump = NewLazyFunction(
	NewSignature([]string{"arg"}, nil, "", nil, nil, ""),
	func(ts ...*Thunk) Value {
		s, err := strictDump(ts[0].Eval())

		if err != nil {
			return err
		}

		return s
	})

func strictDump(v Value) (StringType, *Thunk) {
	switch x := v.(type) {
	case ErrorType:
		return "", Normal(x)
	case dumpable:
		v = x.dump()
	case stringable:
		v = x.string()
	default:
		panic(fmt.Errorf("Invalid value: %#v", x))
	}

	s, ok := ensureNormal(v).(StringType)

	if !ok {
		return "", NotStringError(v)
	}

	return s, nil
}

// StrictDump is a variant of Dump which evaluates input strictly.
func StrictDump(v Value) (string, *Thunk) {
	s, err := strictDump(ensureNormal(v))
	return string(s), err
}

// Equal checks if 2 values are equal or not.
var Equal = NewStrictFunction(
	NewSignature(
		[]string{"x", "y"}, nil, "",
		nil, nil, "",
	),
	func(ts ...*Thunk) (v Value) {
		defer func() {
			if r := recover(); r != nil {
				v = r
			}
		}()

		return BoolType(compare(ts[0].Eval(), ts[1].Eval()) == 0)
	})

// ensureNormal evaluates nested thunks into WHNF values.
// This function must be used with care because it prevents tail call
// elimination.
func ensureNormal(v Value) Value {
	if t, ok := v.(*Thunk); ok {
		return t.Eval()
	}

	return v
}

var identity = NewLazyFunction(
	NewSignature([]string{"arg"}, nil, "", nil, nil, ""),
	func(ts ...*Thunk) Value { return ts[0] })

// TypeOf returns a type name of an argument as a string.
var TypeOf = NewLazyFunction(
	NewSignature([]string{"arg"}, nil, "", nil, nil, ""),
	func(ts ...*Thunk) Value {
		// No case of effectType should be here.
		switch ts[0].Eval().(type) {
		case BoolType:
			return NewString("bool")
		case DictionaryType:
			return NewString("dict")
		case ListType:
			return NewString("list")
		case NilType:
			return NewString("nil")
		case NumberType:
			return NewString("number")
		case StringType:
			return NewString("string")

		case functionType:
			return NewString("function")
		case RawFunctionType:
			return NewString("function")

		case ErrorType:
			return NewString("error")
		}

		panic(fmt.Errorf("Invalid value: %#v", ts[0].Eval()))
	})
<commit_msg>Add comment on TypeOf function<commit_after>package core

import (
	"fmt"
)

func sprint(x interface{}) StringType {
	return StringType(fmt.Sprint(x))
}

type dumpable interface {
	dump() Value
}

// Dump dumps a value into a string type value.
var Dump = NewLazyFunction(
	NewSignature([]string{"arg"}, nil, "", nil, nil, ""),
	func(ts ...*Thunk) Value {
		s, err := strictDump(ts[0].Eval())

		if err != nil {
			return err
		}

		return s
	})

func strictDump(v Value) (StringType, *Thunk) {
	switch x := v.(type) {
	case ErrorType:
		return "", Normal(x)
	case dumpable:
		v = x.dump()
	case stringable:
		v = x.string()
	default:
		panic(fmt.Errorf("Invalid value: %#v", x))
	}

	s, ok := ensureNormal(v).(StringType)

	if !ok {
		return "", NotStringError(v)
	}

	return s, nil
}

// StrictDump is a variant of Dump which evaluates input strictly.
func StrictDump(v Value) (string, *Thunk) {
	s, err := strictDump(ensureNormal(v))
	return string(s), err
}

// Equal checks if 2 values are equal or not.
var Equal = NewStrictFunction(
	NewSignature(
		[]string{"x", "y"}, nil, "",
		nil, nil, "",
	),
	func(ts ...*Thunk) (v Value) {
		defer func() {
			if r := recover(); r != nil {
				v = r
			}
		}()

		return BoolType(compare(ts[0].Eval(), ts[1].Eval()) == 0)
	})

// ensureNormal evaluates nested thunks into WHNF values.
// This function must be used with care because it prevents tail call
// elimination.
func ensureNormal(v Value) Value {
	if t, ok := v.(*Thunk); ok {
		return t.Eval()
	}

	return v
}

var identity = NewLazyFunction(
	NewSignature([]string{"arg"}, nil, "", nil, nil, ""),
	func(ts ...*Thunk) Value { return ts[0] })

// TypeOf returns a type name of an argument as a string.
var TypeOf = NewLazyFunction(
	NewSignature([]string{"arg"}, nil, "", nil, nil, ""),
	func(ts ...*Thunk) Value {
		// No case of effectType should be here.
		switch ts[0].Eval().(type) {
		case BoolType:
			return NewString("bool")
		case DictionaryType:
			return NewString("dict")
		case ListType:
			return NewString("list")
		case NilType:
			return NewString("nil")
		case NumberType:
			return NewString("number")
		case StringType:
			return NewString("string")

		case functionType:
			return NewString("function")
		case RawFunctionType:
			return NewString("function")

		// TODO: Remove this line and use catch function to check if a value is an
		// error or not.
		case ErrorType:
			return NewString("error")
		}

		panic(fmt.Errorf("Invalid value: %#v", ts[0].Eval()))
	})
<|endoftext|>


<commit_before>package lenovo

const ScanCodeIndex = "##1234567890####qwertyuiop####asdfghjkl;####zxcvbnm###### "
var Charset string

type Sha256Digester interface {
    Digest([]byte) [32]byte
}

type Range struct {
    Start byte
    End byte
}

func Hash(hi, ho Sha256Digester, sndPwd, pwd []byte) [32]byte {
    digest := hi.Digest(pwd)
    copy(sndPwd, digest[:12])
    return ho.Digest(sndPwd)
}

func DecodePassword(pwd []byte) (decoded string) {
    for _, i := range pwd {
        decoded += string(ScanCodeIndex[i])
    }
    return decoded
}

func init() {
    Charset = genScanCodeCharset()
}

func genScanCodeCharset() string {
    var charset [38]byte
    c := 0
    ranges := []Range{
        {0x02, 0x0b},
        {0x10, 0x19},
        {0x1e, 0x1f},
        {0x20, 0x27},
        {0x2c, 0x32},
        {0x39, 0x39},
    }
    for _, r := range ranges {
        for i := r.Start; i <= r.End; i++ {
            charset[c] = i
            c += 1
        }
    }
    return string(charset[:])
}
<commit_msg>optimize lenovo specific hash calculation<commit_after>package lenovo

import (
	"github.com/juanmera/lenovata/intelsha"
	"unsafe"
)

const ScanCodeIndex = "##1234567890####qwertyuiop####asdfghjkl;####zxcvbnm###### "
var Charset string

type Range struct {
    Start byte
    End byte
}

type Hash struct {
    dataA [128]byte
    dataB [128]byte
    digestA [128]byte
    Digest [32]byte
    digestAS []byte
    digestBS []byte
    dataAS []byte
    dataBS []byte
    salt []byte
    pDataA unsafe.Pointer
    pDataB unsafe.Pointer
    pDigestA unsafe.Pointer
    pDigestB unsafe.Pointer
}

func New(salt []byte) *Hash {
	m := &Hash{
		salt: salt,
	}
	m.digestAS = m.digestA[:]
	m.digestBS = m.Digest[:]
	m.dataAS = m.dataA[:]
	m.dataBS = m.dataB[:]
	m.pDataA = unsafe.Pointer(&m.dataA)
	m.pDataB = unsafe.Pointer(&m.dataA)
	m.pDigestA = unsafe.Pointer(&m.digestA)
	m.pDigestB = unsafe.Pointer(&m.Digest)

	m.dataA[64] = 0x80
	m.dataA[126] = 0x02
	m.digestA[72] = 0x80
	m.digestA[126] = 0x02
	m.digestA[127] = 0x40
	return m
}

func (m *Hash) Hash(pwd []byte) {
	copy(m.digestAS, intelsha.InitialDigest)
	copy(m.digestBS, intelsha.InitialDigest)
	copy(m.dataAS, pwd)
    intelsha.Sha256RorxX8ms(m.pDataA, m.pDigestA, 2)
    copy(m.digestAS[12:], m.salt)
    intelsha.Sha256RorxX8ms(m.pDigestA, m.pDigestB, 2)
}

func DecodePassword(pwd []byte) (decoded string) {
    for _, i := range pwd {
        decoded += string(ScanCodeIndex[i])
    }
    return decoded
}

func init() {
    Charset = genScanCodeCharset()
}

func genScanCodeCharset() string {
    var charset [38]byte
    c := 0
    ranges := []Range{
        {0x02, 0x0b},
        {0x10, 0x19},
        {0x1e, 0x1f},
        {0x20, 0x27},
        {0x2c, 0x32},
        {0x39, 0x39},
    }
    for _, r := range ranges {
        for i := r.Start; i <= r.End; i++ {
            charset[c] = i
            c += 1
        }
    }
    return string(charset[:])
}
<|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
		}
	}

	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
	}

	// Periodic and parameterized jobs do not return an evaluation and therefore
	// can't perform the evaluationInspector.
	if !job.IsPeriodic() && !job.IsParameterized() {

		// 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)
		}

	case nomadStructs.JobTypeBatch:
		return c.checkJobStatus(job.Name)

	case nomadStructs.JobTypeSystem:
		return c.checkJobStatus(job.Name)

	default:
		logging.Debug("levant/deploy: Levant does not support advanced deployments of job type %s", *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
}
<commit_msg>Check that service jobs have an update stanza.<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
		}
	}

	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
	}

	// Periodic and parameterized jobs do not return an evaluation and therefore
	// can't perform the evaluationInspector.
	if !job.IsPeriodic() && !job.IsParameterized() {

		// 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:

		// If the service job doesn't have an update stanza, the job will not use
		// Nomad deployments.
		if job.Update == nil {
			logging.Info("levant/deploy: job %s is not configured with update stanza, consider adding to use deployments",
				*job.Name)
			return c.checkJobStatus(job.Name)
		}

		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)
		}

	case nomadStructs.JobTypeBatch:
		return c.checkJobStatus(job.Name)

	case nomadStructs.JobTypeSystem:
		return c.checkJobStatus(job.Name)

	default:
		logging.Debug("levant/deploy: Levant does not support advanced deployments of job type %s", *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|>


<commit_before>package lfsapi

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/git-lfs/git-lfs/config"
	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/rubyist/tracerx"
)

const MediaType = "application/vnd.git-lfs+json; charset=utf-8"

var (
	UserAgent = "git-lfs"
	httpRE    = regexp.MustCompile(`\Ahttps?://`)
)

func (c *Client) NewRequest(method string, e Endpoint, suffix string, body interface{}) (*http.Request, error) {
	sshRes, err := c.SSH.Resolve(e, method)
	if err != nil {
		tracerx.Printf("ssh: %s failed, error: %s, message: %s",
			e.SshUserAndHost, err.Error(), sshRes.Message,
		)

		if len(sshRes.Message) > 0 {
			return nil, errors.Wrap(err, sshRes.Message)
		}
		return nil, err
	}

	prefix := e.Url
	if len(sshRes.Href) > 0 {
		prefix = sshRes.Href
	}

	if !httpRE.MatchString(prefix) {
		urlfragment := strings.SplitN(prefix, "?", 2)[0]
		return nil, fmt.Errorf("missing protocol: %q", urlfragment)
	}

	req, err := http.NewRequest(method, joinURL(prefix, suffix), nil)
	if err != nil {
		return req, err
	}

	for key, value := range sshRes.Header {
		req.Header.Set(key, value)
	}
	req.Header.Set("Accept", MediaType)

	if body != nil {
		if merr := MarshalToRequest(req, body); merr != nil {
			return req, merr
		}
		req.Header.Set("Content-Type", MediaType)
	}

	return req, err
}

const slash = "/"

func joinURL(prefix, suffix string) string {
	if strings.HasSuffix(prefix, slash) {
		return prefix + suffix
	}
	return prefix + slash + suffix
}

// Do sends an HTTP request to get an HTTP response. It wraps net/http, adding
// extra headers, redirection handling, and error reporting.
func (c *Client) Do(req *http.Request) (*http.Response, error) {
	req.Header = c.extraHeadersFor(req)

	return c.do(req)
}

// do performs an *http.Request respecting redirects, and handles the response
// as defined in c.handleResponse. Notably, it does not alter the headers for
// the request argument in any way.
func (c *Client) do(req *http.Request) (*http.Response, error) {
	req.Header.Set("User-Agent", UserAgent)

	res, err := c.doWithRedirects(c.httpClient(req.Host), req, nil)
	if err != nil {
		return res, err
	}

	return res, c.handleResponse(res)
}

// Close closes any resources that this client opened.
func (c *Client) Close() error {
	return c.httpLogger.Close()
}

func (c *Client) extraHeadersFor(req *http.Request) http.Header {
	extraHeaders := c.extraHeaders(req.URL)
	if len(extraHeaders) == 0 {
		return req.Header
	}

	copy := make(http.Header, len(req.Header))
	for k, vs := range req.Header {
		copy[k] = vs
	}

	for k, vs := range extraHeaders {
		for _, v := range vs {
			copy[k] = append(copy[k], v)
		}
	}
	return copy
}

func (c *Client) extraHeaders(u *url.URL) map[string][]string {
	hdrs := c.uc.GetAll("http", u.String(), "extraHeader")
	m := make(map[string][]string, len(hdrs))

	for _, hdr := range hdrs {
		parts := strings.SplitN(hdr, ":", 2)
		if len(parts) < 2 {
			continue
		}

		k, v := parts[0], strings.TrimSpace(parts[1])

		m[k] = append(m[k], v)
	}
	return m
}

func (c *Client) doWithRedirects(cli *http.Client, req *http.Request, via []*http.Request) (*http.Response, error) {
	tracedReq, err := c.traceRequest(req)
	if err != nil {
		return nil, err
	}

	var retries int
	if n, ok := Retries(req); ok {
		retries = n
	} else {
		retries = defaultRequestRetries
	}

	var res *http.Response

	requests := tools.MaxInt(0, retries) + 1
	for i := 0; i < requests; i++ {
		res, err = cli.Do(req)
		if err == nil {
			break
		}

		if seek, ok := req.Body.(io.Seeker); ok {
			seek.Seek(0, io.SeekStart)
		}

		c.traceResponse(req, tracedReq, nil)
	}

	if err != nil {
		c.traceResponse(req, tracedReq, nil)
		return nil, err
	}

	if res == nil {
		return nil, nil
	}

	c.traceResponse(req, tracedReq, res)

	if res.StatusCode != 301 &&
		res.StatusCode != 302 &&
		res.StatusCode != 303 &&
		res.StatusCode != 307 &&
		res.StatusCode != 308 {

		// Above are the list of 3xx status codes that we know
		// how to handle below. If the status code contained in
		// the HTTP response was none of them, return the (res,
		// err) tuple as-is, otherwise handle the redirect.
		return res, err
	}

	redirectTo := res.Header.Get("Location")
	locurl, err := url.Parse(redirectTo)
	if err == nil && !locurl.IsAbs() {
		locurl = req.URL.ResolveReference(locurl)
		redirectTo = locurl.String()
	}

	via = append(via, req)
	if len(via) >= 3 {
		return res, errors.New("too many redirects")
	}

	redirectedReq, err := newRequestForRetry(req, redirectTo)
	if err != nil {
		return res, err
	}

	return c.doWithRedirects(cli, redirectedReq, via)
}

func (c *Client) httpClient(host string) *http.Client {
	c.clientMu.Lock()
	defer c.clientMu.Unlock()

	if c.gitEnv == nil {
		c.gitEnv = make(testEnv)
	}

	if c.osEnv == nil {
		c.osEnv = make(testEnv)
	}

	if c.hostClients == nil {
		c.hostClients = make(map[string]*http.Client)
	}

	if client, ok := c.hostClients[host]; ok {
		return client
	}

	concurrentTransfers := c.ConcurrentTransfers
	if concurrentTransfers < 1 {
		concurrentTransfers = 8
	}

	dialtime := c.DialTimeout
	if dialtime < 1 {
		dialtime = 30
	}

	keepalivetime := c.KeepaliveTimeout
	if keepalivetime < 1 {
		keepalivetime = 1800
	}

	tlstime := c.TLSTimeout
	if tlstime < 1 {
		tlstime = 30
	}

	tr := &http.Transport{
		Proxy:               proxyFromClient(c),
		TLSHandshakeTimeout: time.Duration(tlstime) * time.Second,
		MaxIdleConnsPerHost: concurrentTransfers,
	}

	activityTimeout := 30
	if v, ok := c.uc.Get("lfs", fmt.Sprintf("https://%v", host), "activitytimeout"); ok {
		if i, err := strconv.Atoi(v); err == nil {
			activityTimeout = i
		} else {
			activityTimeout = 0
		}
	}

	dialer := &net.Dialer{
		Timeout:   time.Duration(dialtime) * time.Second,
		KeepAlive: time.Duration(keepalivetime) * time.Second,
		DualStack: true,
	}

	if activityTimeout > 0 {
		activityDuration := time.Duration(activityTimeout) * time.Second
		tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
			c, err := dialer.DialContext(ctx, network, addr)
			if c == nil {
				return c, err
			}
			if tc, ok := c.(*net.TCPConn); ok {
				tc.SetKeepAlive(true)
				tc.SetKeepAlivePeriod(dialer.KeepAlive)
			}
			return &deadlineConn{Timeout: activityDuration, Conn: c}, err
		}
	} else {
		tr.DialContext = dialer.DialContext
	}

	tr.TLSClientConfig = &tls.Config{}

	if isClientCertEnabledForHost(c, host) {
		tracerx.Printf("http: client cert for %s", host)
		tr.TLSClientConfig.Certificates = []tls.Certificate{getClientCertForHost(c, host)}
		tr.TLSClientConfig.BuildNameToCertificate()
	}

	if isCertVerificationDisabledForHost(c, host) {
		tr.TLSClientConfig.InsecureSkipVerify = true
	} else {
		tr.TLSClientConfig.RootCAs = getRootCAsForHost(c, host)
	}

	httpClient := &http.Client{
		Transport: tr,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	c.hostClients[host] = httpClient
	if c.VerboseOut == nil {
		c.VerboseOut = os.Stderr
	}

	return httpClient
}

func (c *Client) CurrentUser() (string, string) {
	userName, _ := c.gitEnv.Get("user.name")
	userEmail, _ := c.gitEnv.Get("user.email")
	return userName, userEmail
}

func newRequestForRetry(req *http.Request, location string) (*http.Request, error) {
	newReq, err := http.NewRequest(req.Method, location, nil)
	if err != nil {
		return nil, err
	}

	if req.URL.Scheme == "https" && newReq.URL.Scheme == "http" {
		return nil, errors.New("lfsapi/client: refusing insecure redirect, https->http")
	}

	sameHost := req.URL.Host == newReq.URL.Host
	for key := range req.Header {
		if key == "Authorization" {
			if !sameHost {
				continue
			}
		}
		newReq.Header.Set(key, req.Header.Get(key))
	}

	oldestURL := strings.SplitN(req.URL.String(), "?", 2)[0]
	newURL := strings.SplitN(newReq.URL.String(), "?", 2)[0]
	tracerx.Printf("api: redirect %s %s to %s", req.Method, oldestURL, newURL)

	// This body will have already been rewound from a call to
	// lfsapi.Client.traceRequest().
	newReq.Body = req.Body
	newReq.ContentLength = req.ContentLength

	// Copy the request's context.Context, if any.
	newReq = newReq.WithContext(req.Context())

	return newReq, nil
}

type deadlineConn struct {
	Timeout time.Duration
	net.Conn
}

func (c *deadlineConn) Read(b []byte) (int, error) {
	if err := c.Conn.SetDeadline(time.Now().Add(c.Timeout)); err != nil {
		return 0, err
	}
	return c.Conn.Read(b)
}

func (c *deadlineConn) Write(b []byte) (int, error) {
	if err := c.Conn.SetDeadline(time.Now().Add(c.Timeout)); err != nil {
		return 0, err
	}

	return c.Conn.Write(b)
}

func init() {
	UserAgent = config.VersionDesc
}
<commit_msg>Retry SSH resolution 5 times<commit_after>package lfsapi

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/git-lfs/git-lfs/config"
	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/rubyist/tracerx"
)

const MediaType = "application/vnd.git-lfs+json; charset=utf-8"

var (
	UserAgent                 = "git-lfs"
	httpRE                    = regexp.MustCompile(`\Ahttps?://`)
	maxSshAuthenticateRetries = 5
)

func (c *Client) sshResolveWithRetries(e Endpoint, method string) (sshAuthResponse, error) {
	var sshRes sshAuthResponse
	var err error

	requests := tools.MaxInt(0, maxSshAuthenticateRetries) + 1
	for i := 0; i < requests; i++ {
		sshRes, err = c.SSH.Resolve(e, method)
		if err == nil {
			return sshRes, nil
		}

		tracerx.Printf("ssh: %s failed, error: %s, message: %s",
			e.SshUserAndHost, err.Error(), sshRes.Message,
		)
	}

	if len(sshRes.Message) > 0 {
		return sshAuthResponse{}, errors.Wrap(err, sshRes.Message)
	}

	return sshAuthResponse{}, err
}

func (c *Client) NewRequest(method string, e Endpoint, suffix string, body interface{}) (*http.Request, error) {
	sshRes, err := c.sshResolveWithRetries(e, method)
	if err != nil {
		return nil, err
	}

	prefix := e.Url
	if len(sshRes.Href) > 0 {
		prefix = sshRes.Href
	}

	if !httpRE.MatchString(prefix) {
		urlfragment := strings.SplitN(prefix, "?", 2)[0]
		return nil, fmt.Errorf("missing protocol: %q", urlfragment)
	}

	req, err := http.NewRequest(method, joinURL(prefix, suffix), nil)
	if err != nil {
		return req, err
	}

	for key, value := range sshRes.Header {
		req.Header.Set(key, value)
	}
	req.Header.Set("Accept", MediaType)

	if body != nil {
		if merr := MarshalToRequest(req, body); merr != nil {
			return req, merr
		}
		req.Header.Set("Content-Type", MediaType)
	}

	return req, err
}

const slash = "/"

func joinURL(prefix, suffix string) string {
	if strings.HasSuffix(prefix, slash) {
		return prefix + suffix
	}
	return prefix + slash + suffix
}

// Do sends an HTTP request to get an HTTP response. It wraps net/http, adding
// extra headers, redirection handling, and error reporting.
func (c *Client) Do(req *http.Request) (*http.Response, error) {
	req.Header = c.extraHeadersFor(req)

	return c.do(req)
}

// do performs an *http.Request respecting redirects, and handles the response
// as defined in c.handleResponse. Notably, it does not alter the headers for
// the request argument in any way.
func (c *Client) do(req *http.Request) (*http.Response, error) {
	req.Header.Set("User-Agent", UserAgent)

	res, err := c.doWithRedirects(c.httpClient(req.Host), req, nil)
	if err != nil {
		return res, err
	}

	return res, c.handleResponse(res)
}

// Close closes any resources that this client opened.
func (c *Client) Close() error {
	return c.httpLogger.Close()
}

func (c *Client) extraHeadersFor(req *http.Request) http.Header {
	extraHeaders := c.extraHeaders(req.URL)
	if len(extraHeaders) == 0 {
		return req.Header
	}

	copy := make(http.Header, len(req.Header))
	for k, vs := range req.Header {
		copy[k] = vs
	}

	for k, vs := range extraHeaders {
		for _, v := range vs {
			copy[k] = append(copy[k], v)
		}
	}
	return copy
}

func (c *Client) extraHeaders(u *url.URL) map[string][]string {
	hdrs := c.uc.GetAll("http", u.String(), "extraHeader")
	m := make(map[string][]string, len(hdrs))

	for _, hdr := range hdrs {
		parts := strings.SplitN(hdr, ":", 2)
		if len(parts) < 2 {
			continue
		}

		k, v := parts[0], strings.TrimSpace(parts[1])

		m[k] = append(m[k], v)
	}
	return m
}

func (c *Client) doWithRedirects(cli *http.Client, req *http.Request, via []*http.Request) (*http.Response, error) {
	tracedReq, err := c.traceRequest(req)
	if err != nil {
		return nil, err
	}

	var retries int
	if n, ok := Retries(req); ok {
		retries = n
	} else {
		retries = defaultRequestRetries
	}

	var res *http.Response

	requests := tools.MaxInt(0, retries) + 1
	for i := 0; i < requests; i++ {
		res, err = cli.Do(req)
		if err == nil {
			break
		}

		if seek, ok := req.Body.(io.Seeker); ok {
			seek.Seek(0, io.SeekStart)
		}

		c.traceResponse(req, tracedReq, nil)
	}

	if err != nil {
		c.traceResponse(req, tracedReq, nil)
		return nil, err
	}

	if res == nil {
		return nil, nil
	}

	c.traceResponse(req, tracedReq, res)

	if res.StatusCode != 301 &&
		res.StatusCode != 302 &&
		res.StatusCode != 303 &&
		res.StatusCode != 307 &&
		res.StatusCode != 308 {

		// Above are the list of 3xx status codes that we know
		// how to handle below. If the status code contained in
		// the HTTP response was none of them, return the (res,
		// err) tuple as-is, otherwise handle the redirect.
		return res, err
	}

	redirectTo := res.Header.Get("Location")
	locurl, err := url.Parse(redirectTo)
	if err == nil && !locurl.IsAbs() {
		locurl = req.URL.ResolveReference(locurl)
		redirectTo = locurl.String()
	}

	via = append(via, req)
	if len(via) >= 3 {
		return res, errors.New("too many redirects")
	}

	redirectedReq, err := newRequestForRetry(req, redirectTo)
	if err != nil {
		return res, err
	}

	return c.doWithRedirects(cli, redirectedReq, via)
}

func (c *Client) httpClient(host string) *http.Client {
	c.clientMu.Lock()
	defer c.clientMu.Unlock()

	if c.gitEnv == nil {
		c.gitEnv = make(testEnv)
	}

	if c.osEnv == nil {
		c.osEnv = make(testEnv)
	}

	if c.hostClients == nil {
		c.hostClients = make(map[string]*http.Client)
	}

	if client, ok := c.hostClients[host]; ok {
		return client
	}

	concurrentTransfers := c.ConcurrentTransfers
	if concurrentTransfers < 1 {
		concurrentTransfers = 8
	}

	dialtime := c.DialTimeout
	if dialtime < 1 {
		dialtime = 30
	}

	keepalivetime := c.KeepaliveTimeout
	if keepalivetime < 1 {
		keepalivetime = 1800
	}

	tlstime := c.TLSTimeout
	if tlstime < 1 {
		tlstime = 30
	}

	tr := &http.Transport{
		Proxy:               proxyFromClient(c),
		TLSHandshakeTimeout: time.Duration(tlstime) * time.Second,
		MaxIdleConnsPerHost: concurrentTransfers,
	}

	activityTimeout := 30
	if v, ok := c.uc.Get("lfs", fmt.Sprintf("https://%v", host), "activitytimeout"); ok {
		if i, err := strconv.Atoi(v); err == nil {
			activityTimeout = i
		} else {
			activityTimeout = 0
		}
	}

	dialer := &net.Dialer{
		Timeout:   time.Duration(dialtime) * time.Second,
		KeepAlive: time.Duration(keepalivetime) * time.Second,
		DualStack: true,
	}

	if activityTimeout > 0 {
		activityDuration := time.Duration(activityTimeout) * time.Second
		tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
			c, err := dialer.DialContext(ctx, network, addr)
			if c == nil {
				return c, err
			}
			if tc, ok := c.(*net.TCPConn); ok {
				tc.SetKeepAlive(true)
				tc.SetKeepAlivePeriod(dialer.KeepAlive)
			}
			return &deadlineConn{Timeout: activityDuration, Conn: c}, err
		}
	} else {
		tr.DialContext = dialer.DialContext
	}

	tr.TLSClientConfig = &tls.Config{}

	if isClientCertEnabledForHost(c, host) {
		tracerx.Printf("http: client cert for %s", host)
		tr.TLSClientConfig.Certificates = []tls.Certificate{getClientCertForHost(c, host)}
		tr.TLSClientConfig.BuildNameToCertificate()
	}

	if isCertVerificationDisabledForHost(c, host) {
		tr.TLSClientConfig.InsecureSkipVerify = true
	} else {
		tr.TLSClientConfig.RootCAs = getRootCAsForHost(c, host)
	}

	httpClient := &http.Client{
		Transport: tr,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	c.hostClients[host] = httpClient
	if c.VerboseOut == nil {
		c.VerboseOut = os.Stderr
	}

	return httpClient
}

func (c *Client) CurrentUser() (string, string) {
	userName, _ := c.gitEnv.Get("user.name")
	userEmail, _ := c.gitEnv.Get("user.email")
	return userName, userEmail
}

func newRequestForRetry(req *http.Request, location string) (*http.Request, error) {
	newReq, err := http.NewRequest(req.Method, location, nil)
	if err != nil {
		return nil, err
	}

	if req.URL.Scheme == "https" && newReq.URL.Scheme == "http" {
		return nil, errors.New("lfsapi/client: refusing insecure redirect, https->http")
	}

	sameHost := req.URL.Host == newReq.URL.Host
	for key := range req.Header {
		if key == "Authorization" {
			if !sameHost {
				continue
			}
		}
		newReq.Header.Set(key, req.Header.Get(key))
	}

	oldestURL := strings.SplitN(req.URL.String(), "?", 2)[0]
	newURL := strings.SplitN(newReq.URL.String(), "?", 2)[0]
	tracerx.Printf("api: redirect %s %s to %s", req.Method, oldestURL, newURL)

	// This body will have already been rewound from a call to
	// lfsapi.Client.traceRequest().
	newReq.Body = req.Body
	newReq.ContentLength = req.ContentLength

	// Copy the request's context.Context, if any.
	newReq = newReq.WithContext(req.Context())

	return newReq, nil
}

type deadlineConn struct {
	Timeout time.Duration
	net.Conn
}

func (c *deadlineConn) Read(b []byte) (int, error) {
	if err := c.Conn.SetDeadline(time.Now().Add(c.Timeout)); err != nil {
		return 0, err
	}
	return c.Conn.Read(b)
}

func (c *deadlineConn) Write(b []byte) (int, error) {
	if err := c.Conn.SetDeadline(time.Now().Add(c.Timeout)); err != nil {
		return 0, err
	}

	return c.Conn.Write(b)
}

func init() {
	UserAgent = config.VersionDesc
}
<|endoftext|>


<commit_before>/*
CompileDaemon is a very simple compile daemon for Go.

CompileDaemon watches your .go files in a directory and invokes `go build`
if a file changes.

Examples

In its simplest form, the defaults will do. With the current working directory set
to the source directory you can simply…

    $ CompileDaemon

… and it will recompile your code whenever you save a source file.

If you want it to also run your program each time it builds you might add…

    $ CompileDaemon -command="./MyProgram -my-options"

… and it will also keep a copy of your program running. Killing the old one and
starting a new one each time you build.

You may find that you need to exclude some directories and files from
monitoring, such as a .git repository or emacs temporary files…

    $ CompileDaemon -exclude-dir=.git -exclude=".#*"

If you want to monitor files other than .go and .c files you might…

    $ CompileDaemon -include=Makefile -include="*.less" -include="*.tmpl"

Options

There are command line options.

	FILE SELECTION
	-directory=XXX    – which directory to monitor for changes
	-recursive=XXX    – look into subdirectories
	-exclude-dir=XXX  – exclude directories matching glob pattern XXX
	-exlude=XXX       – exclude files whose basename matches glob pattern XXX
	-include=XXX      – include files whose basename matches glob pattern XXX
	-pattern=XXX      – include files whose path matches regexp XXX

	ACTIONS
	-build=CCC        – Execute CCC to rebuild when a file changes
	-command=CCC      – Run command CCC after a successful build, stops previous command first

*/
package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/fatih/color"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

type globList []string

var excludedDirs globList
var excludedFiles globList
var includedFiles globList

func (g *globList) String() string {
	return fmt.Sprint(*g)
}
func (g *globList) Set(value string) error {
	*g = append(*g, value)
	return nil
}
func (g *globList) Matches(value string) bool {
	for _, v := range *g {
		if match, err := filepath.Match(v, value); err != nil {
			log.Fatalf("Bad pattern \"%s\": %s", v, err.Error())
		} else if match {
			return true
		}
	}
	return false
}

var (
	flag_directory = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern   = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command   = flag.String("command", "", "Command to run and restart after build")
	flag_recursive = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build     = flag.String("build", "go build", "Command to rebuild after changes")
	flag_color     = flag.Bool("color", false, "Colorize output for CompileDaemon status messages")
	flag_logprefix = flag.Bool("log-prefix", true, "Print log timestamps and subprocess stderr/stdout output")
)

var (
	_okColor   = color.GreenString
	_failColor = color.RedString
)

func okColor(format string, args ...interface{}) string {
	if *flag_color {
		return _okColor(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

func failColor(format string, args ...interface{}) string {
	if *flag_color {
		return _failColor(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println(okColor("Running build command!"))

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	cmd.Dir = *flag_directory

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println(okColor("Build ok."))
	} else {
		log.Println(failColor("Error while building:\n"), failColor(string(output)))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildDone chan<- struct{}) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()

	for {
		select {
		case <-jobs:
			threshold = createThreshold()
		case <-threshold:
			if build() {
				buildDone <- struct{}{}
			}
		}
	}
}

func logger(pipeChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			if *flag_logprefix {
				log.Print(prefix, " ", line)
			} else {
				log.Print(line)
			}
		}
	}

	for {
		pipe := <-pipeChan
		go dumper(pipe, "stdout:")

		pipe = <-pipeChan
		go dumper(pipe, "stderr:")
	}
}

// Start the supplied command and return stdout and stderr pipes for logging.
func startCommand(command string) (cmd *exec.Cmd, stdout io.ReadCloser, stderr io.ReadCloser, err error) {
	args := strings.Split(command, " ")
	cmd = exec.Command(args[0], args[1:]...)

	if stdout, err = cmd.StdoutPipe(); err != nil {
		err = fmt.Errorf("can't get stdout pipe for command: %s", err)
		return
	}

	if stderr, err = cmd.StderrPipe(); err != nil {
		err = fmt.Errorf("can't get stderr pipe for command: %s", err)
		return
	}

	if err = cmd.Start(); err != nil {
		err = fmt.Errorf("can't start command: %s", err)
		return
	}

	return
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(command string, buildDone <-chan struct{}) {
	var currentProcess *os.Process
	pipeChan := make(chan io.ReadCloser)

	go logger(pipeChan)

	for {
		<-buildDone

		if currentProcess != nil {
			if err := currentProcess.Kill(); err != nil {
				log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
			}

			_, werr := currentProcess.Wait()

			if werr != nil {
				log.Fatal(failColor("Could not wait for child process. Aborting due to danger of infinite forks."))
			}
		}

		log.Println(okColor("Restarting the given command."))
		cmd, stdoutPipe, stderrPipe, err := startCommand(command)

		if err != nil {
			log.Fatal(failColor("Could not start command:", err))
		}

		pipeChan <- stdoutPipe
		pipeChan <- stderrPipe

		currentProcess = cmd.Process
	}
}

func flusher(buildDone <-chan struct{}) {
	for {
		<-buildDone
	}
}

func main() {
	flag.Var(&excludedDirs, "exclude-dir", " Don't watch directories matching this name")
	flag.Var(&excludedFiles, "exclude", " Don't watch files matching this name")
	flag.Var(&includedFiles, "include", " Watch files matching this name")

	flag.Parse()

	if !*flag_logprefix {
		log.SetFlags(0)
	}

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	watcher, err := fsnotify.NewWatcher()

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

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				if excludedDirs.Matches(info.Name()) {
					return filepath.SkipDir
				} else {
					return watcher.Watch(path)
				}
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

	} else {
		if err := watcher.Watch(*flag_directory); err != nil {
			log.Fatal("watcher.Watch():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildDone := make(chan struct{})

	go builder(jobs, buildDone)

	if *flag_command != "" {
		go runner(*flag_command, buildDone)
	} else {
		go flusher(buildDone)
	}

	for {
		select {
		case ev := <-watcher.Event:
			if ev.Name != "" {
				base := filepath.Base(ev.Name)

				if includedFiles.Matches(base) || matchesPattern(pattern, ev.Name) {
					if !excludedFiles.Matches(base) {
						jobs <- ev.Name
					}
				}
			}

		case err := <-watcher.Error:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<commit_msg>Added -graceful-kill flag to support graceful termination of the child process<commit_after>/*
CompileDaemon is a very simple compile daemon for Go.

CompileDaemon watches your .go files in a directory and invokes `go build`
if a file changes.

Examples

In its simplest form, the defaults will do. With the current working directory set
to the source directory you can simply…

    $ CompileDaemon

… and it will recompile your code whenever you save a source file.

If you want it to also run your program each time it builds you might add…

    $ CompileDaemon -command="./MyProgram -my-options"

… and it will also keep a copy of your program running. Killing the old one and
starting a new one each time you build.

You may find that you need to exclude some directories and files from
monitoring, such as a .git repository or emacs temporary files…

    $ CompileDaemon -exclude-dir=.git -exclude=".#*"

If you want to monitor files other than .go and .c files you might…

    $ CompileDaemon -include=Makefile -include="*.less" -include="*.tmpl"

Options

There are command line options.

	FILE SELECTION
	-directory=XXX    – which directory to monitor for changes
	-recursive=XXX    – look into subdirectories
	-exclude-dir=XXX  – exclude directories matching glob pattern XXX
	-exlude=XXX       – exclude files whose basename matches glob pattern XXX
	-include=XXX      – include files whose basename matches glob pattern XXX
	-pattern=XXX      – include files whose path matches regexp XXX

	ACTIONS
	-build=CCC        – Execute CCC to rebuild when a file changes
	-command=CCC      – Run command CCC after a successful build, stops previous command first

*/
package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/fatih/color"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

type globList []string

var excludedDirs globList
var excludedFiles globList
var includedFiles globList

func (g *globList) String() string {
	return fmt.Sprint(*g)
}
func (g *globList) Set(value string) error {
	*g = append(*g, value)
	return nil
}
func (g *globList) Matches(value string) bool {
	for _, v := range *g {
		if match, err := filepath.Match(v, value); err != nil {
			log.Fatalf("Bad pattern \"%s\": %s", v, err.Error())
		} else if match {
			return true
		}
	}
	return false
}

var (
	flag_directory    = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern      = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command      = flag.String("command", "", "Command to run and restart after build")
	flag_recursive    = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build        = flag.String("build", "go build", "Command to rebuild after changes")
	flag_color        = flag.Bool("color", false, "Colorize output for CompileDaemon status messages")
	flag_logprefix    = flag.Bool("log-prefix", true, "Print log timestamps and subprocess stderr/stdout output")
	flag_gracefulkill = flag.Bool("graceful-kill", false, "Gracefully attempt to kill the child process by sending a SIGTERM first")
)

var (
	_okColor   = color.GreenString
	_failColor = color.RedString
)

func okColor(format string, args ...interface{}) string {
	if *flag_color {
		return _okColor(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

func failColor(format string, args ...interface{}) string {
	if *flag_color {
		return _failColor(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println(okColor("Running build command!"))

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	cmd.Dir = *flag_directory

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println(okColor("Build ok."))
	} else {
		log.Println(failColor("Error while building:\n"), failColor(string(output)))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildDone chan<- struct{}) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()

	for {
		select {
		case <-jobs:
			threshold = createThreshold()
		case <-threshold:
			if build() {
				buildDone <- struct{}{}
			}
		}
	}
}

func logger(pipeChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			if *flag_logprefix {
				log.Print(prefix, " ", line)
			} else {
				log.Print(line)
			}
		}
	}

	for {
		pipe := <-pipeChan
		go dumper(pipe, "stdout:")

		pipe = <-pipeChan
		go dumper(pipe, "stderr:")
	}
}

// Start the supplied command and return stdout and stderr pipes for logging.
func startCommand(command string) (cmd *exec.Cmd, stdout io.ReadCloser, stderr io.ReadCloser, err error) {
	args := strings.Split(command, " ")
	cmd = exec.Command(args[0], args[1:]...)

	if stdout, err = cmd.StdoutPipe(); err != nil {
		err = fmt.Errorf("can't get stdout pipe for command: %s", err)
		return
	}

	if stderr, err = cmd.StderrPipe(); err != nil {
		err = fmt.Errorf("can't get stderr pipe for command: %s", err)
		return
	}

	if err = cmd.Start(); err != nil {
		err = fmt.Errorf("can't start command: %s", err)
		return
	}

	return
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(command string, buildDone <-chan struct{}) {
	var currentProcess *os.Process
	pipeChan := make(chan io.ReadCloser)

	go logger(pipeChan)

	for {
		<-buildDone

		if currentProcess != nil {
			// If enabled, attempt to do a graceful shutdown of the child process.
			done := make(chan error, 1)
			go func() {
				if !*flag_gracefulkill {
					log.Println(okColor("Hard stopping the current process.."))
					done <- currentProcess.Kill()
					return
				}
				log.Println(okColor("Gracefully stopping the current process.."))
				if err := currentProcess.Signal(syscall.SIGTERM); err != nil {
					done <- err
					return
				}
				_, err := currentProcess.Wait()
				done <- err
			}()

			select {
			case <-time.After(3 * time.Second):
				log.Println(failColor("Could not gracefully stop the current process, proceeding to hard stop."))
				if err := currentProcess.Kill(); err != nil {
					log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
				}
				<-done
			case err := <-done:
				if err != nil {
					log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
				}
			}
		}

		log.Println(okColor("Restarting the given command."))
		cmd, stdoutPipe, stderrPipe, err := startCommand(command)

		if err != nil {
			log.Fatal(failColor("Could not start command:", err))
		}

		pipeChan <- stdoutPipe
		pipeChan <- stderrPipe

		currentProcess = cmd.Process
	}
}

func flusher(buildDone <-chan struct{}) {
	for {
		<-buildDone
	}
}

func main() {
	flag.Var(&excludedDirs, "exclude-dir", " Don't watch directories matching this name")
	flag.Var(&excludedFiles, "exclude", " Don't watch files matching this name")
	flag.Var(&includedFiles, "include", " Watch files matching this name")

	flag.Parse()

	if !*flag_logprefix {
		log.SetFlags(0)
	}

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	watcher, err := fsnotify.NewWatcher()

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

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				if excludedDirs.Matches(info.Name()) {
					return filepath.SkipDir
				} else {
					return watcher.Watch(path)
				}
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

	} else {
		if err := watcher.Watch(*flag_directory); err != nil {
			log.Fatal("watcher.Watch():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildDone := make(chan struct{})

	go builder(jobs, buildDone)

	if *flag_command != "" {
		go runner(*flag_command, buildDone)
	} else {
		go flusher(buildDone)
	}

	for {
		select {
		case ev := <-watcher.Event:
			if ev.Name != "" {
				base := filepath.Base(ev.Name)

				if includedFiles.Matches(base) || matchesPattern(pattern, ev.Name) {
					if !excludedFiles.Matches(base) {
						jobs <- ev.Name
					}
				}
			}

		case err := <-watcher.Error:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<|endoftext|>


<commit_before>package task_scheduler

import (
	"errors"
	"math/rand"
	"os"
	"sync"
	"time"

	executorapi "github.com/cloudfoundry-incubator/executor/api"
	"github.com/cloudfoundry-incubator/rep/routes"
	"github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/rata"
)

const ServerCloseErrMsg = "use of closed network connection"

var random = rand.New(rand.NewSource(time.Now().UnixNano()))

type TaskScheduler struct {
	callbackGenerator *rata.RequestGenerator

	executorID string
	bbs        bbs.RepBBS
	logger     lager.Logger
	stack      string
	client     executorapi.Client
	inFlight   *sync.WaitGroup
}

func New(
	executorID string,
	callbackGenerator *rata.RequestGenerator,
	bbs bbs.RepBBS,
	logger lager.Logger,
	stack string,
	executorClient executorapi.Client,
) *TaskScheduler {
	return &TaskScheduler{
		executorID:        executorID,
		callbackGenerator: callbackGenerator,

		bbs:    bbs,
		logger: logger.Session("task-scheduler"),
		stack:  stack,
		client: executorClient,

		inFlight: &sync.WaitGroup{},
	}
}

func (s *TaskScheduler) Run(signals <-chan os.Signal, readyChan chan<- struct{}) error {
	watchLog := s.logger.Session("watching")

	tasks, stopChan, errChan := s.bbs.WatchForDesiredTask()

	watchLog.Info("started")

	close(readyChan)

	for {
		select {
		case err := <-errChan:
			watchLog.Error("failed", err)

			time.Sleep(3 * time.Second)

			tasks, stopChan, errChan = s.bbs.WatchForDesiredTask()

		case task, ok := <-tasks:
			if !ok {
				err := errors.New("task channel closed. This is very unexpected, we did not intented to exit like this.")

				watchLog.Error("task-channel-closed", err)

				s.gracefulShutdown()
				return nil
			}

			s.inFlight.Add(1)
			go func() {
				defer s.inFlight.Done()
				s.handleTaskRequest(task)
			}()

		case <-signals:
			s.gracefulShutdown()
			close(stopChan)
			return nil
		}
	}
}

func (s *TaskScheduler) gracefulShutdown() {
	s.inFlight.Wait()
}

func (s *TaskScheduler) handleTaskRequest(task models.Task) {
	var err error

	if task.Stack != s.stack {
		return
	}

	taskLog := s.logger.Session("task-request")

	_, err = s.client.AllocateContainer(task.Guid, executorapi.ContainerAllocationRequest{
		DiskMB:   task.DiskMB,
		MemoryMB: task.MemoryMB,
	})
	if err != nil {
		taskLog.Error("failed-to-allocate", err)
		return
	}

	s.sleepForARandomInterval()

	err = s.bbs.ClaimTask(task.Guid, s.executorID)
	if err != nil {
		taskLog.Error("failed-to-claim-task", err)
		s.client.DeleteContainer(task.Guid)
		return
	}

	container, err := s.client.InitializeContainer(task.Guid, executorapi.ContainerInitializationRequest{
		CpuPercent: task.CpuPercent,
		Log: executorapi.LogConfig{
			Guid:       task.Log.Guid,
			SourceName: task.Log.SourceName,
		},
	})
	if err != nil {
		taskLog.Error("failed-to-initialize-container", err)
		s.client.DeleteContainer(task.Guid)
		s.markTaskAsFailed(taskLog, task.Guid, err)
		return
	}

	err = s.bbs.StartTask(task.Guid, s.executorID, container.ContainerHandle)
	if err != nil {
		taskLog.Error("failed-to-mark-task-started", err)
		s.client.DeleteContainer(task.Guid)
		return
	}

	callbackRequest, err := s.callbackGenerator.CreateRequest(routes.TaskCompleted, rata.Params{
		"guid": task.Guid,
	}, nil)
	if err != nil {
		taskLog.Error("failed-to-generate-callback-request", err)
		return
	}

	err = s.client.Run(task.Guid, executorapi.ContainerRunRequest{
		Actions:     task.Actions,
		CompleteURL: callbackRequest.URL.String(),
	})
	if err != nil {
		taskLog.Error("failed-to-run-actions", err)
		return
	}

	return
}

func (s *TaskScheduler) markTaskAsFailed(taskLog lager.Logger, taskGuid string, err error) {
	err = s.bbs.CompleteTask(taskGuid, true, "Failed to initialize container - "+err.Error(), "")
	if err != nil {
		taskLog.Error("failed-to-mark-task-failed", err)
	}
}

func (s *TaskScheduler) sleepForARandomInterval() {
	interval := random.Intn(1000)
	time.Sleep(time.Duration(interval) * time.Millisecond)
}
<commit_msg>Add logging to handleTaskRequest<commit_after>package task_scheduler

import (
	"errors"
	"math/rand"
	"os"
	"sync"
	"time"

	executorapi "github.com/cloudfoundry-incubator/executor/api"
	"github.com/cloudfoundry-incubator/rep/routes"
	"github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/rata"
)

const ServerCloseErrMsg = "use of closed network connection"

var random = rand.New(rand.NewSource(time.Now().UnixNano()))

type TaskScheduler struct {
	callbackGenerator *rata.RequestGenerator

	executorID string
	bbs        bbs.RepBBS
	logger     lager.Logger
	stack      string
	client     executorapi.Client
	inFlight   *sync.WaitGroup
}

func New(
	executorID string,
	callbackGenerator *rata.RequestGenerator,
	bbs bbs.RepBBS,
	logger lager.Logger,
	stack string,
	executorClient executorapi.Client,
) *TaskScheduler {
	return &TaskScheduler{
		executorID:        executorID,
		callbackGenerator: callbackGenerator,

		bbs:    bbs,
		logger: logger.Session("task-scheduler"),
		stack:  stack,
		client: executorClient,

		inFlight: &sync.WaitGroup{},
	}
}

func (s *TaskScheduler) Run(signals <-chan os.Signal, readyChan chan<- struct{}) error {
	watchLog := s.logger.Session("watching")

	tasks, stopChan, errChan := s.bbs.WatchForDesiredTask()

	watchLog.Info("started")

	close(readyChan)

	for {
		select {
		case err := <-errChan:
			watchLog.Error("failed", err)

			time.Sleep(3 * time.Second)

			tasks, stopChan, errChan = s.bbs.WatchForDesiredTask()

		case task, ok := <-tasks:
			if !ok {
				err := errors.New("task channel closed. This is very unexpected, we did not intented to exit like this.")

				watchLog.Error("task-channel-closed", err)

				s.gracefulShutdown()
				return nil
			}

			s.inFlight.Add(1)
			go func() {
				defer s.inFlight.Done()
				s.handleTaskRequest(task)
			}()

		case <-signals:
			s.gracefulShutdown()
			close(stopChan)
			return nil
		}
	}
}

func (s *TaskScheduler) gracefulShutdown() {
	s.inFlight.Wait()
}

func (s *TaskScheduler) handleTaskRequest(task models.Task) {
	var err error

	if task.Stack != s.stack {
		return
	}

	taskLog := s.logger.Session("task-request", lager.Data{"taskGuid": task.Guid})

	taskLog.Info("allocating-container")
	_, err = s.client.AllocateContainer(task.Guid, executorapi.ContainerAllocationRequest{
		DiskMB:   task.DiskMB,
		MemoryMB: task.MemoryMB,
	})
	if err != nil {
		taskLog.Error("failed-to-allocate-container", err)
		return
	}
	taskLog.Info("successfully-allocated-container")

	s.sleepForARandomInterval()

	taskLog.Info("claiming-task", lager.Data{"executorID": s.executorID})
	err = s.bbs.ClaimTask(task.Guid, s.executorID)
	if err != nil {
		taskLog.Info("failed-to-claim-task", lager.Data{"error": err.Error()})
		s.client.DeleteContainer(task.Guid)
		return
	}
	taskLog.Info("successfully-claimed-task")

	taskLog.Info("initializing-container")
	container, err := s.client.InitializeContainer(task.Guid, executorapi.ContainerInitializationRequest{
		CpuPercent: task.CpuPercent,
		Log: executorapi.LogConfig{
			Guid:       task.Log.Guid,
			SourceName: task.Log.SourceName,
		},
	})
	if err != nil {
		taskLog.Error("failed-to-initialize-container", err)
		s.client.DeleteContainer(task.Guid)
		s.markTaskAsFailed(taskLog, task.Guid, err)
		return
	}
	taskLog = taskLog.WithData(lager.Data{"containerHandle": container.ContainerHandle})
	taskLog.Info("successfully-initialized-container")

	taskLog.Info("starting-task")
	err = s.bbs.StartTask(task.Guid, s.executorID, container.ContainerHandle)
	if err != nil {
		taskLog.Error("failed-to-mark-task-started", err)
		s.client.DeleteContainer(task.Guid)
		return
	}
	taskLog.Info("successfully-started-task")

	callbackRequest, err := s.callbackGenerator.CreateRequest(routes.TaskCompleted, rata.Params{
		"guid": task.Guid,
	}, nil)
	if err != nil {
		taskLog.Error("failed-to-generate-callback-request", err)
		return
	}

	taskLog.Info("running-task")
	err = s.client.Run(task.Guid, executorapi.ContainerRunRequest{
		Actions:     task.Actions,
		CompleteURL: callbackRequest.URL.String(),
	})
	if err != nil {
		taskLog.Error("failed-to-run-task", err)
		return
	}
	taskLog.Info("successfully-ran-task")

	return
}

func (s *TaskScheduler) markTaskAsFailed(taskLog lager.Logger, taskGuid string, err error) {
	taskLog.Info("complete-task")
	err = s.bbs.CompleteTask(taskGuid, true, "Failed to initialize container - "+err.Error(), "")
	if err != nil {
		taskLog.Error("failed-to-complete-task", err)
	}
	taskLog.Info("successfully-completed-task")
}

func (s *TaskScheduler) sleepForARandomInterval() {
	interval := random.Intn(1000)
	time.Sleep(time.Duration(interval) * time.Millisecond)
}
<|endoftext|>


<commit_before>package miningpool

import (
	"errors"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/bytom/account"
	"github.com/bytom/mining"
	"github.com/bytom/protocol"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

const (
	blockUpdateMS   = 1000
	maxSubmitChSize = 50
)

type submitBlockMsg struct {
	blockHeader *types.BlockHeader
	reply       chan error
}

// MiningPool is the support struct for p2p mine pool
type MiningPool struct {
	mutex    sync.RWMutex
	block    *types.Block
	submitCh chan *submitBlockMsg

	chain          *protocol.Chain
	accountManager *account.Manager
	txPool         *protocol.TxPool
	newBlockCh     chan *bc.Hash
}

// NewMiningPool will create a new MiningPool
func NewMiningPool(c *protocol.Chain, accountManager *account.Manager, txPool *protocol.TxPool, newBlockCh chan *bc.Hash) *MiningPool {
	m := &MiningPool{
		submitCh:       make(chan *submitBlockMsg, maxSubmitChSize),
		chain:          c,
		accountManager: accountManager,
		txPool:         txPool,
		newBlockCh:     newBlockCh,
	}
	go m.blockUpdater()
	return m
}

// blockUpdater is the goroutine for keep update mining block
func (m *MiningPool) blockUpdater() {
	ticker := time.NewTicker(time.Millisecond * blockUpdateMS)
	for {
		select {
		case <-ticker.C:
			m.generateBlock()

		case submitMsg := <-m.submitCh:
			err := m.submitWork(submitMsg.blockHeader)
			if err == nil {
				m.generateBlock()
			}
			submitMsg.reply <- err
		}
	}
}

// generateBlock generates a block template to mine
func (m *MiningPool) generateBlock() {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	if m.block != nil && *m.chain.BestBlockHash() == m.block.PreviousBlockHash {
		m.block.Timestamp = uint64(time.Now().Unix())
		return
	}

	block, err := mining.NewBlockTemplate(m.chain, m.txPool, m.accountManager)
	if err != nil {
		log.Errorf("miningpool: failed on create NewBlockTemplate: %v", err)
		return
	}

	m.block = block
}

// GetWork will return a block header for p2p mining
func (m *MiningPool) GetWork() (*types.BlockHeader, error) {
	if m.block != nil {
		m.mutex.RLock()
		defer m.mutex.RUnlock()
		bh := m.block.BlockHeader
		return &bh, nil
	}
	return nil, errors.New("no block is ready for mining")
}

// SubmitWork will try to submit the result to the blockchain
func (m *MiningPool) SubmitWork(bh *types.BlockHeader) error {
	reply := make(chan error, 1)
	m.submitCh <- &submitBlockMsg{blockHeader: bh, reply: reply}
	err := <-reply
	if err != nil {
		log.WithFields(log.Fields{"err": err, "height": bh.Height}).Warning("submitWork failed")
	}
	return err
}

func (m *MiningPool) submitWork(bh *types.BlockHeader) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if m.block == nil || bh.PreviousBlockHash != m.block.PreviousBlockHash {
		return errors.New("pending mining block has been changed")
	}

	m.block.Nonce = bh.Nonce
	m.block.Timestamp = bh.Timestamp
	isOrphan, err := m.chain.ProcessBlock(m.block)
	if err != nil {
		return err
	}
	if isOrphan {
		return errors.New("submit result is orphan")
	}

	blockHash := bh.Hash()
	m.newBlockCh <- &blockHash
	return nil
}
<commit_msg>:zap: Improve miningpool blockUpdater() (#1264)<commit_after>package miningpool

import (
	"errors"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/bytom/account"
	"github.com/bytom/mining"
	"github.com/bytom/protocol"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

const (
	maxSubmitChSize = 50
)

type submitBlockMsg struct {
	blockHeader *types.BlockHeader
	reply       chan error
}

// MiningPool is the support struct for p2p mine pool
type MiningPool struct {
	mutex    sync.RWMutex
	block    *types.Block
	submitCh chan *submitBlockMsg

	chain          *protocol.Chain
	accountManager *account.Manager
	txPool         *protocol.TxPool
	newBlockCh     chan *bc.Hash
}

// NewMiningPool will create a new MiningPool
func NewMiningPool(c *protocol.Chain, accountManager *account.Manager, txPool *protocol.TxPool, newBlockCh chan *bc.Hash) *MiningPool {
	m := &MiningPool{
		submitCh:       make(chan *submitBlockMsg, maxSubmitChSize),
		chain:          c,
		accountManager: accountManager,
		txPool:         txPool,
		newBlockCh:     newBlockCh,
	}
	m.generateBlock()
	go m.blockUpdater()
	return m
}

// blockUpdater is the goroutine for keep update mining block
func (m *MiningPool) blockUpdater() {
	for {
		select {
		case <-m.chain.BlockWaiter(m.chain.BestBlockHeight() + 1):
			m.generateBlock()

		case submitMsg := <-m.submitCh:
			err := m.submitWork(submitMsg.blockHeader)
			if err == nil {
				m.generateBlock()
			}
			submitMsg.reply <- err
		}
	}
}

// generateBlock generates a block template to mine
func (m *MiningPool) generateBlock() {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	if m.block != nil && *m.chain.BestBlockHash() == m.block.PreviousBlockHash {
		m.block.Timestamp = uint64(time.Now().Unix())
		return
	}

	block, err := mining.NewBlockTemplate(m.chain, m.txPool, m.accountManager)
	if err != nil {
		log.Errorf("miningpool: failed on create NewBlockTemplate: %v", err)
		return
	}

	m.block = block
}

// GetWork will return a block header for p2p mining
func (m *MiningPool) GetWork() (*types.BlockHeader, error) {
	if m.block != nil {
		m.mutex.RLock()
		defer m.mutex.RUnlock()
		bh := m.block.BlockHeader
		return &bh, nil
	}
	return nil, errors.New("no block is ready for mining")
}

// SubmitWork will try to submit the result to the blockchain
func (m *MiningPool) SubmitWork(bh *types.BlockHeader) error {
	reply := make(chan error, 1)
	m.submitCh <- &submitBlockMsg{blockHeader: bh, reply: reply}
	err := <-reply
	if err != nil {
		log.WithFields(log.Fields{"err": err, "height": bh.Height}).Warning("submitWork failed")
	}
	return err
}

func (m *MiningPool) submitWork(bh *types.BlockHeader) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if m.block == nil || bh.PreviousBlockHash != m.block.PreviousBlockHash {
		return errors.New("pending mining block has been changed")
	}

	m.block.Nonce = bh.Nonce
	m.block.Timestamp = bh.Timestamp
	isOrphan, err := m.chain.ProcessBlock(m.block)
	if err != nil {
		return err
	}
	if isOrphan {
		return errors.New("submit result is orphan")
	}

	blockHash := bh.Hash()
	m.newBlockCh <- &blockHash
	return nil
}
<|endoftext|>


<commit_before>package main

import "fmt"

func add(args ...int) int {
    total := 0
    for _, v := range args {
        total += v
    }
    return total
}

func main() {
    fmt.Println(add(1,2,3))

    xs := [] int{1, 2, 3}

    fmt.Println("\nsum of xs = ")
    fmt.Println(add(xs...))
}

<commit_msg>closure: create function inside of functions<commit_after>package main

import "fmt"

func add(args ...int) int {
    total := 0
    for _, v := range args {
        total += v
    }
    return total
}

func main() {
    fmt.Println(add(1,2,3))

    xs := [] int{1, 2, 3}

    fmt.Println("\nsum of xs = ")
    fmt.Println(add(xs...))

    // Closure:  create function inside of function
    add := func(x, y int) int {
        return x + y
    }
    fmt.Println("add(1,1) = ")
    fmt.Println(add(1,1))
}

<|endoftext|>