<commit_before>package graph

import (
	"fmt"
	"strconv"

	plotly "github.com/nboughton/go-plotlytypes"
	"github.com/nboughton/stalotto/lotto"
)

const (
	tScatter = "scatter"
	tBar     = "bar"
	tLine    = "line"
)

// Data groups datasets together to exported and turned into a nice graph
type Data struct {
	Datasets []plotly.Dataset `json:"datasets"`
}

var formatTSLabel = "06/01/02"

// TimeSeries creates a Data struct for a time series graph
func TimeSeries(set lotto.ResultSet) Data {
	var d Data
	d.Datasets = make([]plotly.Dataset, lotto.BALLS+1)

	for _, row := range set {
		for ball := 0; ball < lotto.BALLS; ball++ {
			if d.Datasets[ball].Name == "" {
				d.Datasets[ball].Name = fmt.Sprintf("Ball %d", ball+1)
				d.Datasets[ball].Type = tLine
			}

			d.Datasets[ball].Y.AppendInt(row.Balls[ball])
			d.Datasets[ball].X.AppendStr(fmt.Sprintf("%s:%s:%d", row.Date.Format(formatTSLabel), row.Machine[:3], row.Set))
		}

		if d.Datasets[lotto.BALLS].Name == "" {
			d.Datasets[lotto.BALLS].Name = "Bonus"
		}

		d.Datasets[lotto.BALLS].Y.AppendInt(row.Bonus)
		d.Datasets[lotto.BALLS].X.AppendStr(fmt.Sprintf("%s:%s:%d", row.Date.Format(formatTSLabel), row.Machine[:3], row.Set))
	}

	return d
}

// FreqDist creates a Data struct for a frequency distribution graph
func FreqDist(set lotto.ResultSet) Data {
	var d Data
	d.Datasets = make([]plotly.Dataset, lotto.BALLS+1)
	// Populate Labels
	var labels []string
	for i := 0; i < lotto.MAXBALLVAL; i++ {
		labels = append(labels, fmt.Sprintf("%d", i+1))
	}

	for _, row := range set {
		for ball := 0; ball < lotto.BALLS; ball++ {
			if d.Datasets[ball].Name == "" { // Set Name and create Data
				d.Datasets[ball].Name = fmt.Sprintf("Ball %d", ball+1)
				d.Datasets[ball].Type = tBar
				d.Datasets[ball].Y = make([]string, lotto.MAXBALLVAL)
				d.Datasets[ball].X = labels
			}

			d.Datasets[ball].Y.AddInt(row.Balls[ball]-1, 1)
		}

		if d.Datasets[lotto.BALLS].Name == "" {
			d.Datasets[lotto.BALLS].Name = "Bonus"
			d.Datasets[lotto.BALLS].Type = tBar
			d.Datasets[lotto.BALLS].Y = make([]string, lotto.MAXBALLVAL)
		}

		d.Datasets[lotto.BALLS].Y.AddInt(row.Bonus-1, 1)
	}

	return d
}

// MachineSetDist creates a dataset appropriate for a machine/set distribution bubble graph
/*
func MachineSetDist(records <-chan lotto.Result) Data {
	var (
		d Data
		m = make(map[string]int)
	)

	for row := range records {
		v := fmt.Sprintf("%s:%d", row.Machine, row.Set)
		m[v]++
	}

	return d
}
*/
func addY(cur string, inc int) string {
	a, err := strconv.Atoi(cur)
	if err != nil {
		return cur // effectively ignore errors for now
	}

	return strconv.Itoa(a + inc)
}

func label(ball int) string {
	if ball < 6 {
		return fmt.Sprintf("Ball %d", ball+1)
	}

	return "Bonus"
}
<commit_msg>make full use of plotly.Axis<commit_after>package graph

import (
	"fmt"

	plotly "github.com/nboughton/go-plotlytypes"
	"github.com/nboughton/stalotto/lotto"
)

const (
	tScatter = "scatter"
	tBar     = "bar"
	tLine    = "line"
)

// Data groups datasets together to exported and turned into a nice graph
type Data struct {
	Datasets []plotly.Dataset `json:"datasets"`
}

var formatTSLabel = "06/01/02"

// TimeSeries creates a Data struct for a time series graph
func TimeSeries(set lotto.ResultSet) Data {
	var d Data
	d.Datasets = make([]plotly.Dataset, lotto.BALLS+1)

	for _, row := range set {
		for ball := 0; ball < lotto.BALLS; ball++ {
			if d.Datasets[ball].Name == "" {
				d.Datasets[ball].Name = fmt.Sprintf("Ball %d", ball+1)
				d.Datasets[ball].Type = tLine
			}

			d.Datasets[ball].Y.AppendInt(row.Balls[ball])
			d.Datasets[ball].X.AppendStr(fmt.Sprintf("%s:%s:%d", row.Date.Format(formatTSLabel), row.Machine[:3], row.Set))
		}

		if d.Datasets[lotto.BALLS].Name == "" {
			d.Datasets[lotto.BALLS].Name = "Bonus"
		}

		d.Datasets[lotto.BALLS].Y.AppendInt(row.Bonus)
		d.Datasets[lotto.BALLS].X.AppendStr(fmt.Sprintf("%s:%s:%d", row.Date.Format(formatTSLabel), row.Machine[:3], row.Set))
	}

	return d
}

// FreqDist creates a Data struct for a frequency distribution graph
func FreqDist(set lotto.ResultSet) Data {
	var d Data
	d.Datasets = make([]plotly.Dataset, lotto.BALLS+1)
	// Populate Labels
	var labels plotly.Axis
	for i := 0; i < lotto.MAXBALLVAL; i++ {
		labels.AppendStr(fmt.Sprintf("%d", i+1))
	}

	for _, row := range set {
		for ball := 0; ball < lotto.BALLS; ball++ {
			if d.Datasets[ball].Name == "" { // Set Name and create Data
				d.Datasets[ball].Name = fmt.Sprintf("Ball %d", ball+1)
				d.Datasets[ball].Type = tBar
				d.Datasets[ball].Y = make(plotly.Axis, lotto.MAXBALLVAL)
				d.Datasets[ball].X = labels
			}

			d.Datasets[ball].Y.AddInt(row.Balls[ball]-1, 1)
		}

		if d.Datasets[lotto.BALLS].Name == "" {
			d.Datasets[lotto.BALLS].Name = "Bonus"
			d.Datasets[lotto.BALLS].Type = tBar
			d.Datasets[lotto.BALLS].Y = make(plotly.Axis, lotto.MAXBALLVAL)
		}

		d.Datasets[lotto.BALLS].Y.AddInt(row.Bonus-1, 1)
	}

	return d
}

// MachineSetDist creates a dataset appropriate for a machine/set distribution bubble graph
/*
func MachineSetDist(records <-chan lotto.Result) Data {
	var (
		d Data
		m = make(map[string]int)
	)

	for row := range records {
		v := fmt.Sprintf("%s:%d", row.Machine, row.Set)
		m[v]++
	}

	return d
}
*/

func label(ball int) string {
	if ball < 6 {
		return fmt.Sprintf("Ball %d", ball+1)
	}

	return "Bonus"
}
<|endoftext|>


<commit_before>/*
Graph is used to store a pre-indxed (flattened) version of the
graph.  This is useful when attempting to flatten a large number
of nodes.  We can use some heuristics to determine if we should
reuse the pre-indexed positions or if it is beneficial to flatten
the nodes to find the optimal subgraph for execution.
*/
package graph

import (
	"log"
	"math"
	"runtime"
	"sort"
	"sync"
	"sync/atomic"

	"github.com/Workiva/go-datastructures/bitarray"
	"github.com/Workiva/go-datastructures/queue"
)

const (
	numNodesBeforeSplit int64 = 100
	goldenRatio               = .5
)

// bundle is a helper construct so that we bundle a node
// with its position in the flattened graph.
type bundle struct {
	INode
	position int64
}

// layer helps us order lists of nodes in a dense list of layers
type layer struct {
	position uint64
	nodes    Nodes
}

// layers is a helper struct that implements sort.Interface to help
// us sort a dense list of layers.
type layers []*layer

// Len returns the length of layers.
func (layers layers) Len() int {
	return len(layers)
}

// Less returns a bool indicating the relationship between
// items at the supplied positions.  IE, is layers[i] less than
// layers[j].
func (layers layers) Less(i, j int) bool {
	return layers[i].position < layers[j].position
}

// Swap will swap values at the provided positions.
func (layers layers) Swap(i, j int) {
	layers[i], layers[j] = layers[j], layers[i]
}

// positions stores a list of int64s which represent the layer
// of the subgraph where nodes can live.  The node's id represents
// the placement within the nodes slice (ie, position corresponds
// to node id).  A layer value of -1 indicates the formula is circular.
type positions []*bundle

// extractLayer returns a list of nodes that exist in the map of
// layer indices provided.
func (positions positions) extractLayers(layers map[int64]bool) Nodes {
	nodes := NewNodes()
	for _, node := range positions {
		if node == nil {
			continue
		}
		if _, ok := layers[node.position]; ok {
			nodes = append(nodes, node.INode)
		}
	}

	return nodes
}

// lowest returns the lowest index found in the group of nodes provided
func (positions positions) lowest(nodes Nodes) int64 {
	lowest := int64(math.MaxInt64)
	for _, node := range nodes {
		bundle := positions[node.ID()]
		if bundle == nil {
			continue
		}
		if bundle.position < lowest {
			lowest = bundle.position
		}
	}

	if lowest == int64(math.MaxInt64) {
		lowest = 0
	}

	return lowest
}

func (positions positions) highest(nodes Nodes) int64 {
	if len(nodes) == 0 {
		return -1
	}

	firstRun := true
	highest := int64(-1)

	for _, node := range nodes {
		nb := positions[node.ID()]
		if nb == nil {
			continue
		}

		if firstRun {
			firstRun = false
			highest = nb.position
			continue
		}

		if nb.position > highest {
			highest = nb.position
		}
	}

	return highest
}

// returns the highest position seen in the index
func (positions positions) highestSeen() int64 {
	highest := int64(-1)
	for _, node := range positions {
		if node == nil {
			continue
		}

		if node.position > highest {
			highest = node.position
		}
	}

	return highest
}

// flatten will take nodes and return a flattened list that
// is safe to execute in parallel.  Because this method relies
// on pre-indexed positions, this subgraph may not be optimal.
func (positions positions) flatten(nodes Nodes) (layers, Nodes) {
	var (
		circulars    = make(Nodes, 0, 5)
		layerMap     = make(map[uint64]Nodes, 5)
		position     int64
		nodePosition uint64
	)

	for _, node := range nodes {
		if node == nil {
			continue
		}
		id := node.ID() // reducing number of non-inlined method calls
		if id >= uint64(len(positions)) {
			log.Printf(
				`Trying to flatten a node that hasn't
						been added to the graph.  Node: %+v`, node,
			)
			continue
		}

		bundle := positions[id]
		if bundle == nil {
			continue
		}

		position = bundle.position
		if position == -1 {
			circulars = append(circulars, node)
			continue
		}

		// we've already confirmed position isn't negative so
		// this should be safe
		nodePosition = uint64(position)
		layerMap[nodePosition] = append(layerMap[nodePosition], node)
	}

	layers := make(layers, 0, len(layerMap))
	for position, nodes := range layerMap {
		layers = append(layers, &layer{
			position: position,
			nodes:    nodes,
		})
	}

	sort.Sort(layers)
	return layers, circulars
}

type Graph struct {
	positions positions
	maxLayer  int64
	numItems  uint64
}

// GetSubgraph will return an execution graph that needs to be calculated
// based on the supplied nodes.
func (g *Graph) GetSubgraph(nodes Nodes) *ExecutionGraph {
	// a large proportion of the graph needs to be recalculated
	layers, circulars := g.positions.flatten(nodes)
	flattened := make([]Nodes, 0, len(layers))
	size := int64(len(circulars))
	for _, layer := range layers {
		flattened = append(flattened, layer.nodes)
		size += int64(len(layer.nodes))
	}

	return &ExecutionGraph{
		size:      size,
		toApply:   flattened,
		circulars: circulars,
	}
}

// GetLowestNodes will return a list of nodes that match
// the lowest layer of the nodes provided.
func (g *Graph) GetLowestNodes(nodes Nodes) Nodes {
	if len(nodes) == 0 {
		return nil
	}

	toReturn := make(Nodes, 0, len(nodes))
	isSet := false
	lowest := int64(0)

	for _, node := range nodes {
		nb := g.positions[node.ID()]
		if nb == nil {
			log.Printf(`Node flattened that hasn't been added: %+v`, node)
			continue
		}

		if !isSet {
			isSet = true
			toReturn = append(toReturn, node)
			lowest = nb.position
			continue
		}

		if nb.position == lowest {
			toReturn = append(toReturn, node)
		} else if nb.position < lowest {
			toReturn = toReturn[:0]
			toReturn = append(toReturn, node)
			lowest = nb.position
		}
	}

	return toReturn
}

// AddNodes will add the provided nodes to the flattened index
// of the graph and return an execution graph that is ready to
// be calculated.
func (g *Graph) AddNodes(dp IDependencyProvider, nodes Nodes) *ExecutionGraph {
	dependentNodes := dp.GetDependents(nodes)

	highest := nodes.Highest()

	// want to make sure we don't overflow here
	if highest >= uint64(len(g.positions)) {
		diff := highest - uint64(len(g.positions))
		g.positions = append(g.positions, make(positions, diff+1)...)
	}

	var lowest int64
	if len(g.positions) > 0 && uint64(len(dependentNodes)) < g.numItems {
		dependencyNodes := g.findUniqueDependencies(dp, nodes)
		lowest = g.positions.highest(dependencyNodes) + 1
	}

	dependentNodes = append(dependentNodes, nodes...)
	flattened, circulars, size := flatten(dp, dependentNodes)
	g.insert(flattened, circulars, lowest)

	return &ExecutionGraph{
		size:      size,
		toApply:   flattened,
		circulars: circulars,
	}
}

func (g *Graph) findUniqueDependencies(dp IDependencyProvider, nodes Nodes) Nodes {
	if len(nodes) == 0 {
		return nil
	}

	ba := bitarray.NewSparseBitArray()
	for _, node := range nodes {
		ba.SetBit(node.ID())
	}
	ids := make([]uint64, 0, len(nodes))
	chunks := nodes.Split()
	var lock sync.Mutex
	var wg sync.WaitGroup
	wg.Add(len(chunks))
	for _, chunk := range chunks {
		go func(nodes Nodes) {
			for _, node := range nodes {
				deps := dp.GetDependencies(node)
				if deps == nil {
					continue
				}

				uints := deps.ToNums()
				lock.Lock()
				for _, id := range uints {
					if ok, _ := ba.GetBit(id); ok {
						continue
					}

					ids = append(ids, id)
					ba.SetBit(id)
				}
				lock.Unlock()
			}
			wg.Done()
		}(chunk)
	}

	wg.Wait()

	deps := make(Nodes, 0, len(ids))
	for _, id := range ids {
		nb := g.positions[id]
		if nb == nil {
			continue
		}

		deps = append(deps, nb.INode)
	}

	return deps
}

// RemoveNodes will remove the provided nodes from the graph.
func (g *Graph) RemoveNodes(dp IDependencyProvider, nodes Nodes) *ExecutionGraph {
	if len(nodes) == 0 {
		return &ExecutionGraph{}
	}

	lowest := g.positions.lowest(nodes)
	for _, node := range nodes {
		if g.positions[node.ID()] != nil {
			g.numItems--
		}
		g.positions[node.ID()] = nil // kill these first
	}

	maxLayer := g.maxLayer
	if maxLayer == -1 {
		maxLayer = 1
	}

	toExtract := make(map[int64]bool, maxLayer)
	if lowest == -1 { // we end up reordering the whole damn thing
		toExtract[-1] = true
		lowest = 0
	}

	for i := lowest; i <= g.maxLayer; i++ {
		toExtract[i] = true
	}

	nodesToFlatten := g.positions.extractLayers(toExtract)
	flattened, circulars, size := flatten(dp, nodesToFlatten)

	g.insert(flattened, circulars, lowest)
	// have to set new max layer, only in the delete case
	g.maxLayer = g.positions.highestSeen()

	return &ExecutionGraph{
		toApply:   flattened,
		circulars: circulars,
		size:      size,
	}
}

func (g *Graph) insert(flattened []Nodes, circulars Nodes, offset int64) {
	maxLayer := int64(-1)
	for i, nodes := range flattened {
		for _, node := range nodes {
			if g.positions[node.ID()] == nil {
				g.numItems++
			}
			g.positions[node.ID()] = &bundle{
				INode:    node,
				position: int64(i) + offset,
			}
		}
		maxLayer = int64(i) + offset
	}

	for _, node := range circulars {
		g.positions[node.ID()] = &bundle{
			INode: node, position: -1,
		}
	}

	if maxLayer > g.maxLayer {
		g.maxLayer = maxLayer
	}
}

// FromNodes will create a new graph from the given nodes.
func FromNodes(dp IDependencyProvider, nodes Nodes) *Graph {
	if len(nodes) == 0 {
		return &Graph{}
	}

	highest := nodes.Highest()
	positions := make(positions, highest+1)
	g := &Graph{
		positions: positions,
		maxLayer:  -1, // it's possible they are all circular
	}
	flattened, circulars, _ := flatten(dp, nodes)

	g.insert(flattened, circulars, 0)

	return g
}

func flatten(dp IDependencyProvider, nodes Nodes) ([]Nodes, Nodes, int64) {
	if len(nodes) == 0 {
		return nil, nil, 0
	}

	ba := bitarray.NewBitArray(dp.MaxNode(), true)
	remaining := make(Nodes, 0, len(nodes))
	helper := make(Nodes, len(nodes))
	for _, node := range nodes {
		if node != nil {
			ba.ClearBit(node.ID())
			remaining = append(remaining, node)
		}
	}

	var circulars Nodes
	results := make([]Nodes, 0, 5) // 5 is a guess
	layer := make([]INode, 0, len(nodes))
	size := int64(0)

	var wg sync.WaitGroup
	var lock sync.Mutex

	for {
		if len(remaining) == 0 && len(layer) == 0 {
			break
		}

		q := queue.New(int64(len(remaining)))
		for _, node := range remaining {
			q.Put(node)
		}

		completed := uint64(0)
		required := uint64(len(remaining))

		wg.Add(1)
		for i := 0; i < runtime.NumCPU(); i++ {
			go func() {
				for {
					items, err := q.Get(1)
					if err != nil { // queue was disposed
						break
					}

					node := items[0].(INode)
					okToAdd := true
					deps := dp.GetDependencies(node)
					if deps != nil && !ba.Intersects(deps) {
						okToAdd = false
					}

					if okToAdd {
						if node.IsCircular() {
							node.SetCircular(false)
						}
						lock.Lock()
						layer = append(layer, node)
						lock.Unlock()
					}

					if atomic.AddUint64(&completed, 1) == required {
						wg.Done()
						break
					}
				}
			}()
		}

		wg.Wait()
		q.Dispose()

		if len(layer) == 0 { // this is a problem, nothing new found
			circulars = make(Nodes, 0, len(remaining))
			for _, n := range remaining {
				n.SetCircular(true)
				circulars = append(circulars, n)
				size++
			}
			break
		}

		for _, node := range layer {
			size++
			ba.SetBit(node.ID())
		}

		results = append(results, layer)
		layer = make([]INode, 0, len(layer))
		copy(helper, remaining)
		remaining = remaining[:0]
		for i := 0; i < len(helper); i++ {
			node := helper[i]
			if node == nil {
				break
			}
			if ok, _ := ba.GetBit(node.ID()); ok {
				continue
			}

			remaining = append(remaining, node)
			helper[i] = nil
		}

		helper = helper[:len(remaining)]
	}

	return results, circulars, size
}

func New() *Graph {
	return &Graph{
		maxLayer: -1,
	}
}
<commit_msg>Added some optimizations to the graph code.<commit_after>/*
Graph is used to store a pre-indxed (flattened) version of the
graph.  This is useful when attempting to flatten a large number
of nodes.  We can use some heuristics to determine if we should
reuse the pre-indexed positions or if it is beneficial to flatten
the nodes to find the optimal subgraph for execution.
*/
package graph

import (
	"log"
	"math"
	"runtime"
	"sort"
	"sync"
	"sync/atomic"

	"github.com/Workiva/go-datastructures/bitarray"
	"github.com/Workiva/go-datastructures/queue"
)

const (
	numNodesBeforeSplit int64 = 100
	goldenRatio               = .5
)

// bundle is a helper construct so that we bundle a node
// with its position in the flattened graph.
type bundle struct {
	INode
	position int64
}

// layer helps us order lists of nodes in a dense list of layers
type layer struct {
	position uint64
	nodes    Nodes
}

// layers is a helper struct that implements sort.Interface to help
// us sort a dense list of layers.
type layers []*layer

// Len returns the length of layers.
func (layers layers) Len() int {
	return len(layers)
}

// Less returns a bool indicating the relationship between
// items at the supplied positions.  IE, is layers[i] less than
// layers[j].
func (layers layers) Less(i, j int) bool {
	return layers[i].position < layers[j].position
}

// Swap will swap values at the provided positions.
func (layers layers) Swap(i, j int) {
	layers[i], layers[j] = layers[j], layers[i]
}

// positions stores a list of int64s which represent the layer
// of the subgraph where nodes can live.  The node's id represents
// the placement within the nodes slice (ie, position corresponds
// to node id).  A layer value of -1 indicates the formula is circular.
type positions []*bundle

// extractLayer returns a list of nodes that exist in the map of
// layer indices provided.
func (positions positions) extractLayers(layers map[int64]bool) Nodes {
	nodes := NewNodes()
	for _, node := range positions {
		if node == nil {
			continue
		}
		if _, ok := layers[node.position]; ok {
			nodes = append(nodes, node.INode)
		}
	}

	return nodes
}

// lowest returns the lowest index found in the group of nodes provided
func (positions positions) lowest(nodes Nodes) int64 {
	lowest := int64(math.MaxInt64)
	for _, node := range nodes {
		bundle := positions[node.ID()]
		if bundle == nil {
			continue
		}
		if bundle.position < lowest {
			lowest = bundle.position
		}
	}

	if lowest == int64(math.MaxInt64) {
		lowest = 0
	}

	return lowest
}

func (positions positions) highest(nodes Nodes) int64 {
	if len(nodes) == 0 {
		return -1
	}

	firstRun := true
	highest := int64(-1)

	for _, node := range nodes {
		nb := positions[node.ID()]
		if nb == nil {
			continue
		}

		if firstRun {
			firstRun = false
			highest = nb.position
			continue
		}

		if nb.position > highest {
			highest = nb.position
		}
	}

	return highest
}

// returns the highest position seen in the index
func (positions positions) highestSeen() int64 {
	highest := int64(-1)
	for _, node := range positions {
		if node == nil {
			continue
		}

		if node.position > highest {
			highest = node.position
		}
	}

	return highest
}

// flatten will take nodes and return a flattened list that
// is safe to execute in parallel.  Because this method relies
// on pre-indexed positions, this subgraph may not be optimal.
func (positions positions) flatten(nodes Nodes) (layers, Nodes) {
	var (
		circulars    = make(Nodes, 0, 5)
		layerMap     = make(map[uint64]Nodes, 5)
		position     int64
		nodePosition uint64
	)

	for _, node := range nodes {
		if node == nil {
			continue
		}
		id := node.ID() // reducing number of non-inlined method calls
		if id >= uint64(len(positions)) {
			log.Printf(
				`Trying to flatten a node that hasn't
						been added to the graph.  Node: %+v`, node,
			)
			continue
		}

		bundle := positions[id]
		if bundle == nil {
			continue
		}

		position = bundle.position
		if position == -1 {
			circulars = append(circulars, node)
			continue
		}

		// we've already confirmed position isn't negative so
		// this should be safe
		nodePosition = uint64(position)
		layerMap[nodePosition] = append(layerMap[nodePosition], node)
	}

	layers := make(layers, 0, len(layerMap))
	for position, nodes := range layerMap {
		layers = append(layers, &layer{
			position: position,
			nodes:    nodes,
		})
	}

	sort.Sort(layers)
	return layers, circulars
}

type Graph struct {
	positions positions
	maxLayer  int64
	numItems  uint64
}

// GetSubgraph will return an execution graph that needs to be calculated
// based on the supplied nodes.
func (g *Graph) GetSubgraph(nodes Nodes) *ExecutionGraph {
	// a large proportion of the graph needs to be recalculated
	layers, circulars := g.positions.flatten(nodes)
	flattened := make([]Nodes, 0, len(layers))
	size := int64(len(circulars))
	for _, layer := range layers {
		flattened = append(flattened, layer.nodes)
		size += int64(len(layer.nodes))
	}

	return &ExecutionGraph{
		size:      size,
		toApply:   flattened,
		circulars: circulars,
	}
}

// GetLowestNodes will return a list of nodes that match
// the lowest layer of the nodes provided.
func (g *Graph) GetLowestNodes(nodes Nodes) Nodes {
	if len(nodes) == 0 {
		return nil
	}

	toReturn := make(Nodes, 0, len(nodes))
	isSet := false
	lowest := int64(0)

	for _, node := range nodes {
		nb := g.positions[node.ID()]
		if nb == nil {
			log.Printf(`Node flattened that hasn't been added: %+v`, node)
			continue
		}

		if !isSet {
			isSet = true
			toReturn = append(toReturn, node)
			lowest = nb.position
			continue
		}

		if nb.position == lowest {
			toReturn = append(toReturn, node)
		} else if nb.position < lowest {
			toReturn = toReturn[:0]
			toReturn = append(toReturn, node)
			lowest = nb.position
		}
	}

	return toReturn
}

// AddNodes will add the provided nodes to the flattened index
// of the graph and return an execution graph that is ready to
// be calculated.
func (g *Graph) AddNodes(dp IDependencyProvider, nodes Nodes) *ExecutionGraph {
	var dependentNodes Nodes
	// if there are no items in the graph then no node added
	// has a dependent in the graph
	if g.numItems > 0 {
		dependentNodes = dp.GetDependents(nodes)
	}

	highest := nodes.Highest()

	// want to make sure we don't overflow here
	if highest >= uint64(len(g.positions)) {
		diff := highest - uint64(len(g.positions))
		g.positions = append(g.positions, make(positions, diff+1)...)
	}

	var lowest int64
	if len(g.positions) > 0 && uint64(len(dependentNodes)) < g.numItems {
		dependencyNodes := g.findUniqueDependencies(dp, nodes)
		lowest = g.positions.highest(dependencyNodes) + 1
	}

	dependentNodes = append(dependentNodes, nodes...)
	flattened, circulars, size := flatten(dp, dependentNodes)
	g.insert(flattened, circulars, lowest)

	return &ExecutionGraph{
		size:      size,
		toApply:   flattened,
		circulars: circulars,
	}
}

func (g *Graph) findUniqueDependencies(dp IDependencyProvider, nodes Nodes) Nodes {
	if len(nodes) == 0 {
		return nil
	}

	ba := bitarray.NewSparseBitArray()
	for _, node := range nodes {
		ba.SetBit(node.ID())
	}
	ids := make([]uint64, 0, len(nodes))
	chunks := nodes.Split()
	var lock sync.Mutex
	var wg sync.WaitGroup
	wg.Add(len(chunks))
	for _, chunk := range chunks {
		go func(nodes Nodes) {
			for _, node := range nodes {
				deps := dp.GetDependencies(node)
				if deps == nil {
					continue
				}

				uints := deps.ToNums()
				lock.Lock()
				for _, id := range uints {
					if ok, _ := ba.GetBit(id); ok {
						continue
					}

					ids = append(ids, id)
					ba.SetBit(id)
				}
				lock.Unlock()
			}
			wg.Done()
		}(chunk)
	}

	wg.Wait()

	deps := make(Nodes, 0, len(ids))
	for _, id := range ids {
		nb := g.positions[id]
		if nb == nil {
			continue
		}

		deps = append(deps, nb.INode)
	}

	return deps
}

// RemoveNodes will remove the provided nodes from the graph.
func (g *Graph) RemoveNodes(dp IDependencyProvider, nodes Nodes) *ExecutionGraph {
	if len(nodes) == 0 {
		return &ExecutionGraph{}
	}

	lowest := g.positions.lowest(nodes)
	for _, node := range nodes {
		if g.positions[node.ID()] != nil {
			g.numItems--
		}
		g.positions[node.ID()] = nil // kill these first
	}

	maxLayer := g.maxLayer
	if maxLayer == -1 {
		maxLayer = 1
	}

	toExtract := make(map[int64]bool, maxLayer)
	if lowest == -1 { // we end up reordering the whole damn thing
		toExtract[-1] = true
		lowest = 0
	}

	for i := lowest; i <= g.maxLayer; i++ {
		toExtract[i] = true
	}

	nodesToFlatten := g.positions.extractLayers(toExtract)
	flattened, circulars, size := flatten(dp, nodesToFlatten)

	g.insert(flattened, circulars, lowest)
	// have to set new max layer, only in the delete case
	g.maxLayer = g.positions.highestSeen()

	return &ExecutionGraph{
		toApply:   flattened,
		circulars: circulars,
		size:      size,
	}
}

// Len returns the number of items in the graph.
func (g *Graph) Len() uint64 {
	return g.numItems
}

func (g *Graph) insert(flattened []Nodes, circulars Nodes, offset int64) {
	maxLayer := int64(-1)
	for i, nodes := range flattened {
		for _, node := range nodes {
			if g.positions[node.ID()] == nil {
				g.numItems++
			}
			g.positions[node.ID()] = &bundle{
				INode:    node,
				position: int64(i) + offset,
			}
		}
		maxLayer = int64(i) + offset
	}

	for _, node := range circulars {
		g.positions[node.ID()] = &bundle{
			INode: node, position: -1,
		}
	}

	if maxLayer > g.maxLayer {
		g.maxLayer = maxLayer
	}
}

// FromNodes will create a new graph from the given nodes.
func FromNodes(dp IDependencyProvider, nodes Nodes) *Graph {
	if len(nodes) == 0 {
		return &Graph{}
	}

	highest := nodes.Highest()
	positions := make(positions, highest+1)
	g := &Graph{
		positions: positions,
		maxLayer:  -1, // it's possible they are all circular
	}
	flattened, circulars, _ := flatten(dp, nodes)

	g.insert(flattened, circulars, 0)

	return g
}

func flatten(dp IDependencyProvider, nodes Nodes) ([]Nodes, Nodes, int64) {
	if len(nodes) == 0 {
		return nil, nil, 0
	}

	ba := bitarray.NewBitArray(dp.MaxNode(), true)
	remaining := make(Nodes, 0, len(nodes))
	helper := make(Nodes, len(nodes))
	for _, node := range nodes {
		if node != nil {
			ba.ClearBit(node.ID())
			remaining = append(remaining, node)
		}
	}

	var circulars Nodes
	results := make([]Nodes, 0, 5) // 5 is a guess
	layer := make([]INode, 0, len(nodes))
	size := int64(0)

	var wg sync.WaitGroup
	var lock sync.Mutex

	for {
		if len(remaining) == 0 && len(layer) == 0 {
			break
		}

		q := queue.New(int64(len(remaining)))
		for _, node := range remaining {
			q.Put(node)
		}

		completed := uint64(0)
		required := uint64(len(remaining))

		wg.Add(1)
		for i := 0; i < runtime.NumCPU(); i++ {
			go func() {
				for {
					items, err := q.Get(1)
					if err != nil { // queue was disposed
						break
					}

					node := items[0].(INode)
					okToAdd := true
					deps := dp.GetDependencies(node)
					if deps != nil && !ba.Intersects(deps) {
						okToAdd = false
					}

					if okToAdd {
						if node.IsCircular() {
							node.SetCircular(false)
						}
						lock.Lock()
						layer = append(layer, node)
						lock.Unlock()
					}

					if atomic.AddUint64(&completed, 1) == required {
						wg.Done()
						break
					}
				}
			}()
		}

		wg.Wait()
		q.Dispose()

		if len(layer) == 0 { // this is a problem, nothing new found
			circulars = make(Nodes, 0, len(remaining))
			for _, n := range remaining {
				n.SetCircular(true)
				circulars = append(circulars, n)
				size++
			}
			break
		}

		for _, node := range layer {
			size++
			ba.SetBit(node.ID())
		}

		results = append(results, layer)
		layer = make([]INode, 0, len(layer))
		copy(helper, remaining)
		remaining = remaining[:0]
		for i := 0; i < len(helper); i++ {
			node := helper[i]
			if node == nil {
				break
			}
			if ok, _ := ba.GetBit(node.ID()); ok {
				continue
			}

			remaining = append(remaining, node)
			helper[i] = nil
		}

		helper = helper[:len(remaining)]
	}

	return results, circulars, size
}

func New() *Graph {
	return &Graph{
		maxLayer: -1,
	}
}
<|endoftext|>


<commit_before>package henchman

import (
	"fmt"
	"io/ioutil"
	"strings"

	"gopkg.in/yaml.v2"
)

type InventoryConfig map[string]string
type InventoryInterface interface {
	Load(ic InventoryConfig) (Inventory, error)
}

//type Inventory map[string][]*Machine
type Inventory struct {
	//GroupHosts map[string][]*Machine
	Groups     map[string]HostGroup                   `yaml:"groups"`
	HostVars   map[string]map[interface{}]interface{} `yaml:"host_vars"`
	GlobalVars map[interface{}]interface{}            `yaml:"global_vars"`
}

type HostGroup struct {
	Hosts []string                    `yaml:"hosts"`
	Vars  map[interface{}]interface{} `yaml:"vars"`
}

func (inv Inventory) Count() int {
	count := 0
	for _, hostGroup := range inv.Groups {
		count += len(hostGroup.Hosts)
	}
	return count
}

// FIXME: Have a way to provide specifics
type YAMLInventory struct {
	Groups     map[string]HostGroup                   `yaml:"groups"`
	HostVars   map[string]map[interface{}]interface{} `yaml:"host_vars"`
	GlobalVars map[interface{}]interface{}            `yaml:"global_vars"`
}

func (yi *YAMLInventory) Load(ic InventoryConfig) (Inventory, error) {
	fname, present := ic["path"]
	if !present {
		return Inventory{}, fmt.Errorf("Missing 'path' in the config")
	}
	buf, err := ioutil.ReadFile(fname)
	if err != nil {
		return Inventory{}, HenchErr(err, map[string]interface{}{
			"file":     fname,
			"solution": "make sure directory exists, correct permissions, or is not corrupted",
		}, "While Reading File")
	}
	err = yaml.Unmarshal(buf, &yi)
	if err != nil {
		return Inventory{}, HenchErr(err, map[string]interface{}{
			"file":     fname,
			"solution": "Make sure inventory follows proper formatting.  Also check for tabs when there should be spaces.",
		}, "While unmarshalling inventory")
	}

	if yi.Groups == nil {
		return Inventory{}, HenchErr(fmt.Errorf("Groups field is required."), map[string]interface{}{
			"file":     fname,
			"solution": "Refer to the wiki for proper formatting.",
		}, "")
	}

	for key, val := range yi.Groups {
		if key == "hosts" {
			return Inventory{}, HenchErr(fmt.Errorf("'hosts' is not a valid group name"), map[string]interface{}{
				"file":     fname,
				"solution": "Change a group name away from hosts",
			}, "")
		}
		if val.Hosts == nil {
			return Inventory{}, HenchErr(fmt.Errorf("%v requires a hosts field.", key), map[string]interface{}{
				"file":     fname,
				"group":    key,
				"solution": "Refet to the wiki for proper formatting.",
			}, "")
		}
	}

	iv := &Inventory{}
	iv.HostVars = yi.HostVars
	iv.Groups = yi.Groups
	iv.GlobalVars = yi.GlobalVars
	return *iv, nil
}

//NOTE: Not being used.  Remove at 1/4/16
/*
func (inv *Inventory) MergeHostVars(hostname string, taskVars map[interface{}]interface{}) {
	if len(inv.HostVars) == 0 {
		return
	}
	if _, present := inv.HostVars[hostname]; present {
		MergeMap(inv.HostVars[hostname], taskVars, true)
	}
}
*/

/**
 * gets the groups under hosts section from plan file
 */
func GetInventoryGroups(planBuf []byte) ([]string, error) {
	hostsProxy := struct {
		Groups []string `yaml:"hosts"`
	}{}

	err := yaml.Unmarshal(planBuf, &hostsProxy)
	if err != nil {
		return nil, HenchErr(err, map[string]interface{}{
			"solution": "Check if hosts section exists",
		}, "While unmarshalling hosts section")
	}

	return hostsProxy.Groups, nil
}

/**
 * Returns a filtered inventory depending on which groups are specified
 * Keeps the groups specified in the filtered inventory
 */
func (inv *Inventory) GetInventoryForGroups(groups []string) Inventory {
	// FIXME: Support globbing in the groups
	// No groups? No problem. Just return the full inventory
	//	return fullInventory
	if len(groups) == 0 {
		return *inv
	} else {
		filtered := Inventory{}
		filtered.Groups = make(map[string]HostGroup)
		//filtered.HostVars = fullInventory.HostVars
		//log.Println(fullInventory)
		for _, group := range groups {
			machines, present := inv.Groups[group]
			if present {
				filtered.Groups[group] = machines
			}
		}
		filtered.HostVars = inv.HostVars
		filtered.GlobalVars = inv.GlobalVars
		return filtered
	}
}

func (inv *Inventory) GetMachines(tc TransportConfig) ([]*Machine, error) {
	var machines []*Machine
	machineSet := make(map[string]bool)
	for _, hostGroup := range inv.Groups {
		for _, hostname := range hostGroup.Hosts {
			if _, present := machineSet[hostname]; !present {
				machine := &Machine{}
				machineSet[hostname] = true
				machine.Hostname = hostname
				machine.Vars = hostGroup.Vars
				machines = append(machines, machine)
			} else {
				//machine part of multiple groups
				//update vars if any
				//latter group's vars overrides prev. groups vars
				for _, machine := range machines {
					if machine.Hostname == hostname {
						MergeMap(hostGroup.Vars, machine.Vars, true)
					}
					//get second groups henchman tc vars
				}
			}
		}
	}

	// gets henchman specific vars from global_vars
	globalInvHenchmanVars := GetHenchmanVars(inv.GlobalVars)

	// update hostvars
	for _, machine := range machines {
		for hostname, vars := range inv.HostVars {
			if machine.Hostname == hostname {
				MergeMap(vars, machine.Vars, true)
			}
		}
		// now open ssh connection for each machine
		tcCurr := make(TransportConfig)
		tcCurr["hostname"] = machine.Hostname
		for k, v := range tc {
			tcCurr[k] = v
		}

		// gets henchman specific vars from each machine and merges with global
		henchmanVars := GetHenchmanVars(machine.Vars)
		MergeMap(globalInvHenchmanVars, henchmanVars, false)

		for k, v := range henchmanVars {
			tcCurr[k.(string)] = v.(string)
		}

		Debug(map[string]interface{}{
			"host":   machine.Hostname,
			"config": tcCurr,
		}, "Transport Config for machine")

		// FIXME: This is frigging wrong
		// See #47
		ssht, err := NewSSH(&tcCurr)
		if err != nil {
			return nil, err
		}
		machine.Transport = ssht
	}

	return machines, nil
}

/**
 * Gets any vars with the prefix henchman
 */
func GetHenchmanVars(vars VarsMap) VarsMap {
	henchmanVars := VarsMap{}

	for k, v := range vars {
		if strings.Contains(k.(string), HENCHMAN_PREFIX) {
			parts := strings.Split(k.(string), HENCHMAN_PREFIX)
			henchmanVars[parts[len(parts)-1]] = v
		}
	}

	return henchmanVars
}
<commit_msg>Fixed a nil Vars map issue with repeat machines in multiple groups<commit_after>package henchman

import (
	"fmt"
	"io/ioutil"
	"strings"

	"gopkg.in/yaml.v2"
)

type InventoryConfig map[string]string
type InventoryInterface interface {
	Load(ic InventoryConfig) (Inventory, error)
}

//type Inventory map[string][]*Machine
type Inventory struct {
	//GroupHosts map[string][]*Machine
	Groups     map[string]HostGroup                   `yaml:"groups"`
	HostVars   map[string]map[interface{}]interface{} `yaml:"host_vars"`
	GlobalVars map[interface{}]interface{}            `yaml:"global_vars"`
}

type HostGroup struct {
	Hosts []string                    `yaml:"hosts"`
	Vars  map[interface{}]interface{} `yaml:"vars"`
}

func (inv Inventory) Count() int {
	count := 0
	for _, hostGroup := range inv.Groups {
		count += len(hostGroup.Hosts)
	}
	return count
}

// FIXME: Have a way to provide specifics
type YAMLInventory struct {
	Groups     map[string]HostGroup                   `yaml:"groups"`
	HostVars   map[string]map[interface{}]interface{} `yaml:"host_vars"`
	GlobalVars map[interface{}]interface{}            `yaml:"global_vars"`
}

func (yi *YAMLInventory) Load(ic InventoryConfig) (Inventory, error) {
	fname, present := ic["path"]
	if !present {
		return Inventory{}, fmt.Errorf("Missing 'path' in the config")
	}
	buf, err := ioutil.ReadFile(fname)
	if err != nil {
		return Inventory{}, HenchErr(err, map[string]interface{}{
			"file":     fname,
			"solution": "make sure directory exists, correct permissions, or is not corrupted",
		}, "While Reading File")
	}
	err = yaml.Unmarshal(buf, &yi)
	if err != nil {
		return Inventory{}, HenchErr(err, map[string]interface{}{
			"file":     fname,
			"solution": "Make sure inventory follows proper formatting.  Also check for tabs when there should be spaces.",
		}, "While unmarshalling inventory")
	}

	if yi.Groups == nil {
		return Inventory{}, HenchErr(fmt.Errorf("Groups field is required."), map[string]interface{}{
			"file":     fname,
			"solution": "Refer to the wiki for proper formatting.",
		}, "")
	}

	for key, val := range yi.Groups {
		if key == "hosts" {
			return Inventory{}, HenchErr(fmt.Errorf("'hosts' is not a valid group name"), map[string]interface{}{
				"file":     fname,
				"solution": "Change a group name away from hosts",
			}, "")
		}
		if val.Hosts == nil {
			return Inventory{}, HenchErr(fmt.Errorf("%v requires a hosts field.", key), map[string]interface{}{
				"file":     fname,
				"group":    key,
				"solution": "Refet to the wiki for proper formatting.",
			}, "")
		}
	}

	iv := &Inventory{}
	iv.HostVars = yi.HostVars
	iv.Groups = yi.Groups
	iv.GlobalVars = yi.GlobalVars
	return *iv, nil
}

//NOTE: Not being used.  Remove at 1/4/16
/*
func (inv *Inventory) MergeHostVars(hostname string, taskVars map[interface{}]interface{}) {
	if len(inv.HostVars) == 0 {
		return
	}
	if _, present := inv.HostVars[hostname]; present {
		MergeMap(inv.HostVars[hostname], taskVars, true)
	}
}
*/

/**
 * gets the groups under hosts section from plan file
 */
func GetInventoryGroups(planBuf []byte) ([]string, error) {
	hostsProxy := struct {
		Groups []string `yaml:"hosts"`
	}{}

	err := yaml.Unmarshal(planBuf, &hostsProxy)
	if err != nil {
		return nil, HenchErr(err, map[string]interface{}{
			"solution": "Check if hosts section exists",
		}, "While unmarshalling hosts section")
	}

	return hostsProxy.Groups, nil
}

/**
 * Returns a filtered inventory depending on which groups are specified
 * Keeps the groups specified in the filtered inventory
 */
func (inv *Inventory) GetInventoryForGroups(groups []string) Inventory {
	// FIXME: Support globbing in the groups
	// No groups? No problem. Just return the full inventory
	//	return fullInventory
	if len(groups) == 0 {
		return *inv
	} else {
		filtered := Inventory{}
		filtered.Groups = make(map[string]HostGroup)
		//filtered.HostVars = fullInventory.HostVars
		//log.Println(fullInventory)
		for _, group := range groups {
			machines, present := inv.Groups[group]
			if present {
				filtered.Groups[group] = machines
			}
		}
		filtered.HostVars = inv.HostVars
		filtered.GlobalVars = inv.GlobalVars
		return filtered
	}
}

func (inv *Inventory) GetMachines(tc TransportConfig) ([]*Machine, error) {
	var machines []*Machine
	machineSet := make(map[string]bool)
	for _, hostGroup := range inv.Groups {
		for _, hostname := range hostGroup.Hosts {
			if _, present := machineSet[hostname]; !present {
				machine := &Machine{}
				machineSet[hostname] = true
				machine.Hostname = hostname
				machine.Vars = hostGroup.Vars
				if machine.Vars == nil {
					machine.Vars = make(VarsMap)
				}
				machines = append(machines, machine)
			} else {
				//machine part of multiple groups
				//update vars if any
				//latter group's vars overrides prev. groups vars
				for _, machine := range machines {
					if machine.Hostname == hostname {
						MergeMap(hostGroup.Vars, machine.Vars, true)
					}
					//get second groups henchman tc vars
				}
			}
		}
	}

	// gets henchman specific vars from global_vars
	globalInvHenchmanVars := GetHenchmanVars(inv.GlobalVars)

	// update hostvars
	for _, machine := range machines {
		for hostname, vars := range inv.HostVars {
			if machine.Hostname == hostname {
				MergeMap(vars, machine.Vars, true)
			}
		}
		// now open ssh connection for each machine
		tcCurr := make(TransportConfig)
		tcCurr["hostname"] = machine.Hostname
		for k, v := range tc {
			tcCurr[k] = v
		}

		// gets henchman specific vars from each machine and merges with global
		henchmanVars := GetHenchmanVars(machine.Vars)
		MergeMap(globalInvHenchmanVars, henchmanVars, false)

		for k, v := range henchmanVars {
			tcCurr[k.(string)] = v.(string)
		}

		Debug(map[string]interface{}{
			"host":   machine.Hostname,
			"config": tcCurr,
		}, "Transport Config for machine")

		// FIXME: This is frigging wrong
		// See #47
		ssht, err := NewSSH(&tcCurr)
		if err != nil {
			return nil, err
		}
		machine.Transport = ssht
	}

	return machines, nil
}

/**
 * Gets any vars with the prefix henchman
 */
func GetHenchmanVars(vars VarsMap) VarsMap {
	henchmanVars := VarsMap{}

	for k, v := range vars {
		if strings.Contains(k.(string), HENCHMAN_PREFIX) {
			parts := strings.Split(k.(string), HENCHMAN_PREFIX)
			henchmanVars[parts[len(parts)-1]] = v
		}
	}

	return henchmanVars
}
<|endoftext|>


<commit_before>/*
Package hamt64 implements a 64 node wide Hashed Array Mapped Trie. The hash key
is 60 bits wide and broken into ten numbers of 6bits each. Those 6bit numbers
allows us to index into a 64 node array. Each node is either a leaf or another
64 node table. So the 60bit hash allows us to index into a B+ Tree with a
branching factor of 64 and a Maximum depth of 6.

The basic insertion operation is to calculate a 60 bit hash value from your key
(a string in the case you use hamt.StringKey), then split it into ten 6bit
 numbers. These ten numbers represent a path thru the tree. For each level we
use the coresponding number as an index into the 64 cell array. If the cell is
empty we create a  leaf node there. If the cell is occupide by another table
we continue walking up the tree. If the cell is occupide by a leaf we promote
that cell to a new table and put the current leaf and new one into that table
in cells corresponding to that new level. If we are at the maxDepth of tree
and there is already a leaf there we insert our key,value pair into that leaf.

The retrieval operation is a similar tree walk guided by the ten 6bit numbers
till we find a leaf with the key,value pair in it.

The deletion operation is a walk to find the key, then delete the key from the
leaf. An empty leaf is removed from it's table. If the table has only one other
leaf in that level we will remove that leaf, replace the table in it's parent
table placing that last leaf down one level.

*/
package hamt64

import (
	"fmt"
	"strings"

	"github.com/lleo/go-hamt/key"
)

// nBits constant is the number of bits(6) a 60bit hash value is split into,
// to provied the indexes of a HAMT.
const nBits uint = 6

// maxDepth constant is the maximum depth(9) of nBits values that constitute
// the path in a HAMT, from [0..maxDepth] for a total of maxDepth+1(10) levels.
// nBits*(maxDepth+1) == HASHBITS (ie 6*(9+1) == 60).
const maxDepth uint = 9

// tableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 1<<nBits (ie 2^6 == 64).
const tableCapacity uint = uint(1 << nBits)

// downgradeThreshold constant is the number of nodes a fullTable has shrunk to,
// before it is converted to a compressedTable.
const downgradeThreshold uint = tableCapacity / 8

// upgradeThreshold constan is the number of nodes a compressedTable has grown to,
// before it is converted to a fullTable.
const upgradeThreshold uint = tableCapacity / 2

func indexMask(depth uint) uint64 {
	return uint64(uint8(1<<nBits)-1) << (depth * nBits)
}

func index(h60 uint64, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h60 & idxMask) >> (depth * nBits))
	return idx
}

func hashPathString(hashPath uint64, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func hash60String(h60 uint64) string {
	return hashPathString(h60, maxDepth)
}

func hashPathMask(depth uint) uint64 {
	return uint64(1<<(depth*nBits)) - 1
}

func buildHashPath(hashPath uint64, idx, depth uint) uint64 {
	return hashPath | uint64(idx<<(depth*nBits))
}

// Configuration contants to be passed to `hamt64.New(int) *Hamt`.
const (
	// HybridTables indicates the structure should use compressedTable
	// initially, then upgrad to fullTable when appropriate.
	HybridTables = iota
	// CompTablesOnly indicates the structure should use compressedTables ONLY.
	// This was intended just save space, but also seems to be faster; CPU cache
	// locality maybe?
	CompTablesOnly
	// FullTableOnly indicates the structure should use fullTables ONLY.
	// This was intended to be for speed, as compressed tables use a software
	// bitCount function to access individual cells. Turns out, not so much.
	FullTablesOnly
)

// TableOptionName is a pedantic lookup table; given the configuration option
// it maps to the configuration option's name.
var TableOptionName = make(map[int]string, 3)

func init() {
	TableOptionName[0] = "HybridTables"
	TableOptionName[1] = "CompTablesOnly"
	TableOptionName[2] = "FullTablesOnly"
}

// Hamt is a Hashed Array Map Trie data structure. It has a branching factor of
// 64 and is at most 10 nodes deep. See:
// https://en.wikipedia.org/wiki/Hash_array_mapped_trie
type Hamt struct {
	root            tableI
	nentries        int
	grade, fullinit bool
}

// New creates a new hamt64.Hamt data structure with the table option set to
// either:
//
// `hamt64.HybridTables`:
// Initially start out with compressedTable, but when the table is half full
// upgrade to fullTable. If a fullTable shrinks to tableCapacity/8(4) entries
// downgrade to compressedTable.
//
// `hamt64.CompTablesOnly`:
// Use compressedTable ONLY with no up/downgrading to/from fullTable. This
// uses the least amount of space.
//
// `hamt64.FullTablesOnly`:
// Only use fullTable no up/downgrading from/to compressedTables. This is
// the fastest performance.
//
func New(opt int) *Hamt {
	var h = new(Hamt)
	if opt == CompTablesOnly {
		h.grade = false
		h.fullinit = false
	} else if opt == FullTablesOnly {
		h.grade = false
		h.fullinit = true
	} else /* opt == HybridTables */ {
		h.grade = true
		h.fullinit = false
	}
	return h
}

// IsEmpty returns a boolean indicating if this Hamt structure has no entries.
func (h *Hamt) IsEmpty() bool {
	return h.root == nil
}

// Get Hamt method looks up a given key in the Hamt data structure.
func (h *Hamt) Get(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h60 = k.Hash60()

	var curTable = h.root //ISA tableI

	for depth := uint(0); depth <= maxDepth; depth++ {
		var idx = index(h60, depth)
		var curNode = curTable.get(idx) //nodeI

		if curNode == nil {
			break
		}

		if leaf, isLeaf := curNode.(leafI); isLeaf {
			var val, found = leaf.get(k)
			return val, found
		}

		//else curNode MUST BE A tableI
		curTable = curNode.(tableI)
	}
	// curNode == nil || depth > maxDepth

	return nil, false
}

// Put Hamt method inserts a given key/val pair into the Hamt data structure.
// It returns a boolean indicating if the key/val was inserted or whether or
// not the key already existed and the val was merely overwritten.
func (h *Hamt) Put(k key.Key, v interface{}) bool {
	var depth uint
	var hashPath uint64

	if h.IsEmpty() {
		h.root = h.newRootTable(depth, hashPath, newFlatLeaf(k, v))
		h.nentries++
		return true
	}

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			curTable.set(idx, newFlatLeaf(k, v))
			h.nentries++

			// upgrade?
			if h.grade {
				_, isCompressedTable := curTable.(*compressedTable)
				if isCompressedTable && curTable.nentries() >= upgradeThreshold {
					curTable = upgradeToFullTable(hashPath, curTable.entries())
					if depth == 0 {
						h.root = curTable
					} else {
						parentTable := path.peek()
						parentIdx := index(k.Hash60(), depth-1)
						parentTable.set(parentIdx, curTable)
					}
				}
			}

			return true //inserted
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			if curLeaf.Hash60() == k.Hash60() {
				// This is a minor optimization but since these two leaves
				// will collide all the way up the to maxDepth, we can
				// choose to create the collisionLeaf hear and now.

				// Accumulate collisionLeaf
				newLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, newLeaf)
					h.nentries++
				}
				return inserted
			}

			if depth == maxDepth {
				// this test should be delete cuz it is logically impossible
				if curLeaf.Hash60() != k.Hash60() {
					// This should not happen cuz we had to walk up maxDepth
					// levels to get here.
					panic("WTF!!!")
				}

				// Accumulate collisionLeaf
				insLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, insLeaf)
					h.nentries++
				}
				return inserted
			}

			hashPath = buildHashPath(hashPath, idx, depth)
			var collisionTable = h.newTable(depth+1, hashPath, curLeaf, newFlatLeaf(k, v))
			curTable.set(idx, collisionTable)
			h.nentries++

			return true
		}

		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	}

	panic("WTF!")
}

// Del Hamt Method removes a given key from the Hamt data structure. It returns
// a pair of values: the value stored and a boolean indicating if the key was
// even found and deleted.
func (h *Hamt) Del(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h60 = k.Hash60()
	var depth uint
	var hashPath uint64

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(h60, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			return nil, false
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			val, delLeaf, deleted := curLeaf.del(k)
			if !deleted {
				return nil, false
			}
			// else a leaf key/value was deleted
			h.nentries--

			// If curLeaf was a collisionLeaf,
			// then delLeaf is either a slimmed down collisionLeaf or a flatLeaf.
			// If curLeaf was a flatLeaf then delLeaf is nil.
			curTable.set(idx, delLeaf)

			// downgrade?
			if h.grade {
				if delLeaf == nil {
					_, isFullTable := curTable.(*fullTable)
					if isFullTable && curTable.nentries() <= downgradeThreshold {
						curTable = downgradeToCompressedTable(hashPath, curTable.entries())
						if depth == 0 {
							h.root = curTable
						} else {
							parentTable := path.peek()
							parentIdx := index(h60, depth-1)
							parentTable.set(parentIdx, curTable)
						}
					}
				}
			}

			// If curTable has only one entry and that entry is a leafI,
			// then collapse that leafI down to the position curTable holds
			// in the parent Table; repeat test and collapse for parent table.
			//
			// These are identical for conditionals:
			//  curTable != h.root AND len(path) > 0 AND depth > 0
			//
			for curTable.nentries() == 1 && depth > 0 {
				// _ = ASSERT && Assert(curTable != h.root, "curTable == h.root")
				// _ = ASSERT && Assert(depth == len(path), "depth != len(path)")

				var node = curTable.entries()[0].node
				var leaf, isLeaf = node.(leafI)
				if !isLeaf {
					// We only collapse leafs
					break
				}

				// Collapse leaf down to where curTable is in parentTable

				var parentTable = path.pop()
				depth-- // OR depth = len(path)

				parentIdx := index(curTable.Hash60(), depth)
				parentTable.set(parentIdx, leaf)

				curTable = parentTable
			}

			// TODO: I should keep this table rather than throwing it away.
			// Instead using h.root == nil to detect emptyness, we should
			// trust our accounting and use h.nentries == 0.
			if curTable == h.root && curTable.nentries() == 0 {
				h.root = nil
			}

			return val, true
		} //if isLeaf

		// curNode is not nil
		// curNode is not a leafI
		// curNode MUST be a tableI
		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	} //for depth loop

	//log.Printf("Hamt.Del: WTF! this should never be called; k=%s", k)
	return nil, false
}

// String returns a string representation of the Hamt string.
func (h *Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root.LongString("", 0))
}

// LongString returns a complete listing of the entire Hamt data structure.
func (h *Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent, 0)
		str += indent + "} //Hamt"
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<commit_msg>Copied over the new hamt32/hamt.go code and used convert32to64.sh on it. See last commit for info on hamt32/hamt.go code change (added find() method).<commit_after>/*
Package hamt64 implements a 64 node wide Hashed Array Mapped Trie. The hash key
is 60 bits wide and broken into 9 numbers of 6 bits each. Those 6bit numbers
allows us to index into a 64 node array. Each node is either a leaf or another
64 node table. So the 60bit hash allows us to index into a B+ Tree with a
branching factor of 64 and a Maximum depth of 9.

The basic insertion operation is to calculate a 60 bit hash value from your key
(a string in the case you use hamt.StringKey), then split it into ten 6bit
 numbers. These ten numbers represent a path thru the tree. For each level we
use the coresponding number as an index into the 64 cell array. If the cell is
empty we create a  leaf node there. If the cell is occupide by another table
we continue walking up the tree. If the cell is occupide by a leaf we promote
that cell to a new table and put the current leaf and new one into that table
in cells corresponding to that new level. If we are at the maximun depth of
the  tree and there is already a leaf there we insert our key,value pair into
that leaf.

The retrieval operation is a similar tree walk guided by the ten 6bit numbers
till we find a leaf with the key,value pair in it.

The deletion operation is a walk to find the key, then delete the key from the
leaf. An empty leaf is removed from it's table. If the table has only one other
leaf in that level we will remove that leaf, replace the table in it's parent
table placing that last leaf down one level.

*/
package hamt64

import (
	"fmt"
	"log"
	"strings"

	"github.com/lleo/go-hamt/key"
)

// nBits constant is the number of bits(6) a 60bit hash value is split into
// to provied the index of a HAMT.
const nBits uint = 6

// maxDepth constant is the maximum depth(6) of nBits values that constitute
// the path in a HAMT, from [0..maxDepth] for a total of maxDepth+1(9) levels.
// nBits*(maxDepth+1) == HASHBITS (ie 6*(6+1) == 60).
const maxDepth uint = 6

// tableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 2^6 == 64.
const tableCapacity uint = uint(1 << nBits)

// downgradeThreshold constant is the number of nodes a fullTable has shrunk to,
// before it is converted to a compressedTable.
const downgradeThreshold uint = tableCapacity / 8

// upgradeThreshold constan is the number of nodes a compressedTable has grown to,
// before it is converted to a fullTable.
const upgradeThreshold uint = tableCapacity / 2

func indexMask(depth uint) uint64 {
	return uint64(uint8(1<<nBits)-1) << (depth * nBits)
}

func index(h60 uint64, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h60 & idxMask) >> (depth * nBits))
	return idx
}

func hashPathString(hashPath uint64, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func hash60String(h60 uint64) string {
	return hashPathString(h60, maxDepth)
}

func hashPathMask(depth uint) uint64 {
	return uint64(1<<(depth*nBits)) - 1
}

func buildHashPath(hashPath uint64, idx, depth uint) uint64 {
	return hashPath | uint64(idx<<(depth*nBits))
}

// Configuration contants to be passed to `hamt64.New(int) *Hamt`.
const (
	// HybridTables indicates the structure should use compressedTable
	// initially, then upgrad to fullTable when appropriate.
	HybridTables = iota //0
	// CompTablesOnly indicates the structure should use compressedTables ONLY.
	// This was intended just save space, but also seems to be faster; CPU cache
	// locality maybe?
	CompTablesOnly //1
	// FullTableOnly indicates the structure should use fullTables ONLY.
	// This was intended to be for speed, as compressed tables use a software
	// bitCount function to access individual cells. Turns out, not so much.
	FullTablesOnly //2
)

// TableOptionName is a map of the table option value Hybrid, CompTablesOnly,
// or FullTableOnly to a string representing that option.
//      var options = hamt64.FullTablesOnly
//      hamt64.TableOptionName[hamt64.FullTablesOnly] == "FullTablesOnly"
var TableOptionName = make(map[int]string, 3)

func init() {
	TableOptionName[HybridTables] = "HybridTables"
	TableOptionName[CompTablesOnly] = "CompTablesOnly"
	TableOptionName[FullTablesOnly] = "FullTablesOnly"
}

//Hamt is a Hashed Array Map Trie data structure. It has a branching factor of
//64 and is at most 9 nodes deep. See:
//https://en.wikipedia.org/wiki/Hash_array_mapped_trie
type Hamt struct {
	root            tableI
	nentries        int
	grade, fullinit bool
}

//New creates a new hamt64.Hamt data structure with the table option set to
//either:
//
//`hamt64.HybridTables`:
//Initially start out with compressedTable, but when the table is half full
//upgrade to fullTable. If a fullTable shrinks to tableCapacity/8(4) entries
//downgrade to compressedTable.
//
//`hamt64.CompTablesOnly`:
//Use compressedTable ONLY with no up/downgrading to/from fullTable. This
//uses the least amount of space.
//
//`hamt64.FullTablesOnly`:
//Only use fullTable no up/downgrading from/to compressedTables. This is
//the fastest performance.
func New(opt int) *Hamt {
	var h = new(Hamt)
	if opt == CompTablesOnly {
		h.grade = false
		h.fullinit = false
	} else if opt == FullTablesOnly {
		h.grade = false
		h.fullinit = true
	} else /* opt == HybridTables */ {
		h.grade = true
		h.fullinit = false
	}
	return h
}

// IsEmpty Hamt method returns a boolean indicating if this Hamt structure has
// no entries.
func (h *Hamt) IsEmpty() bool {
	return h.root == nil
}

func (h *Hamt) Nentries() int {
	return h.nentries
}

func (h Hamt) find(k key.Key) (path pathT, leaf leafI, idx uint) {
	if h.IsEmpty() {
		return nil, nil, 0
	}

	path = newPathT()
	var curTable = h.root

	var h60 = k.Hash60()
	var depth uint
	var curNode nodeI

DepthIter:
	for depth = 0; depth <= maxDepth; depth++ {
		path.push(curTable)
		idx = index(h60, depth)
		curNode = curTable.get(idx)

		switch n := curNode.(type) {
		case nil:
			leaf = nil
			break DepthIter
		case leafI:
			leaf = n
			break DepthIter
		case tableI:
			if depth == maxDepth {
				log.Panicf("SHOULD NOT BE REACHED; depth,%d == maxDepth,%d & tableI entry found; %s", depth, maxDepth, n)
			}
			curTable = n
			// exit switch then loop for
		default:
			log.Panicf("SHOULD NOT BE REACHED: depth=%d; curNode unknown type=%T;", depth, curNode)
		}
	}

	return
}

// Get Hamt method looks up a given key in the Hamt data structure.
// BenchHamt64:
//func (h *Hamt) Get(k key.Key) (val interface{}, found bool) {
//	var _, leaf, _ = h.find(k)
//
//	if leaf == nil {
//		return //nil, false
//	}
//
//	val, found = leaf.get(k)
//	return
//}

func (h *Hamt) Get(k key.Key) (val interface{}, found bool) {
	if h.IsEmpty() {
		return //nil, false
	}

	var h60 = k.Hash60()

	var curTable = h.root //ISA tableI

	for depth := uint(0); depth <= maxDepth; depth++ {
		var idx = index(h60, depth)
		var curNode = curTable.get(idx) //nodeI

		if curNode == nil {
			return //nil, false
		}

		if leaf, isLeaf := curNode.(leafI); isLeaf {
			val, found = leaf.get(k)
			return //val, found
		}

		if depth == maxDepth {
			panic("SHOULD NOT HAPPEN")
		}
		curTable = curNode.(tableI)
	}

	panic("SHOULD NEVER BE REACHED")
}

// Put Hamt method inserts a given key/val pair into the Hamt data structure.
// It returns a boolean indicating if the key/val was inserted or whether or
// not the key already existed and the val was merely overwritten.
func (h *Hamt) Put(k key.Key, v interface{}) bool {
	var depth uint
	var hashPath uint64

	if h.IsEmpty() {
		h.root = h.newRootTable(depth, hashPath, newFlatLeaf(k, v))
		h.nentries++
		return true
	}

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			curTable.set(idx, newFlatLeaf(k, v))
			h.nentries++

			// upgrade?
			if h.grade {
				_, isCompressedTable := curTable.(*compressedTable)
				if isCompressedTable && curTable.nentries() >= upgradeThreshold {
					curTable = upgradeToFullTable(hashPath, curTable.entries())
					if depth == 0 {
						h.root = curTable
					} else {
						parentTable := path.peek()
						parentIdx := index(k.Hash60(), depth-1)
						parentTable.set(parentIdx, curTable)
					}
				}
			}

			return true //inserted
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			if curLeaf.Hash60() == k.Hash60() {
				// This is a minor optimization but since these two leaves
				// will collide all the way up the to maxDepth, we can
				// choose to create the collisionLeaf hear and now.

				// Accumulate collisionLeaf
				colLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, colLeaf)
					h.nentries++
				}
				return inserted
			}

			if depth == maxDepth {
				// this test should be delete cuz it is logically impossible
				if curLeaf.Hash60() != k.Hash60() {
					// This should not happen cuz we had to walk up maxDepth
					// levels to get here.
					panic("WTF!!!")
				}

				// Accumulate collisionLeaf
				insLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, insLeaf)
					h.nentries++
				}
				return inserted
			}

			hashPath = buildHashPath(hashPath, idx, depth)
			var collisionTable = h.newTable(depth+1, hashPath, curLeaf, newFlatLeaf(k, v))
			curTable.set(idx, collisionTable)
			h.nentries++

			return true
		}

		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	}

	//log.Println(path)
	//log.Printf("k=%s, v=%v", k, v)

	panic("WTF!!")
}

// Del Hamt Method removes a given key from the Hamt data structure. It returns
// a pair of values: the value stored and a boolean indicating if the key was
// even found and deleted.
func (h *Hamt) Del(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h60 = k.Hash60()
	var depth uint
	var hashPath uint64

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(h60, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			return nil, false
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			val, delLeaf, deleted := curLeaf.del(k)
			if !deleted {
				return nil, false
			}
			// else a leaf key/value was deleted
			h.nentries--

			// If curLeaf was a collisionLeaf,
			// then delLeaf is either a slimmed down collisionLeaf or a flatLeaf.
			// If curLeaf was a flatLeaf then delLeaf is nil.
			curTable.set(idx, delLeaf)

			// downgrade?
			if h.grade {
				if delLeaf == nil {
					_, isFullTable := curTable.(*fullTable)
					if isFullTable && curTable.nentries() <= downgradeThreshold {
						curTable = downgradeToCompressedTable(hashPath, curTable.entries())
						if depth == 0 {
							h.root = curTable
						} else {
							parentTable := path.peek()
							parentIdx := index(h60, depth-1)
							parentTable.set(parentIdx, curTable)
						}
					}
				}
			}

			// If curTable has only one entry and that entry is a leafI,
			// then collapse that leafI down to the position curTable holds
			// in the parent Table; repeat test and collapse for parent table.
			//
			// These are identical for conditionals:
			//  curTable != h.root AND len(path) > 0 AND depth > 0
			//
			for curTable.nentries() == 1 && depth > 0 {
				// _ = ASSERT && Assert(curTable != h.root, "curTable == h.root")
				// _ = ASSERT && Assert(depth == len(path), "depth != len(path)")

				var node = curTable.entries()[0].node
				var leaf, isLeaf = node.(leafI)
				if !isLeaf {
					// We only collapse leafs
					break
				}

				// Collapse leaf down to where curTable is in parentTable

				var parentTable = path.pop()
				depth-- // OR depth = len(path)

				parentIdx := index(curTable.Hash60(), depth)
				parentTable.set(parentIdx, leaf)

				curTable = parentTable
			}

			// TODO: I should keep this table rather than throwing it away.
			// Instead using h.root == nil to detect emptyness, we should
			// trust our accounting and use h.nentries == 0.
			if curTable == h.root && curTable.nentries() == 0 {
				h.root = nil
			}

			return val, true
		} //if isLeaf

		// curNode is not nil
		// curNode is not a leafI
		// curNode MUST be a tableI
		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	} //for depth loop

	//log.Printf("WTF! this should never be called; k=%s", k)
	return nil, false
}

// String returns a string representation of the Hamt string.
func (h *Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root.LongString("", 0))
}

// LongString returns a complete listing of the entire Hamt data structure.
func (h *Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent, 0)
		str += indent + "} //Hamt"
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"
)

type rCommand struct {
	command   *exec.Cmd
	output    string
	startTime int64
	endTime   int64
	err       error
	mainOut   string
	status    string
	crash1    bool
	crash2    bool
	base      string
}

//Scheduler the main task scheduler
type Scheduler struct {
	commands []*rCommand
	cMutex   *sync.Mutex
	rMutex   *sync.Mutex
	rMap     map[string]*rCommand
	Log      func(string)
	lastRun  string
}

func (s *Scheduler) clean() {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	for key, command := range s.rMap {
		if command.endTime > 0 && time.Now().Sub(time.Unix(command.endTime, 0)) > time.Minute*5 {
			delete(s.rMap, key)
		}
	}
}

func (s *Scheduler) getState(key string) string {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if _, ok := s.rMap[key]; ok {
		return fmt.Sprintf("%v -> %v", s.rMap[key].endTime, s.rMap[key].output)
	}

	return "UNKNOWN"
}

func (s *Scheduler) markComplete(key string) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		val.endTime = time.Now().Unix()
		val.output = key
	} else {
		s.rMap[key] = &rCommand{endTime: time.Now().Unix(), output: key, base: key}
	}
}

// Schedule schedules a task
func (s *Scheduler) Schedule(c *rCommand) string {
	key := fmt.Sprintf("%v-%v", time.Now().Nanosecond(), c.command.Path)
	s.commands = append(s.commands, c)
	s.rMutex.Lock()
	s.rMap[key] = c
	s.rMutex.Unlock()
	s.processCommands()
	return key
}

func (s *Scheduler) getOutput(key string) (string, error) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		return val.output, nil
	}

	return "", fmt.Errorf("KEY NOT_IN_MAP: %v", key)
}

func (s *Scheduler) getErrOutput(key string) (string, error) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		return val.mainOut, nil
	}

	return "", fmt.Errorf("KEY NOT_IN_MAP: %v", key)
}

func (s *Scheduler) getStatus(key string) string {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		return val.status
	}

	return fmt.Sprintf("KEY NOT_IN_MAP: %v", key)
}

func (s *Scheduler) killJob(key string) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		if val.command.Process != nil {
			val.command.Process.Kill()
			val.command.Process.Wait()
		}
	}
}

func (s *Scheduler) removeJob(key string) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	delete(s.rMap, key)
}

func (s *Scheduler) schedulerComplete(key string) bool {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		//Validate that this in the process queue
		if val.status == "InQueue" && len(s.commands) == 0 {
			s.commands = append(s.commands, val)
			s.processCommands()
		}

		return val.endTime > 0
	}

	// Default to true if the key is not found
	return true
}

func (s *Scheduler) processCommands() {
	s.cMutex.Lock()
	if len(s.commands) > 0 {
		c := s.commands[0]
		s.commands = s.commands[1:]
		s.lastRun = c.command.Path + " -> " + fmt.Sprintf("%v", c.command.Args)
		err := run(c)
		if err != nil {
			c.endTime = time.Now().Unix()
		}
	}
	s.cMutex.Unlock()
}

func run(c *rCommand) error {
	c.status = "Running"
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}
	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	pathbin := fmt.Sprintf("GOBIN=" + home + "/gobuild/bin")
	found := false
	for i, blah := range env {
		if strings.HasPrefix(blah, "GOPATH") {
			env[i] = path
			found = true
		}
		if strings.HasPrefix(blah, "GOBIN") {
			env[i] = pathbin
			found = true
		}
	}
	if !found {
		env = append(env, path)
	}
	c.command.Env = env

	out, err1 := c.command.StderrPipe()
	outr, err2 := c.command.StdoutPipe()

	if c.crash1 || err1 != nil {
		return err1
	}

	if c.crash2 || err2 != nil {
		return err2
	}

	scanner := bufio.NewScanner(out)
	go func() {
		for scanner != nil && scanner.Scan() {
			c.output += scanner.Text()
		}
		out.Close()
	}()

	scanner2 := bufio.NewScanner(outr)
	go func() {
		for scanner2 != nil && scanner2.Scan() {
			c.mainOut += scanner2.Text()
		}
		outr.Close()
	}()

	c.status = "StartCommand"
	err := c.command.Start()
	if err != nil {
		return err
	}
	c.startTime = time.Now().Unix()

	// Monitor the job and report completion
	go func() {
		c.status = "Entering Wait"
		err := c.command.Wait()
		c.status = "Completed Wait"
		if err != nil {
			c.err = err
		}
		c.endTime = time.Now().Unix()
	}()

	return nil
}
<commit_msg>Also this<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"
)

type rCommand struct {
	command   *exec.Cmd
	output    string
	startTime int64
	endTime   int64
	err       error
	mainOut   string
	status    string
	crash1    bool
	crash2    bool
	base      string
}

//Scheduler the main task scheduler
type Scheduler struct {
	commands []*rCommand
	cMutex   *sync.Mutex
	rMutex   *sync.Mutex
	rMap     map[string]*rCommand
	Log      func(string)
	lastRun  string
}

func (s *Scheduler) clean() {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	for key, command := range s.rMap {
		if command.endTime > 0 && time.Now().Sub(time.Unix(command.endTime, 0)) > time.Minute*5 {
			delete(s.rMap, key)
		}
	}
}

func (s *Scheduler) getState(key string) string {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if _, ok := s.rMap[key]; ok {
		return fmt.Sprintf("%v -> %v", s.rMap[key].endTime, s.rMap[key].output)
	}

	return "UNKNOWN"
}

func (s *Scheduler) markComplete(key string) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		val.endTime = time.Now().Unix()
		val.output = key
	} else {
		s.rMap[key] = &rCommand{endTime: time.Now().Unix(), output: key, base: key}
	}
}

// Schedule schedules a task
func (s *Scheduler) Schedule(c *rCommand) string {
	key := fmt.Sprintf("%v-%v", time.Now().Nanosecond(), c.command.Path)
	s.commands = append(s.commands, c)
	c.status = "InQueue"
	s.rMutex.Lock()
	s.rMap[key] = c
	s.rMutex.Unlock()
	s.processCommands()
	return key
}

func (s *Scheduler) getOutput(key string) (string, error) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		return val.output, nil
	}

	return "", fmt.Errorf("KEY NOT_IN_MAP: %v", key)
}

func (s *Scheduler) getErrOutput(key string) (string, error) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		return val.mainOut, nil
	}

	return "", fmt.Errorf("KEY NOT_IN_MAP: %v", key)
}

func (s *Scheduler) getStatus(key string) string {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		return val.status
	}

	return fmt.Sprintf("KEY NOT_IN_MAP: %v", key)
}

func (s *Scheduler) killJob(key string) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		if val.command.Process != nil {
			val.command.Process.Kill()
			val.command.Process.Wait()
		}
	}
}

func (s *Scheduler) removeJob(key string) {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	delete(s.rMap, key)
}

func (s *Scheduler) schedulerComplete(key string) bool {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()
	if val, ok := s.rMap[key]; ok {
		//Validate that this in the process queue
		if val.status == "InQueue" && len(s.commands) == 0 {
			s.commands = append(s.commands, val)
			s.processCommands()
		}

		return val.endTime > 0
	}

	// Default to true if the key is not found
	return true
}

func (s *Scheduler) processCommands() {
	s.cMutex.Lock()
	if len(s.commands) > 0 {
		c := s.commands[0]
		s.commands = s.commands[1:]
		s.lastRun = c.command.Path + " -> " + fmt.Sprintf("%v", c.command.Args)
		err := run(c)
		if err != nil {
			c.endTime = time.Now().Unix()
		}
	}
	s.cMutex.Unlock()
}

func run(c *rCommand) error {
	c.status = "Running"
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}
	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	pathbin := fmt.Sprintf("GOBIN=" + home + "/gobuild/bin")
	found := false
	for i, blah := range env {
		if strings.HasPrefix(blah, "GOPATH") {
			env[i] = path
			found = true
		}
		if strings.HasPrefix(blah, "GOBIN") {
			env[i] = pathbin
			found = true
		}
	}
	if !found {
		env = append(env, path)
	}
	c.command.Env = env

	out, err1 := c.command.StderrPipe()
	outr, err2 := c.command.StdoutPipe()

	if c.crash1 || err1 != nil {
		return err1
	}

	if c.crash2 || err2 != nil {
		return err2
	}

	scanner := bufio.NewScanner(out)
	go func() {
		for scanner != nil && scanner.Scan() {
			c.output += scanner.Text()
		}
		out.Close()
	}()

	scanner2 := bufio.NewScanner(outr)
	go func() {
		for scanner2 != nil && scanner2.Scan() {
			c.mainOut += scanner2.Text()
		}
		outr.Close()
	}()

	c.status = "StartCommand"
	err := c.command.Start()
	if err != nil {
		return err
	}
	c.startTime = time.Now().Unix()

	// Monitor the job and report completion
	go func() {
		c.status = "Entering Wait"
		err := c.command.Wait()
		c.status = "Completed Wait"
		if err != nil {
			c.err = err
		}
		c.endTime = time.Now().Unix()
	}()

	return nil
}
<|endoftext|>


<commit_before>package native

import (
	"errors"
	"time"

	"github.com/CyCoreSystems/ari"
	uuid "github.com/satori/go.uuid"
)

// Bridge provides the ARI Bridge accessors for the native client
type Bridge struct {
	client *Client
}

// Create creates a bridge and returns the lazy handle for the bridge
func (b *Bridge) Create(key *ari.Key, t string, name string) (bh *ari.BridgeHandle, err error) {
	bh, err = b.StageCreate(key, t, name)
	if err != nil {
		return nil, err
	}
	return bh, bh.Exec()
}

// StageCreate creates a new bridge handle, staged with a bridge `Create` operation.
func (b *Bridge) StageCreate(key *ari.Key, btype, name string) (*ari.BridgeHandle, error) {
	if key.ID == "" {
		key.ID = uuid.NewV1().String()
	}

	req := struct {
		ID   string `json:"bridgeId,omitempty"`
		Type string `json:"type,omitempty"`
		Name string `json:"name,omitempty"`
	}{
		ID:   key.ID,
		Type: btype,
		Name: name,
	}

	return ari.NewBridgeHandle(b.client.stamp(key), b, func(bh *ari.BridgeHandle) (err error) {
		return b.client.post("/bridges/"+key.ID, &req, nil)
	}), nil
}

// Get gets the lazy handle for the given bridge id
func (b *Bridge) Get(key *ari.Key) *ari.BridgeHandle {
	return ari.NewBridgeHandle(b.client.stamp(key), b, nil)
}

// List lists the current bridges and returns a list of lazy handles
func (b *Bridge) List(filter *ari.Key) (bx []*ari.Key, err error) {
	// native client ignores filter

	var bridges = []struct {
		ID string `json:"id"`
	}{}

	err = b.client.get("/bridges", &bridges)
	for _, i := range bridges {
		k := b.client.stamp(ari.NewKey(ari.BridgeKey, i.ID))
		if filter.Match(k) {
			bx = append(bx, k)
		}
	}
	return
}

// Data returns the details of a bridge
// Equivalent to Get /bridges/{bridgeId}
func (b *Bridge) Data(key *ari.Key) (*ari.BridgeData, error) {
	if key == nil || key.ID == "" {
		return nil, errors.New("bridge key not supplied")
	}

	var data = new(ari.BridgeData)
	if err := b.client.get("/bridges/"+key.ID, data); err != nil {
		return nil, dataGetError(err, "bridge", "%v", key.ID)
	}

	data.Key = b.client.stamp(key)

	return data, nil
}

// AddChannel adds a channel to a bridge
// Equivalent to Post /bridges/{id}/addChannel
func (b *Bridge) AddChannel(key *ari.Key, channelID string) (err error) {
	id := key.ID

	type request struct {
		ChannelID string `json:"channel"`
		Role      string `json:"role,omitempty"`
	}

	req := request{channelID, ""}
	err = b.client.post("/bridges/"+id+"/addChannel", nil, &req)
	return
}

// RemoveChannel removes the specified channel from a bridge
// Equivalent to Post /bridges/{id}/removeChannel
func (b *Bridge) RemoveChannel(key *ari.Key, channelID string) (err error) {
	id := key.ID

	req := struct {
		ChannelID string `json:"channel"`
	}{
		ChannelID: channelID,
	}

	//pass request
	err = b.client.post("/bridges/"+id+"/removeChannel", nil, &req)
	return
}

// Delete shuts down a bridge. If any channels are in this bridge,
// they will be removed and resume whatever they were doing beforehand.
// This means that the channels themselves are not deleted.
// Equivalent to DELETE /bridges/{id}
func (b *Bridge) Delete(key *ari.Key) (err error) {
	id := key.ID
	err = b.client.del("/bridges/"+id, nil, "")
	return
}

// Play attempts to play the given mediaURI on the bridge, using the playbackID
// as the identifier to the created playback handle
func (b *Bridge) Play(key *ari.Key, playbackID string, mediaURI string) (*ari.PlaybackHandle, error) {
	h, err := b.StagePlay(key, playbackID, mediaURI)
	if err != nil {
		return nil, err
	}
	return h, h.Exec()
}

// StagePlay stages a `Play` operation on the bridge
func (b *Bridge) StagePlay(key *ari.Key, playbackID string, mediaURI string) (*ari.PlaybackHandle, error) {
	if playbackID == "" {
		playbackID = uuid.NewV1().String()
	}

	resp := make(map[string]interface{})
	req := struct {
		Media string `json:"media"`
	}{
		Media: mediaURI,
	}
	playbackKey := b.client.stamp(ari.NewKey(ari.PlaybackKey, playbackID))

	return ari.NewPlaybackHandle(playbackKey, b.client.Playback(), func(h *ari.PlaybackHandle) error {
		return b.client.post("/bridges/"+key.ID+"/play/"+playbackID, &resp, &req)
	}), nil
}

// Record attempts to record audio on the bridge, using name as the identifier for
// the created live recording handle
func (b *Bridge) Record(key *ari.Key, name string, opts *ari.RecordingOptions) (*ari.LiveRecordingHandle, error) {
	h, err := b.StageRecord(key, name, opts)
	if err != nil {
		return nil, err
	}
	return h, h.Exec()
}

// StageRecord stages a `Record` opreation
func (b *Bridge) StageRecord(key *ari.Key, name string, opts *ari.RecordingOptions) (*ari.LiveRecordingHandle, error) {
	if opts == nil {
		opts = &ari.RecordingOptions{}
	}

	resp := make(map[string]interface{})
	req := struct {
		Name        string `json:"name"`
		Format      string `json:"format"`
		MaxDuration int    `json:"maxDurationSeconds"`
		MaxSilence  int    `json:"maxSilenceSeconds"`
		IfExists    string `json:"ifExists,omitempty"`
		Beep        bool   `json:"beep"`
		TerminateOn string `json:"terminateOn,omitempty"`
	}{
		Name:        name,
		Format:      opts.Format,
		MaxDuration: int(opts.MaxDuration / time.Second),
		MaxSilence:  int(opts.MaxSilence / time.Second),
		IfExists:    opts.Exists,
		Beep:        opts.Beep,
		TerminateOn: opts.Terminate,
	}

	recordingKey := b.client.stamp(ari.NewKey(ari.LiveRecordingKey, name))

	return ari.NewLiveRecordingHandle(recordingKey, b.client.LiveRecording(), func(h *ari.LiveRecordingHandle) error {
		return b.client.post("/bridges/"+key.ID+"/record", &resp, &req)
	}), nil
}

// Subscribe creates an event subscription for events related to the given
// bridge␃entity
func (b *Bridge) Subscribe(key *ari.Key, n ...string) ari.Subscription {
	return b.client.Bus().Subscribe(key, n...)
}
<commit_msg>Fix parameter order for bridge.StageCreate().<commit_after>package native

import (
	"errors"
	"time"

	"github.com/CyCoreSystems/ari"
	uuid "github.com/satori/go.uuid"
)

// Bridge provides the ARI Bridge accessors for the native client
type Bridge struct {
	client *Client
}

// Create creates a bridge and returns the lazy handle for the bridge
func (b *Bridge) Create(key *ari.Key, t string, name string) (bh *ari.BridgeHandle, err error) {
	bh, err = b.StageCreate(key, t, name)
	if err != nil {
		return nil, err
	}
	return bh, bh.Exec()
}

// StageCreate creates a new bridge handle, staged with a bridge `Create` operation.
func (b *Bridge) StageCreate(key *ari.Key, btype, name string) (*ari.BridgeHandle, error) {
	if key.ID == "" {
		key.ID = uuid.NewV1().String()
	}

	req := struct {
		ID   string `json:"bridgeId,omitempty"`
		Type string `json:"type,omitempty"`
		Name string `json:"name,omitempty"`
	}{
		ID:   key.ID,
		Type: btype,
		Name: name,
	}

	return ari.NewBridgeHandle(b.client.stamp(key), b, func(bh *ari.BridgeHandle) (err error) {
		return b.client.post("/bridges/"+key.ID, nil, &req)
	}), nil
}

// Get gets the lazy handle for the given bridge id
func (b *Bridge) Get(key *ari.Key) *ari.BridgeHandle {
	return ari.NewBridgeHandle(b.client.stamp(key), b, nil)
}

// List lists the current bridges and returns a list of lazy handles
func (b *Bridge) List(filter *ari.Key) (bx []*ari.Key, err error) {
	// native client ignores filter

	var bridges = []struct {
		ID string `json:"id"`
	}{}

	err = b.client.get("/bridges", &bridges)
	for _, i := range bridges {
		k := b.client.stamp(ari.NewKey(ari.BridgeKey, i.ID))
		if filter.Match(k) {
			bx = append(bx, k)
		}
	}
	return
}

// Data returns the details of a bridge
// Equivalent to Get /bridges/{bridgeId}
func (b *Bridge) Data(key *ari.Key) (*ari.BridgeData, error) {
	if key == nil || key.ID == "" {
		return nil, errors.New("bridge key not supplied")
	}

	var data = new(ari.BridgeData)
	if err := b.client.get("/bridges/"+key.ID, data); err != nil {
		return nil, dataGetError(err, "bridge", "%v", key.ID)
	}

	data.Key = b.client.stamp(key)

	return data, nil
}

// AddChannel adds a channel to a bridge
// Equivalent to Post /bridges/{id}/addChannel
func (b *Bridge) AddChannel(key *ari.Key, channelID string) (err error) {
	id := key.ID

	type request struct {
		ChannelID string `json:"channel"`
		Role      string `json:"role,omitempty"`
	}

	req := request{channelID, ""}
	err = b.client.post("/bridges/"+id+"/addChannel", nil, &req)
	return
}

// RemoveChannel removes the specified channel from a bridge
// Equivalent to Post /bridges/{id}/removeChannel
func (b *Bridge) RemoveChannel(key *ari.Key, channelID string) (err error) {
	id := key.ID

	req := struct {
		ChannelID string `json:"channel"`
	}{
		ChannelID: channelID,
	}

	//pass request
	err = b.client.post("/bridges/"+id+"/removeChannel", nil, &req)
	return
}

// Delete shuts down a bridge. If any channels are in this bridge,
// they will be removed and resume whatever they were doing beforehand.
// This means that the channels themselves are not deleted.
// Equivalent to DELETE /bridges/{id}
func (b *Bridge) Delete(key *ari.Key) (err error) {
	id := key.ID
	err = b.client.del("/bridges/"+id, nil, "")
	return
}

// Play attempts to play the given mediaURI on the bridge, using the playbackID
// as the identifier to the created playback handle
func (b *Bridge) Play(key *ari.Key, playbackID string, mediaURI string) (*ari.PlaybackHandle, error) {
	h, err := b.StagePlay(key, playbackID, mediaURI)
	if err != nil {
		return nil, err
	}
	return h, h.Exec()
}

// StagePlay stages a `Play` operation on the bridge
func (b *Bridge) StagePlay(key *ari.Key, playbackID string, mediaURI string) (*ari.PlaybackHandle, error) {
	if playbackID == "" {
		playbackID = uuid.NewV1().String()
	}

	resp := make(map[string]interface{})
	req := struct {
		Media string `json:"media"`
	}{
		Media: mediaURI,
	}
	playbackKey := b.client.stamp(ari.NewKey(ari.PlaybackKey, playbackID))

	return ari.NewPlaybackHandle(playbackKey, b.client.Playback(), func(h *ari.PlaybackHandle) error {
		return b.client.post("/bridges/"+key.ID+"/play/"+playbackID, &resp, &req)
	}), nil
}

// Record attempts to record audio on the bridge, using name as the identifier for
// the created live recording handle
func (b *Bridge) Record(key *ari.Key, name string, opts *ari.RecordingOptions) (*ari.LiveRecordingHandle, error) {
	h, err := b.StageRecord(key, name, opts)
	if err != nil {
		return nil, err
	}
	return h, h.Exec()
}

// StageRecord stages a `Record` opreation
func (b *Bridge) StageRecord(key *ari.Key, name string, opts *ari.RecordingOptions) (*ari.LiveRecordingHandle, error) {
	if opts == nil {
		opts = &ari.RecordingOptions{}
	}

	resp := make(map[string]interface{})
	req := struct {
		Name        string `json:"name"`
		Format      string `json:"format"`
		MaxDuration int    `json:"maxDurationSeconds"`
		MaxSilence  int    `json:"maxSilenceSeconds"`
		IfExists    string `json:"ifExists,omitempty"`
		Beep        bool   `json:"beep"`
		TerminateOn string `json:"terminateOn,omitempty"`
	}{
		Name:        name,
		Format:      opts.Format,
		MaxDuration: int(opts.MaxDuration / time.Second),
		MaxSilence:  int(opts.MaxSilence / time.Second),
		IfExists:    opts.Exists,
		Beep:        opts.Beep,
		TerminateOn: opts.Terminate,
	}

	recordingKey := b.client.stamp(ari.NewKey(ari.LiveRecordingKey, name))

	return ari.NewLiveRecordingHandle(recordingKey, b.client.LiveRecording(), func(h *ari.LiveRecordingHandle) error {
		return b.client.post("/bridges/"+key.ID+"/record", &resp, &req)
	}), nil
}

// Subscribe creates an event subscription for events related to the given
// bridge␃entity
func (b *Bridge) Subscribe(key *ari.Key, n ...string) ari.Subscription {
	return b.client.Bus().Subscribe(key, n...)
}
<|endoftext|>


<commit_before>package client

import (
	"testing"
	"time"

	"github.com/hashicorp/nomad/nomad/structs"
)

func testPolicy(success bool, mode string) *structs.RestartPolicy {
	return &structs.RestartPolicy{
		Interval:         2 * time.Minute,
		Delay:            1 * time.Second,
		Attempts:         3,
		Mode:             mode,
		RestartOnSuccess: success,
	}
}

// withinJitter is a helper that returns whether the returned delay is within
// the jitter.
func withinJitter(expected, actual time.Duration) bool {
	return float64((actual.Nanoseconds()-expected.Nanoseconds())/
		expected.Nanoseconds()) <= jitter
}

func TestClient_RestartTracker_ModeDelay(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeDelay)
	rt := newRestartTracker(p)
	for i := 0; i < p.Attempts; i++ {
		actual, when := rt.NextRestart(127)
		if !actual {
			t.Fatalf("NextRestart() returned %v, want %v", actual, true)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Follow up restarts should cause delay.
	for i := 0; i < 3; i++ {
		actual, when := rt.NextRestart(127)
		if !actual {
			t.Fail()
		}
		if !(when > p.Delay && when < p.Interval) {
			t.Fatalf("NextRestart() returned %v; want less than %v and more than %v", when, p.Interval, p.Delay)
		}
	}
}

func TestClient_RestartTracker_ModeFail(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	rt := newRestartTracker(p)
	for i := 0; i < p.Attempts; i++ {
		actual, when := rt.NextRestart(127)
		if !actual {
			t.Fatalf("NextRestart() returned %v, want %v", actual, true)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause fail
	if actual, _ := rt.NextRestart(127); actual {
		t.Fail()
	}
}

func TestClient_RestartTracker_NoRestartOnSuccess(t *testing.T) {
	t.Parallel()
	p := testPolicy(false, structs.RestartPolicyModeDelay)
	rt := newRestartTracker(p)
	if shouldRestart, _ := rt.NextRestart(0); shouldRestart {
		t.Fatalf("NextRestart() returned %v, expected: %v", shouldRestart, false)
	}

}
<commit_msg>Allow for low precision time.Now on Windows.<commit_after>package client

import (
	"testing"
	"time"

	"github.com/hashicorp/nomad/nomad/structs"
)

func testPolicy(success bool, mode string) *structs.RestartPolicy {
	return &structs.RestartPolicy{
		Interval:         2 * time.Minute,
		Delay:            1 * time.Second,
		Attempts:         3,
		Mode:             mode,
		RestartOnSuccess: success,
	}
}

// withinJitter is a helper that returns whether the returned delay is within
// the jitter.
func withinJitter(expected, actual time.Duration) bool {
	return float64((actual.Nanoseconds()-expected.Nanoseconds())/
		expected.Nanoseconds()) <= jitter
}

func TestClient_RestartTracker_ModeDelay(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeDelay)
	rt := newRestartTracker(p)
	for i := 0; i < p.Attempts; i++ {
		actual, when := rt.NextRestart(127)
		if !actual {
			t.Fatalf("NextRestart() returned %v, want %v", actual, true)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Follow up restarts should cause delay.
	for i := 0; i < 3; i++ {
		actual, when := rt.NextRestart(127)
		if !actual {
			t.Fail()
		}
		if !(when > p.Delay && when <= p.Interval) {
			t.Fatalf("NextRestart() returned %v; want > %v and <= %v", when, p.Delay, p.Interval)
		}
	}
}

func TestClient_RestartTracker_ModeFail(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	rt := newRestartTracker(p)
	for i := 0; i < p.Attempts; i++ {
		actual, when := rt.NextRestart(127)
		if !actual {
			t.Fatalf("NextRestart() returned %v, want %v", actual, true)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause fail
	if actual, _ := rt.NextRestart(127); actual {
		t.Fail()
	}
}

func TestClient_RestartTracker_NoRestartOnSuccess(t *testing.T) {
	t.Parallel()
	p := testPolicy(false, structs.RestartPolicyModeDelay)
	rt := newRestartTracker(p)
	if shouldRestart, _ := rt.NextRestart(0); shouldRestart {
		t.Fatalf("NextRestart() returned %v, expected: %v", shouldRestart, false)
	}

}
<|endoftext|>