<commit_before>// Copyright 2016 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Raphael 'kena' Poss (knz@cockroachlabs.com)

package sql

import (
	"fmt"
	"math"

	"github.com/cockroachdb/cockroach/pkg/sql/parser"
	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
	"github.com/cockroachdb/cockroach/pkg/util/encoding"
)

// expandPlan finalizes type checking of placeholders and expands
// the query plan to its final form, including index selection and
// expansion of sub-queries. Returns an error if the initialization
// fails.
func (p *planner) expandPlan(plan planNode) (planNode, error) {
	params := expandParameters{p: p}
	return doExpandPlan(params.clear(), plan)
}

// expandParameters propagates the known row limit and desired ordering at
// a given level to the levels under it (upstream).
type expandParameters struct {
	p               *planner
	numRowsHint     int64
	desiredOrdering sqlbase.ColumnOrdering
}

func (p expandParameters) clear() expandParameters {
	p.numRowsHint = math.MaxInt64
	p.desiredOrdering = nil
	return p
}

// doExpandPlan is the algorithm that supports expandPlan().
func doExpandPlan(params expandParameters, plan planNode) (planNode, error) {
	if plan == nil {
		return nil, nil
	}

	var err error
	switch n := plan.(type) {
	case *createTableNode:
		n.sourcePlan, err = doExpandPlan(params.clear(), n.sourcePlan)

	case *updateNode:
		n.run.rows, err = doExpandPlan(params.clear(), n.run.rows)

	case *insertNode:
		n.run.rows, err = doExpandPlan(params.clear(), n.run.rows)

	case *deleteNode:
		n.run.rows, err = doExpandPlan(params.clear(), n.run.rows)

	case *createViewNode:
		n.sourcePlan, err = doExpandPlan(params.clear(), n.sourcePlan)

	case *explainDebugNode:
		n.plan, err = doExpandPlan(params.clear(), n.plan)
		if err != nil {
			return plan, err
		}
		n.plan.MarkDebug(explainDebug)

	case *explainTraceNode:
		n.plan, err = doExpandPlan(params.clear(), n.plan)
		if err != nil {
			return plan, err
		}
		n.plan.MarkDebug(explainDebug)

	case *explainPlanNode:
		if n.expanded {
			n.plan, err = doExpandPlan(params.clear(), n.plan)
			if err != nil {
				return plan, err
			}
			// Trigger limit hint propagation, which would otherwise only occur
			// during the plan's Start() phase. This may trigger additional
			// optimizations (eg. in sortNode) which the user of EXPLAIN will be
			// interested in.
			setUnlimited(n.plan)
		}

	case *indexJoinNode:
		// We ignore the return value because we know the scanNode is preserved.
		_, err = doExpandPlan(params, n.index)
		if err != nil {
			return plan, err
		}

		// The row limit and desired ordering, if any, only propagates on
		// the index side.
		_, err = doExpandPlan(params.clear(), n.table)

	case *unionNode:
		n.right, err = doExpandPlan(params, n.right)
		if err != nil {
			return plan, err
		}
		n.left, err = doExpandPlan(params, n.left)

	case *filterNode:
		n.source.plan, err = doExpandPlan(params, n.source.plan)

	case *joinNode:
		params = params.clear()
		n.left.plan, err = doExpandPlan(params, n.left.plan)
		if err != nil {
			return plan, err
		}
		n.right.plan, err = doExpandPlan(params, n.right.plan)

	case *ordinalityNode:
		// If there's a desired ordering on the ordinality column, drop it.
		if len(params.desiredOrdering) > 0 {
			newDesired := make(sqlbase.ColumnOrdering, 0, len(params.desiredOrdering))
			for _, ordInfo := range params.desiredOrdering {
				if ordInfo.ColIdx == len(n.columns)-1 {
					break
				}
				newDesired = append(newDesired, ordInfo)
			}
			params.desiredOrdering = newDesired
		}
		n.source, err = doExpandPlan(params, n.source)
		if err != nil {
			return plan, err
		}

		// We are going to "optimize" the ordering. We had an ordering
		// initially from the source, but expand() may have caused it to
		// change. So here retrieve the ordering of the source again.
		origOrdering := n.source.Ordering()

		if len(origOrdering.ordering) > 0 {
			// TODO(knz/radu) we basically have two simultaneous orderings.
			// What we really want is something that orderingInfo cannot
			// currently express: that the rows are ordered by a set of
			// columns AND at the same time they are also ordered by a
			// different set of columns. However since ordinalityNode is
			// currently the only case where this happens we consider it's not
			// worth the hassle and just use the source ordering.
			n.ordering = origOrdering
		} else {
			// No ordering defined in the source, so create a new one.
			n.ordering.exactMatchCols = origOrdering.exactMatchCols
			n.ordering.ordering = sqlbase.ColumnOrdering{
				sqlbase.ColumnOrderInfo{
					ColIdx:    len(n.columns) - 1,
					Direction: encoding.Ascending,
				},
			}
			n.ordering.unique = true
		}

	case *limitNode:
		// Estimate the limit parameters. We can't full eval them just yet,
		// because evaluation requires running potential sub-queries, which
		// cannot occur during expand.
		n.estimateLimit()
		params.numRowsHint = getLimit(n.count, n.offset)
		n.plan, err = doExpandPlan(params, n.plan)

	case *groupNode:
		params.desiredOrdering = n.desiredOrdering
		// Under a group node, there may be arbitrarily more rows
		// than those required by the context.
		params.numRowsHint = math.MaxInt64
		n.plan, err = doExpandPlan(params, n.plan)

	case *windowNode:
		// Under a window node, there may be arbitrarily more rows
		// than those required by the context.
		params.numRowsHint = math.MaxInt64
		n.plan, err = doExpandPlan(params, n.plan)

	case *sortNode:
		params.desiredOrdering = n.ordering
		n.plan, err = doExpandPlan(params, n.plan)
		if err != nil {
			return plan, err
		}

		// Check to see if the requested ordering is compatible with the existing
		// ordering.
		existingOrdering := n.plan.Ordering()
		match := computeOrderingMatch(n.ordering, existingOrdering, false)
		if match < len(n.ordering) {
			n.needSort = true
		} else if len(n.columns) < len(n.plan.Columns()) {
			// No sorting required, but we have to strip off the extra render
			// expressions we added. So keep the sort node.
		} else {
			// Sort node fully disappears.
			plan = n.plan
		}

	case *distinctNode:
		// TODO(radu/knz) perhaps we can propagate the DISTINCT
		// clause as desired ordering/exact match for the source node.
		n.plan, err = doExpandPlan(params, n.plan)
		if err != nil {
			return plan, err
		}

		ordering := n.plan.Ordering()
		if !ordering.isEmpty() {
			n.columnsInOrder = make([]bool, len(n.plan.Columns()))
			for colIdx := range ordering.exactMatchCols {
				if colIdx >= len(n.columnsInOrder) {
					// If the exact-match column is not part of the output, we can safely ignore it.
					continue
				}
				n.columnsInOrder[colIdx] = true
			}
			for _, c := range ordering.ordering {
				if c.ColIdx >= len(n.columnsInOrder) {
					// Cannot use sort order. This happens when the
					// columns used for sorting are not part of the output.
					// e.g. SELECT a FROM t ORDER BY c.
					n.columnsInOrder = nil
					break
				}
				n.columnsInOrder[c.ColIdx] = true
			}
		}

	case *scanNode:
		plan, err = expandScanNode(params, n)

	case *renderNode:
		plan, err = expandRenderNode(params, n)

	case *delayedNode:
		n.plan, err = n.constructor(params.p)
		if err == nil {
			n.plan, err = doExpandPlan(params, n.plan)
		}

	case *valuesNode:
	case *alterTableNode:
	case *copyNode:
	case *createDatabaseNode:
	case *createIndexNode:
	case *createUserNode:
	case *dropDatabaseNode:
	case *dropIndexNode:
	case *dropTableNode:
	case *dropViewNode:
	case *emptyNode:
	case *hookFnNode:
	case *splitNode:
	case *valueGenerator:

	default:
		panic(fmt.Sprintf("unhandled node type: %T", plan))
	}
	return plan, err
}

func expandScanNode(params expandParameters, s *scanNode) (planNode, error) {
	var analyzeOrdering analyzeOrderingFn
	if len(params.desiredOrdering) > 0 {
		analyzeOrdering = func(indexOrdering orderingInfo) (matchingCols, totalCols int) {
			match := computeOrderingMatch(params.desiredOrdering, indexOrdering, false)
			return match, len(params.desiredOrdering)
		}
	}

	// If we have a reasonable limit, prefer an order matching index even if
	// it is not covering.
	var preferOrderMatchingIndex bool
	if len(params.desiredOrdering) > 0 && params.numRowsHint <= 1000 {
		preferOrderMatchingIndex = true
	}

	plan, err := params.p.selectIndex(s, analyzeOrdering, preferOrderMatchingIndex)
	if err != nil {
		return s, err
	}
	return plan, nil
}

func expandRenderNode(params expandParameters, r *renderNode) (planNode, error) {
	params.desiredOrdering = translateOrdering(params.desiredOrdering, r)

	var err error
	r.source.plan, err = doExpandPlan(params, r.source.plan)
	if err != nil {
		return r, err
	}

	// Elide the render node if it renders its source as-is.

	sourceCols := r.source.plan.Columns()
	if len(r.columns) == len(sourceCols) && r.source.info.viewDesc == nil {
		// 1) we don't drop renderNodes which also interface to a view, because
		// CREATE VIEW needs it.
		// TODO(knz) make this optimization conditional on a flag, which can
		// be set to false by CREATE VIEW.
		//
		// 2) we don't drop renderNodes which have a different number of
		// columns than their sources, because some nodes currently assume
		// the number of source columns doesn't change between
		// instantiation and Start() (e.g. groupNode).
		// TODO(knz) investigate this further and enable the optimization fully.

		foundNonTrivialRender := false
		for i, e := range r.render {
			if r.columns[i].omitted {
				continue
			}
			if iv, ok := e.(*parser.IndexedVar); ok && i < len(sourceCols) &&
				(iv.Idx == i && sourceCols[i].Name == r.columns[i].Name) {
				continue
			}
			foundNonTrivialRender = true
			break
		}
		if !foundNonTrivialRender {
			// Nothing special rendered, remove the render node entirely.
			return r.source.plan, nil
		}
	}

	r.ordering = r.computeOrdering(r.source.plan.Ordering())

	return r, nil
}

// translateOrdering modifies a desired ordering on the output of the
// renderNode to a desired ordering on its its input.
//
// For example, it translates a desired ordering [@2 asc, @1 desc] for
// a render node that renders [@4, @3, @2] into a desired ordering [@3
// asc, @4 desc].
func translateOrdering(desiredDown sqlbase.ColumnOrdering, r *renderNode) sqlbase.ColumnOrdering {
	var desiredUp sqlbase.ColumnOrdering

	for _, colOrder := range desiredDown {
		rendered := r.render[colOrder.ColIdx]
		if _, ok := rendered.(parser.Datum); ok {
			// Simple constants do not participate in ordering. Just ignore.
			continue
		}
		if v, ok := rendered.(*parser.IndexedVar); ok {
			// This is a simple render, so we can propagate the desired ordering.
			// However take care of avoiding duplicate ordering requests in
			// case there is more than one render for the same source column.
			duplicate := false
			for _, desiredOrderCol := range desiredUp {
				if desiredOrderCol.ColIdx == v.Idx {
					duplicate = true
					break
				}
			}
			if !duplicate {
				desiredUp = append(desiredUp,
					sqlbase.ColumnOrderInfo{ColIdx: v.Idx, Direction: colOrder.Direction})
			}
			continue
		}
		// Anything else and we can't propagate the desired order.
		break
	}

	return desiredUp
}
<commit_msg>sql: minor cleanup for expand plan parameters<commit_after>// Copyright 2016 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Raphael 'kena' Poss (knz@cockroachlabs.com)

package sql

import (
	"fmt"
	"math"

	"github.com/cockroachdb/cockroach/pkg/sql/parser"
	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
	"github.com/cockroachdb/cockroach/pkg/util/encoding"
)

// expandPlan finalizes type checking of placeholders and expands
// the query plan to its final form, including index selection and
// expansion of sub-queries. Returns an error if the initialization
// fails.
func (p *planner) expandPlan(plan planNode) (planNode, error) {
	return doExpandPlan(p, noParams, plan)
}

// expandParameters propagates the known row limit and desired ordering at
// a given level to the levels under it (upstream).
type expandParameters struct {
	numRowsHint     int64
	desiredOrdering sqlbase.ColumnOrdering
}

var noParams = expandParameters{numRowsHint: math.MaxInt64, desiredOrdering: nil}

// doExpandPlan is the algorithm that supports expandPlan().
func doExpandPlan(p *planner, params expandParameters, plan planNode) (planNode, error) {
	if plan == nil {
		return nil, nil
	}

	var err error
	switch n := plan.(type) {
	case *createTableNode:
		n.sourcePlan, err = doExpandPlan(p, noParams, n.sourcePlan)

	case *updateNode:
		n.run.rows, err = doExpandPlan(p, noParams, n.run.rows)

	case *insertNode:
		n.run.rows, err = doExpandPlan(p, noParams, n.run.rows)

	case *deleteNode:
		n.run.rows, err = doExpandPlan(p, noParams, n.run.rows)

	case *createViewNode:
		n.sourcePlan, err = doExpandPlan(p, noParams, n.sourcePlan)

	case *explainDebugNode:
		n.plan, err = doExpandPlan(p, noParams, n.plan)
		if err != nil {
			return plan, err
		}
		n.plan.MarkDebug(explainDebug)

	case *explainTraceNode:
		n.plan, err = doExpandPlan(p, noParams, n.plan)
		if err != nil {
			return plan, err
		}
		n.plan.MarkDebug(explainDebug)

	case *explainPlanNode:
		if n.expanded {
			n.plan, err = doExpandPlan(p, noParams, n.plan)
			if err != nil {
				return plan, err
			}
			// Trigger limit hint propagation, which would otherwise only occur
			// during the plan's Start() phase. This may trigger additional
			// optimizations (eg. in sortNode) which the user of EXPLAIN will be
			// interested in.
			setUnlimited(n.plan)
		}

	case *indexJoinNode:
		// We ignore the return value because we know the scanNode is preserved.
		_, err = doExpandPlan(p, params, n.index)
		if err != nil {
			return plan, err
		}

		// The row limit and desired ordering, if any, only propagates on
		// the index side.
		_, err = doExpandPlan(p, noParams, n.table)

	case *unionNode:
		n.right, err = doExpandPlan(p, params, n.right)
		if err != nil {
			return plan, err
		}
		n.left, err = doExpandPlan(p, params, n.left)

	case *filterNode:
		n.source.plan, err = doExpandPlan(p, params, n.source.plan)

	case *joinNode:
		n.left.plan, err = doExpandPlan(p, noParams, n.left.plan)
		if err != nil {
			return plan, err
		}
		n.right.plan, err = doExpandPlan(p, noParams, n.right.plan)

	case *ordinalityNode:
		// If there's a desired ordering on the ordinality column, drop it.
		if len(params.desiredOrdering) > 0 {
			newDesired := make(sqlbase.ColumnOrdering, 0, len(params.desiredOrdering))
			for _, ordInfo := range params.desiredOrdering {
				if ordInfo.ColIdx == len(n.columns)-1 {
					break
				}
				newDesired = append(newDesired, ordInfo)
			}
			params.desiredOrdering = newDesired
		}
		n.source, err = doExpandPlan(p, params, n.source)
		if err != nil {
			return plan, err
		}

		// We are going to "optimize" the ordering. We had an ordering
		// initially from the source, but expand() may have caused it to
		// change. So here retrieve the ordering of the source again.
		origOrdering := n.source.Ordering()

		if len(origOrdering.ordering) > 0 {
			// TODO(knz/radu) we basically have two simultaneous orderings.
			// What we really want is something that orderingInfo cannot
			// currently express: that the rows are ordered by a set of
			// columns AND at the same time they are also ordered by a
			// different set of columns. However since ordinalityNode is
			// currently the only case where this happens we consider it's not
			// worth the hassle and just use the source ordering.
			n.ordering = origOrdering
		} else {
			// No ordering defined in the source, so create a new one.
			n.ordering.exactMatchCols = origOrdering.exactMatchCols
			n.ordering.ordering = sqlbase.ColumnOrdering{
				sqlbase.ColumnOrderInfo{
					ColIdx:    len(n.columns) - 1,
					Direction: encoding.Ascending,
				},
			}
			n.ordering.unique = true
		}

	case *limitNode:
		// Estimate the limit parameters. We can't full eval them just yet,
		// because evaluation requires running potential sub-queries, which
		// cannot occur during expand.
		n.estimateLimit()
		params.numRowsHint = getLimit(n.count, n.offset)
		n.plan, err = doExpandPlan(p, params, n.plan)

	case *groupNode:
		params.desiredOrdering = n.desiredOrdering
		// Under a group node, there may be arbitrarily more rows
		// than those required by the context.
		params.numRowsHint = math.MaxInt64
		n.plan, err = doExpandPlan(p, params, n.plan)

	case *windowNode:
		// Under a window node, there may be arbitrarily more rows
		// than those required by the context.
		params.numRowsHint = math.MaxInt64
		n.plan, err = doExpandPlan(p, params, n.plan)

	case *sortNode:
		params.desiredOrdering = n.ordering
		n.plan, err = doExpandPlan(p, params, n.plan)
		if err != nil {
			return plan, err
		}

		// Check to see if the requested ordering is compatible with the existing
		// ordering.
		existingOrdering := n.plan.Ordering()
		match := computeOrderingMatch(n.ordering, existingOrdering, false)
		if match < len(n.ordering) {
			n.needSort = true
		} else if len(n.columns) < len(n.plan.Columns()) {
			// No sorting required, but we have to strip off the extra render
			// expressions we added. So keep the sort node.
		} else {
			// Sort node fully disappears.
			plan = n.plan
		}

	case *distinctNode:
		// TODO(radu/knz) perhaps we can propagate the DISTINCT
		// clause as desired ordering/exact match for the source node.
		n.plan, err = doExpandPlan(p, params, n.plan)
		if err != nil {
			return plan, err
		}

		ordering := n.plan.Ordering()
		if !ordering.isEmpty() {
			n.columnsInOrder = make([]bool, len(n.plan.Columns()))
			for colIdx := range ordering.exactMatchCols {
				if colIdx >= len(n.columnsInOrder) {
					// If the exact-match column is not part of the output, we can safely ignore it.
					continue
				}
				n.columnsInOrder[colIdx] = true
			}
			for _, c := range ordering.ordering {
				if c.ColIdx >= len(n.columnsInOrder) {
					// Cannot use sort order. This happens when the
					// columns used for sorting are not part of the output.
					// e.g. SELECT a FROM t ORDER BY c.
					n.columnsInOrder = nil
					break
				}
				n.columnsInOrder[c.ColIdx] = true
			}
		}

	case *scanNode:
		plan, err = expandScanNode(p, params, n)

	case *renderNode:
		plan, err = expandRenderNode(p, params, n)

	case *delayedNode:
		n.plan, err = n.constructor(p)
		if err == nil {
			n.plan, err = doExpandPlan(p, params, n.plan)
		}

	case *valuesNode:
	case *alterTableNode:
	case *copyNode:
	case *createDatabaseNode:
	case *createIndexNode:
	case *createUserNode:
	case *dropDatabaseNode:
	case *dropIndexNode:
	case *dropTableNode:
	case *dropViewNode:
	case *emptyNode:
	case *hookFnNode:
	case *splitNode:
	case *valueGenerator:

	default:
		panic(fmt.Sprintf("unhandled node type: %T", plan))
	}
	return plan, err
}

func expandScanNode(p *planner, params expandParameters, s *scanNode) (planNode, error) {
	var analyzeOrdering analyzeOrderingFn
	if len(params.desiredOrdering) > 0 {
		analyzeOrdering = func(indexOrdering orderingInfo) (matchingCols, totalCols int) {
			match := computeOrderingMatch(params.desiredOrdering, indexOrdering, false)
			return match, len(params.desiredOrdering)
		}
	}

	// If we have a reasonable limit, prefer an order matching index even if
	// it is not covering.
	var preferOrderMatchingIndex bool
	if len(params.desiredOrdering) > 0 && params.numRowsHint <= 1000 {
		preferOrderMatchingIndex = true
	}

	plan, err := p.selectIndex(s, analyzeOrdering, preferOrderMatchingIndex)
	if err != nil {
		return s, err
	}
	return plan, nil
}

func expandRenderNode(p *planner, params expandParameters, r *renderNode) (planNode, error) {
	params.desiredOrdering = translateOrdering(params.desiredOrdering, r)

	var err error
	r.source.plan, err = doExpandPlan(p, params, r.source.plan)
	if err != nil {
		return r, err
	}

	// Elide the render node if it renders its source as-is.

	sourceCols := r.source.plan.Columns()
	if len(r.columns) == len(sourceCols) && r.source.info.viewDesc == nil {
		// 1) we don't drop renderNodes which also interface to a view, because
		// CREATE VIEW needs it.
		// TODO(knz) make this optimization conditional on a flag, which can
		// be set to false by CREATE VIEW.
		//
		// 2) we don't drop renderNodes which have a different number of
		// columns than their sources, because some nodes currently assume
		// the number of source columns doesn't change between
		// instantiation and Start() (e.g. groupNode).
		// TODO(knz) investigate this further and enable the optimization fully.

		foundNonTrivialRender := false
		for i, e := range r.render {
			if r.columns[i].omitted {
				continue
			}
			if iv, ok := e.(*parser.IndexedVar); ok && i < len(sourceCols) &&
				(iv.Idx == i && sourceCols[i].Name == r.columns[i].Name) {
				continue
			}
			foundNonTrivialRender = true
			break
		}
		if !foundNonTrivialRender {
			// Nothing special rendered, remove the render node entirely.
			return r.source.plan, nil
		}
	}

	r.ordering = r.computeOrdering(r.source.plan.Ordering())

	return r, nil
}

// translateOrdering modifies a desired ordering on the output of the
// renderNode to a desired ordering on its its input.
//
// For example, it translates a desired ordering [@2 asc, @1 desc] for
// a render node that renders [@4, @3, @2] into a desired ordering [@3
// asc, @4 desc].
func translateOrdering(desiredDown sqlbase.ColumnOrdering, r *renderNode) sqlbase.ColumnOrdering {
	var desiredUp sqlbase.ColumnOrdering

	for _, colOrder := range desiredDown {
		rendered := r.render[colOrder.ColIdx]
		if _, ok := rendered.(parser.Datum); ok {
			// Simple constants do not participate in ordering. Just ignore.
			continue
		}
		if v, ok := rendered.(*parser.IndexedVar); ok {
			// This is a simple render, so we can propagate the desired ordering.
			// However take care of avoiding duplicate ordering requests in
			// case there is more than one render for the same source column.
			duplicate := false
			for _, desiredOrderCol := range desiredUp {
				if desiredOrderCol.ColIdx == v.Idx {
					duplicate = true
					break
				}
			}
			if !duplicate {
				desiredUp = append(desiredUp,
					sqlbase.ColumnOrderInfo{ColIdx: v.Idx, Direction: colOrder.Direction})
			}
			continue
		}
		// Anything else and we can't propagate the desired order.
		break
	}

	return desiredUp
}
<|endoftext|>


<commit_before>/*
Copyright The Helm Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package version // import "k8s.io/helm/pkg/version"

import "k8s.io/helm/pkg/proto/hapi/version"

var (
	// Version is the current version of the Helm.
	// Update this whenever making a new release.
	// The version is of the format Major.Minor.Patch[-Prerelease][+BuildMetadata]
	//
	// Increment major number for new feature additions and behavioral changes.
	// Increment minor number for bug fixes and performance enhancements.
	// Increment patch number for critical fixes to existing releases.
	Version = "v2.15"

	// BuildMetadata is extra build time data
	BuildMetadata = "unreleased"
	// GitCommit is the git sha1
	GitCommit = ""
	// GitTreeState is the state of the git tree
	GitTreeState = ""
)

// GetVersion returns the semver string of the version
func GetVersion() string {
	if BuildMetadata == "" {
		return Version
	}
	return Version + "+" + BuildMetadata
}

// GetVersionProto returns protobuf representing the version
func GetVersionProto() *version.Version {
	return &version.Version{
		SemVer:       GetVersion(),
		GitCommit:    GitCommit,
		GitTreeState: GitTreeState,
	}
}
<commit_msg>bump version to v2.16.0-rc.1<commit_after>/*
Copyright The Helm Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package version // import "k8s.io/helm/pkg/version"

import "k8s.io/helm/pkg/proto/hapi/version"

var (
	// Version is the current version of the Helm.
	// Update this whenever making a new release.
	// The version is of the format Major.Minor.Patch[-Prerelease][+BuildMetadata]
	//
	// Increment major number for new feature additions and behavioral changes.
	// Increment minor number for bug fixes and performance enhancements.
	// Increment patch number for critical fixes to existing releases.
	Version = "v2.16"

	// BuildMetadata is extra build time data
	BuildMetadata = "unreleased"
	// GitCommit is the git sha1
	GitCommit = ""
	// GitTreeState is the state of the git tree
	GitTreeState = ""
)

// GetVersion returns the semver string of the version
func GetVersion() string {
	if BuildMetadata == "" {
		return Version
	}
	return Version + "+" + BuildMetadata
}

// GetVersionProto returns protobuf representing the version
func GetVersionProto() *version.Version {
	return &version.Version{
		SemVer:       GetVersion(),
		GitCommit:    GitCommit,
		GitTreeState: GitTreeState,
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 the LinuxBoot Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package visitors

import (
	"errors"
	"fmt"
	"io/ioutil"

	"github.com/linuxboot/fiano/pkg/uefi"
)

// InsertType defines the insert type operation that is requested
type InsertType int

// Insert Types
const (
	// These first two specify a firmware volume.

	// InsertFront inserts a file at the beginning of the firmware volume,
	// which is specified by 1) FVname GUID, or (File GUID/File name) of a file
	// inside that FV.
	InsertFront InsertType = iota
	// InsertEnd inserts a file at the end of the specified firmware volume.
	InsertEnd

	// These two specify a File to insert before or after
	// InsertAfter inserts after the specified file,
	// which is specified by a File GUID or File name.
	InsertAfter
	// InsertBefore inserts before the specified file.
	InsertBefore
	// InsertDXE inserts into the Dxe Firmware Volume.
	InsertDXE
	// TODO: Add InsertIn
)

var insertTypeNames = map[InsertType]string{
	InsertFront:  "insert_front",
	InsertEnd:    "insert_end",
	InsertAfter:  "insert_after",
	InsertBefore: "insert_before",
	InsertDXE:    "insert_dxe",
}

// String creates a string representation for the insert type.
func (i InsertType) String() string {
	if t, ok := insertTypeNames[i]; ok {
		return t
	}
	return "UNKNOWN"
}

// Insert inserts a firmware file into an FV
type Insert struct {
	// Input
	Predicate func(f uefi.Firmware) bool
	NewFile   *uefi.File
	InsertType

	// Matched File
	FileMatch uefi.Firmware
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *Insert) Run(f uefi.Firmware) error {
	// First run "find" to generate a position to insert into.
	find := Find{
		Predicate: v.Predicate,
	}
	if err := find.Run(f); err != nil {
		return err
	}

	if numMatch := len(find.Matches); numMatch > 1 {
		return fmt.Errorf("more than one match, only one match allowed! got %v", find.Matches)
	} else if numMatch == 0 {
		return errors.New("no matches found")
	}

	// Find should only match a file or a firmware volume. If it's an FV, we can
	// edit the FV directly.
	if fvMatch, ok := find.Matches[0].(*uefi.FirmwareVolume); ok {
		switch v.InsertType {
		case InsertFront:
			fvMatch.Files = append([]*uefi.File{v.NewFile}, fvMatch.Files...)
		case InsertEnd:
			fvMatch.Files = append(fvMatch.Files, v.NewFile)
		default:
			return fmt.Errorf("matched FV but insert operation was %s, which only matches Files",
				v.InsertType.String())
		}
		return nil
	}
	var ok bool
	if v.FileMatch, ok = find.Matches[0].(*uefi.File); !ok {
		return fmt.Errorf("match was not a file or a firmware volume: got %T, unable to insert", find.Matches[0])
	}
	// Match is a file, apply visitor.
	return f.Apply(v)
}

// Visit applies the Insert visitor to any Firmware type.
func (v *Insert) Visit(f uefi.Firmware) error {
	switch f := f.(type) {
	case *uefi.FirmwareVolume:
		for i := 0; i < len(f.Files); i++ {
			if f.Files[i] == v.FileMatch {
				switch v.InsertType {
				case InsertFront:
					f.Files = append([]*uefi.File{v.NewFile}, f.Files...)
				case InsertDXE:
					fallthrough
				case InsertEnd:
					f.Files = append(f.Files, v.NewFile)
				case InsertAfter:
					f.Files = append(f.Files[:i+1], append([]*uefi.File{v.NewFile}, f.Files[i+1:]...)...)
				case InsertBefore:
					f.Files = append(f.Files[:i], append([]*uefi.File{v.NewFile}, f.Files[i:]...)...)
				}
				return nil
			}
		}
	}

	return f.ApplyChildren(v)
}

func genInsertCLI(iType InsertType) func(args []string) (uefi.Visitor, error) {
	return func(args []string) (uefi.Visitor, error) {
		var pred FindPredicate
		var err error
		var filename string

		if iType == InsertDXE {
			pred = FindFileTypePredicate(uefi.FVFileTypeDXECore)
			filename = args[0]
		} else {
			pred, err = FindFileFVPredicate(args[0])
			if err != nil {
				return nil, err
			}
			filename = args[1]
		}

		newFileBuf, err := ioutil.ReadFile(filename)
		if err != nil {
			return nil, err
		}
		// Parse the file.
		file, err := uefi.NewFile(newFileBuf)
		if err != nil {
			return nil, err
		}

		// Insert File.
		return &Insert{
			Predicate:  pred,
			NewFile:    file,
			InsertType: iType,
		}, nil
	}
}

func init() {
	RegisterCLI(insertTypeNames[InsertFront],
		"insert a file at the beginning of a firmware volume", 2, genInsertCLI(InsertFront))
	RegisterCLI(insertTypeNames[InsertEnd],
		"insert a file at the end of a firmware volume", 2, genInsertCLI(InsertEnd))
	RegisterCLI(insertTypeNames[InsertDXE],
		"insert a file at the end of the DXE firmware volume", 1, genInsertCLI(InsertDXE))
	RegisterCLI(insertTypeNames[InsertAfter],
		"insert a file after another file", 2, genInsertCLI(InsertAfter))
	RegisterCLI(insertTypeNames[InsertBefore],
		"insert a file before another file", 2, genInsertCLI(InsertBefore))
}
<commit_msg>Add replace_ffs command.<commit_after>// Copyright 2018 the LinuxBoot Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package visitors

import (
	"errors"
	"fmt"
	"io/ioutil"

	"github.com/linuxboot/fiano/pkg/uefi"
)

// InsertType defines the insert type operation that is requested
type InsertType int

// Insert Types
const (
	// These first two specify a firmware volume.

	// InsertFront inserts a file at the beginning of the firmware volume,
	// which is specified by 1) FVname GUID, or (File GUID/File name) of a file
	// inside that FV.
	InsertFront InsertType = iota
	// InsertEnd inserts a file at the end of the specified firmware volume.
	InsertEnd

	// These two specify a File to insert before or after
	// InsertAfter inserts after the specified file,
	// which is specified by a File GUID or File name.
	InsertAfter
	// InsertBefore inserts before the specified file.
	InsertBefore
	// InsertDXE inserts into the Dxe Firmware Volume. This works by searching
	// for the DxeCore first to identify the Dxe Firmware Volume.
	InsertDXE

	// ReplaceFFS replaces the found file with the new FFS. This is used
	// as a shortcut for remove and insert combined, but also when we want to make
	// sure that the starting offset of the new file is the same as the old.
	ReplaceFFS
	// TODO: Add InsertIn
)

var insertTypeNames = map[InsertType]string{
	InsertFront:  "insert_front",
	InsertEnd:    "insert_end",
	InsertAfter:  "insert_after",
	InsertBefore: "insert_before",
	InsertDXE:    "insert_dxe",
	ReplaceFFS:   "replace_ffs",
}

// String creates a string representation for the insert type.
func (i InsertType) String() string {
	if t, ok := insertTypeNames[i]; ok {
		return t
	}
	return "UNKNOWN"
}

// Insert inserts a firmware file into an FV
type Insert struct {
	// Input
	Predicate func(f uefi.Firmware) bool
	NewFile   *uefi.File
	InsertType

	// Matched File
	FileMatch uefi.Firmware
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *Insert) Run(f uefi.Firmware) error {
	// First run "find" to generate a position to insert into.
	find := Find{
		Predicate: v.Predicate,
	}
	if err := find.Run(f); err != nil {
		return err
	}

	if numMatch := len(find.Matches); numMatch > 1 {
		return fmt.Errorf("more than one match, only one match allowed! got %v", find.Matches)
	} else if numMatch == 0 {
		return errors.New("no matches found")
	}

	// Find should only match a file or a firmware volume. If it's an FV, we can
	// edit the FV directly.
	if fvMatch, ok := find.Matches[0].(*uefi.FirmwareVolume); ok {
		switch v.InsertType {
		case InsertFront:
			fvMatch.Files = append([]*uefi.File{v.NewFile}, fvMatch.Files...)
		case InsertEnd:
			fvMatch.Files = append(fvMatch.Files, v.NewFile)
		default:
			return fmt.Errorf("matched FV but insert operation was %s, which only matches Files",
				v.InsertType.String())
		}
		return nil
	}
	var ok bool
	if v.FileMatch, ok = find.Matches[0].(*uefi.File); !ok {
		return fmt.Errorf("match was not a file or a firmware volume: got %T, unable to insert", find.Matches[0])
	}
	// Match is a file, apply visitor.
	return f.Apply(v)
}

// Visit applies the Insert visitor to any Firmware type.
func (v *Insert) Visit(f uefi.Firmware) error {
	switch f := f.(type) {
	case *uefi.FirmwareVolume:
		for i := 0; i < len(f.Files); i++ {
			if f.Files[i] == v.FileMatch {
				switch v.InsertType {
				case InsertFront:
					f.Files = append([]*uefi.File{v.NewFile}, f.Files...)
				case InsertDXE:
					fallthrough
				case InsertEnd:
					f.Files = append(f.Files, v.NewFile)
				case InsertAfter:
					f.Files = append(f.Files[:i+1], append([]*uefi.File{v.NewFile}, f.Files[i+1:]...)...)
				case InsertBefore:
					f.Files = append(f.Files[:i], append([]*uefi.File{v.NewFile}, f.Files[i:]...)...)
				case ReplaceFFS:
					f.Files = append(f.Files[:i], append([]*uefi.File{v.NewFile}, f.Files[i+1:]...)...)
				}
				return nil
			}
		}
	}

	return f.ApplyChildren(v)
}

func genInsertCLI(iType InsertType) func(args []string) (uefi.Visitor, error) {
	return func(args []string) (uefi.Visitor, error) {
		var pred FindPredicate
		var err error
		var filename string

		if iType == InsertDXE {
			pred = FindFileTypePredicate(uefi.FVFileTypeDXECore)
			filename = args[0]
		} else {
			pred, err = FindFileFVPredicate(args[0])
			if err != nil {
				return nil, err
			}
			filename = args[1]
		}

		newFileBuf, err := ioutil.ReadFile(filename)
		if err != nil {
			return nil, err
		}
		// Parse the file.
		file, err := uefi.NewFile(newFileBuf)
		if err != nil {
			return nil, err
		}

		// Insert File.
		return &Insert{
			Predicate:  pred,
			NewFile:    file,
			InsertType: iType,
		}, nil
	}
}

func init() {
	RegisterCLI(insertTypeNames[InsertFront],
		"insert a file at the beginning of a firmware volume", 2, genInsertCLI(InsertFront))
	RegisterCLI(insertTypeNames[InsertEnd],
		"insert a file at the end of a firmware volume", 2, genInsertCLI(InsertEnd))
	RegisterCLI(insertTypeNames[InsertDXE],
		"insert a file at the end of the DXE firmware volume", 1, genInsertCLI(InsertDXE))
	RegisterCLI(insertTypeNames[InsertAfter],
		"insert a file after another file", 2, genInsertCLI(InsertAfter))
	RegisterCLI(insertTypeNames[InsertBefore],
		"insert a file before another file", 2, genInsertCLI(InsertBefore))
	RegisterCLI(insertTypeNames[ReplaceFFS],
		"replace a file with another file", 2, genInsertCLI(ReplaceFFS))
}
<|endoftext|>


<commit_before>package unit

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/innovate-technologies/Dispatch/dispatchd/etcdcache"

	"strconv"

	"github.com/spf13/afero"

	"github.com/innovate-technologies/Dispatch/dispatchd/config"
	"github.com/innovate-technologies/Dispatch/dispatchd/etcdclient"
	"github.com/innovate-technologies/Dispatch/dispatchd/unit/state"
	"github.com/innovate-technologies/Dispatch/interfaces"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"github.com/coreos/go-systemd/dbus"
	"golang.org/x/net/context"
)

const unitPath = "/var/run/dispatch/"

var (
	// Config is a pointer need to be set to the main configuration
	Config *config.ConfigurationInfo
	ctx    = context.Background()
	//EtcdAPI is the etcd keys api
	EtcdAPI interfaces.EtcdAPI = etcdclient.GetEtcdv3()
	// DBusConnection is the connection to the system's D-Bus
	DBusConnection DBusConnectionInterface
	// FS is the file system to be used
	FS = afero.NewOsFs()
)

// UnitInterface is the interface to a Unit
type UnitInterface interface {
	Start()
	Stop()
	Restart()
	Create()
	PutOnQueue()
	SaveOnEtcd()
	Destroy()
	LoadAndWatch()
	Watch()
	SetState(state.State)
	SetDesiredState(s state.State)
}

// Unit is a struct containing all info of a specific unit
type Unit struct {
	Name         string `json:"name"`
	Machine      string `json:"machine"`
	Template     string `json:"template,omitempty"` // is set with template name if from Template
	Global       string `json:"global,omitempty"`   // is set with global name if from global
	State        state.State
	DesiredState state.State
	Ports        []int64 `json:"ports"`
	Constraints  map[string]string
	UnitContent  string `json:"unitContent"`
	onEtcd       bool
	onDisk       bool
	etcdName     string
	etcdCache    *etcdcache.EtcdCache
	disableCache bool
	runContext   context.Context
	runCancel    context.CancelFunc
}

// GetAll returns all units in our zone
func GetAll() ([]Unit, error) {
	cache, err := etcdcache.NewForPrefix(fmt.Sprintf("/dispatch/%s/units", Config.Zone))
	if err != nil {
		return nil, err
	}

	units := []Unit{}
	hadUnitNames := map[string]bool{}

	for _, kv := range cache.GetAll() {
		pathParts := strings.Split(string(kv.Key), "/")
		if _, ok := hadUnitNames[pathParts[4]]; !ok {
			units = append(units, NewFromEtcd(pathParts[4]))
			hadUnitNames[pathParts[4]] = true
		}

	}

	return units, nil
}

// New returns a new Unit
func New() Unit {
	setUpDBus()
	return Unit{}
}

// NewFromEtcd creates a new unit with info from etcd
func NewFromEtcd(name string) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}

	unit := New()
	fillFromEtcd(&unit, name)

	return unit
}

// NewFromEtcdWithCache creates a new unit with info from etcd using a specified cache
func NewFromEtcdWithCache(name string, cache *etcdcache.EtcdCache) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}
	unit := New()
	unit.etcdCache = cache
	fillFromEtcd(&unit, name)

	return unit
}

func fillFromEtcd(unit *Unit, name string) {
	unit.etcdName = name
	unit.onEtcd = true
	unit.Name = unit.getKeyFromEtcd("name")
	unit.Machine = unit.getKeyFromEtcd("machine")
	unit.Template = unit.getKeyFromEtcd("template")
	unit.Global = unit.getKeyFromEtcd("global")
	unit.UnitContent = unit.getKeyFromEtcd("unit")
	unit.DesiredState = state.ForString(unit.getKeyFromEtcd("desiredState"))
	unit.State = state.ForString(unit.getKeyFromEtcd("state"))

	unit.Ports = []int64{}
	portsStringArray := strings.Split(unit.getKeyFromEtcd("ports"), ",")
	for _, portString := range portsStringArray {
		port, err := strconv.ParseInt(portString, 10, 64)
		if err == nil {
			unit.Ports = append(unit.Ports, port)
		}
	}
}

// Start starts the specific unit
func (unit *Unit) Start() {
	log.Println("Starting unit", unit.Name)
	unit.SetState(state.Starting)
	c := make(chan string)
	_, err := DBusConnection.StartUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error starting unit", unit.Name, err)
		return
	}
	<-c
	unit.SetState(state.Active)
}

// Stop stops the unit
func (unit *Unit) Stop() {
	log.Println("Stopping unit", unit.Name)
	c := make(chan string)
	_, err := DBusConnection.StopUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error stopping unit", unit.Name, err)
		log.Println("Killing unit", unit.Name)
		DBusConnection.KillUnit(unit.Name, 9) // the big guns!
		unit.SetState(state.Dead)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Stopped unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Restart restarts the unit
func (unit *Unit) Restart() {
	unit.Stop()
	unit.Start()
}

// Create writes the unit to the disk
func (unit *Unit) Create() {
	if unit.Name == "" {
		log.Println("Error starting unit with no name")
		if unit.etcdName != "" {
			// Faulty unit
			unit.Name = unit.etcdName
			unit.Destroy()
		}
		return // can't create file without a name
	}
	unit.runContext, unit.runCancel = context.WithCancel(context.Background())

	thisUnitPath := unitPath + unit.Name

	fileContent := []byte(unit.UnitContent)

	FS.MkdirAll(unitPath, 0755)

	FS.Remove(thisUnitPath) // make sure the old unit is gone
	file, err := FS.Create(thisUnitPath)
	if err != nil {
		panic(err)
	}
	_, err = file.Write(fileContent)
	if err != nil {
		panic(err)
	}

	c := make(chan string)
	_, stopErr := DBusConnection.StopUnit(unit.Name, "fail", c) // stop unit to make sure new one is loaded
	if stopErr == nil {
		<-c // wait on completion
	}

	unit.onDisk = true
	_, dberr := DBusConnection.LinkUnitFiles([]string{thisUnitPath}, true, true)
	fmt.Println(dberr)
	DBusConnection.Reload()
}

// PutOnQueue places a specific unit on the queue
func (unit *Unit) PutOnQueue() {
	if unit.Global != "" {
		return
	}
	log.Println("Placing", unit.Name, "on queue")
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), unit.Name)
}

func (unit *Unit) normalizeName() {
	if !strings.HasSuffix(unit.Name, ".service") {
		unit.Name += ".service"
	}
}

// SaveOnEtcd saves the unit to etcd
func (unit *Unit) SaveOnEtcd() {
	unit.normalizeName()

	log.Println("Saving", unit.Name, "to etcd")

	unit.setKeyOnEtcd("name", unit.Name)
	unit.setKeyOnEtcd("unit", unit.UnitContent)
	unit.setKeyOnEtcd("template", unit.Template)
	unit.setKeyOnEtcd("desiredState", unit.DesiredState.String())

	portStrings := []string{}
	for _, port := range unit.Ports {
		portStrings = append(portStrings, strconv.FormatInt(port, 10))
	}
	unit.setKeyOnEtcd("ports", strings.Join(portStrings, ","))

	if unit.Global != "" {
		unit.setKeyOnEtcd("global", unit.Global)
		EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), unit.Name)
	}

	unit.onEtcd = true
}

// Destroy destroys the given unit
func (unit *Unit) Destroy() {
	log.Println("Destroying unit", unit.Name)
	if unit.runCancel != nil {
		unit.runCancel()
	}

	unit.Stop() // just making sure

	FS.Remove(unitPath + unit.Name)
	unit.onDisk = false
	DBusConnection.Reload()

	if unit.Name == "" {
		// oopsie
		unit.onEtcd = false //probably not
		return
	}

	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/units/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/machines/%s/units/%s", Config.Zone, unit.Machine, unit.Name), etcd.WithPrefix())
	if unit.Global != "" {
		EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	}
	unit.etcdCache = etcdcache.New() // clear out all old cache!
	unit.onEtcd = false
}

// LoadAndWatch loads the unit to the system and follows the desired state
func (unit *Unit) LoadAndWatch() {
	unit.Create()
	unit.becomeDesiredState()
	go unit.Watch()
}

func (unit *Unit) becomeDesiredState() {
	fmt.Println("desiredState:", unit.DesiredState)
	if unit.DesiredState == state.Active {
		unit.Start()
	} else if unit.DesiredState == state.Dead {
		unit.Stop()
	} else if unit.DesiredState == state.Destroy {
		unit.Destroy()
	}
}

// Watch creates and etcd watcher for the desired state of a specific unit
func (unit *Unit) Watch() {
	chans := EtcdAPI.Watch(unit.runContext, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name), etcd.WithPrefix())
	for resp := range chans {
		for _, ev := range resp.Events {
			if ev.IsModify() || ev.IsCreate() {
				unit.DesiredState = state.ForString(string(ev.Kv.Value))
				unit.becomeDesiredState()
			}
			if ev.Type == mvccpb.DELETE {
				break
			}
		}
	}
}

// WaitOnDestroy waits for the unit to enter a destroyed state
func (unit *Unit) WaitOnDestroy() {
	unit.disableCache = true
	for unit.getKeyFromEtcd("name") != "" {
		time.Sleep(100 * time.Millisecond)
	}
}

// WaitOnState waits for the unit to enter a specific state
func (unit *Unit) WaitOnState(s state.State) {
	stateInfo, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name))
	if err != nil || stateInfo.Count == 0 { // hmmm
		return
	}
	if string(stateInfo.Kvs[0].Value) == s.String() {
		return
	}

	chans := EtcdAPI.Watch(context.Background(), fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name), etcd.WithPrefix())
	for resp := range chans {
		for _, ev := range resp.Events {
			if string(ev.Kv.Value) == s.String() {
				break
			}
		}
	}
}

// SetState sets the state of the unit and updates etcd
func (unit *Unit) SetState(s state.State) {
	if unit.Global != "" || !unit.isHealthy() {
		return
	}
	unit.State = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name), s.String())
}

// SetDesiredState sets the desiredstate of the unit and updates etcd
func (unit *Unit) SetDesiredState(s state.State) {
	unit.DesiredState = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name), s.String())
}

func (unit *Unit) isHealthy() bool {
	unit.disableCache = true
	name := unit.getKeyFromEtcd("name")
	if name == "" {
		unit.Destroy()
		return false
	}
	return true
}

func (unit *Unit) getKeyFromEtcd(key string) string {
	if unit.etcdCache != nil {
		if kv, err := unit.etcdCache.Get(key); err == nil {
			return string(kv.Value)
		}
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	response, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key))
	if err != nil || response.Count == 0 {
		return ""
	}
	return string(response.Kvs[0].Value)
}

func (unit *Unit) setKeyOnEtcd(key, content string) {
	if unit.etcdCache != nil {
		unit.etcdCache.Invalidate(key)
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key), content)
}

func setUpDBus() {
	if DBusConnection == nil {
		var err error
		DBusConnection, err = dbus.NewSystemdConnection()
		if err != nil {
			panic(err)
		}
	}
}

// KillAllOldUnits makes sure all old Dispatch spawned unit files on the system are deleted
func KillAllOldUnits() {
	setUpDBus()

	FS.MkdirAll(unitPath, 0755) // maybe we have a first run
	files, err := afero.Afero{Fs: FS}.ReadDir(unitPath)
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		log.Println("Stopping unit", file.Name())
		DBusConnection.KillUnit(file.Name(), 9) // do we care at this point?
		FS.Remove(unitPath + file.Name())
	}

	DBusConnection.Reload()
}
<commit_msg>fix infinite destroy loop<commit_after>package unit

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/innovate-technologies/Dispatch/dispatchd/etcdcache"

	"strconv"

	"github.com/spf13/afero"

	"github.com/innovate-technologies/Dispatch/dispatchd/config"
	"github.com/innovate-technologies/Dispatch/dispatchd/etcdclient"
	"github.com/innovate-technologies/Dispatch/dispatchd/unit/state"
	"github.com/innovate-technologies/Dispatch/interfaces"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"github.com/coreos/go-systemd/dbus"
	"golang.org/x/net/context"
)

const unitPath = "/var/run/dispatch/"

var (
	// Config is a pointer need to be set to the main configuration
	Config *config.ConfigurationInfo
	ctx    = context.Background()
	//EtcdAPI is the etcd keys api
	EtcdAPI interfaces.EtcdAPI = etcdclient.GetEtcdv3()
	// DBusConnection is the connection to the system's D-Bus
	DBusConnection DBusConnectionInterface
	// FS is the file system to be used
	FS = afero.NewOsFs()
)

// UnitInterface is the interface to a Unit
type UnitInterface interface {
	Start()
	Stop()
	Restart()
	Create()
	PutOnQueue()
	SaveOnEtcd()
	Destroy()
	LoadAndWatch()
	Watch()
	SetState(state.State)
	SetDesiredState(s state.State)
}

// Unit is a struct containing all info of a specific unit
type Unit struct {
	Name         string `json:"name"`
	Machine      string `json:"machine"`
	Template     string `json:"template,omitempty"` // is set with template name if from Template
	Global       string `json:"global,omitempty"`   // is set with global name if from global
	State        state.State
	DesiredState state.State
	Ports        []int64 `json:"ports"`
	Constraints  map[string]string
	UnitContent  string `json:"unitContent"`
	onEtcd       bool
	onDisk       bool
	etcdName     string
	etcdCache    *etcdcache.EtcdCache
	disableCache bool
	runContext   context.Context
	runCancel    context.CancelFunc
}

// GetAll returns all units in our zone
func GetAll() ([]Unit, error) {
	cache, err := etcdcache.NewForPrefix(fmt.Sprintf("/dispatch/%s/units", Config.Zone))
	if err != nil {
		return nil, err
	}

	units := []Unit{}
	hadUnitNames := map[string]bool{}

	for _, kv := range cache.GetAll() {
		pathParts := strings.Split(string(kv.Key), "/")
		if _, ok := hadUnitNames[pathParts[4]]; !ok {
			units = append(units, NewFromEtcd(pathParts[4]))
			hadUnitNames[pathParts[4]] = true
		}

	}

	return units, nil
}

// New returns a new Unit
func New() Unit {
	setUpDBus()
	return Unit{}
}

// NewFromEtcd creates a new unit with info from etcd
func NewFromEtcd(name string) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}

	unit := New()
	fillFromEtcd(&unit, name)

	return unit
}

// NewFromEtcdWithCache creates a new unit with info from etcd using a specified cache
func NewFromEtcdWithCache(name string, cache *etcdcache.EtcdCache) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}
	unit := New()
	unit.etcdCache = cache
	fillFromEtcd(&unit, name)

	return unit
}

func fillFromEtcd(unit *Unit, name string) {
	unit.etcdName = name
	unit.onEtcd = true
	unit.Name = unit.getKeyFromEtcd("name")
	unit.Machine = unit.getKeyFromEtcd("machine")
	unit.Template = unit.getKeyFromEtcd("template")
	unit.Global = unit.getKeyFromEtcd("global")
	unit.UnitContent = unit.getKeyFromEtcd("unit")
	unit.DesiredState = state.ForString(unit.getKeyFromEtcd("desiredState"))
	unit.State = state.ForString(unit.getKeyFromEtcd("state"))

	unit.Ports = []int64{}
	portsStringArray := strings.Split(unit.getKeyFromEtcd("ports"), ",")
	for _, portString := range portsStringArray {
		port, err := strconv.ParseInt(portString, 10, 64)
		if err == nil {
			unit.Ports = append(unit.Ports, port)
		}
	}
}

// Start starts the specific unit
func (unit *Unit) Start() {
	log.Println("Starting unit", unit.Name)
	unit.SetState(state.Starting)
	c := make(chan string)
	_, err := DBusConnection.StartUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error starting unit", unit.Name, err)
		return
	}
	<-c
	unit.SetState(state.Active)
}

// Stop stops the unit
func (unit *Unit) Stop() {
	log.Println("Stopping unit", unit.Name)
	c := make(chan string)
	_, err := DBusConnection.StopUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error stopping unit", unit.Name, err)
		log.Println("Killing unit", unit.Name)
		DBusConnection.KillUnit(unit.Name, 9) // the big guns!
		unit.SetState(state.Dead)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Stopped unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Restart restarts the unit
func (unit *Unit) Restart() {
	unit.Stop()
	unit.Start()
}

// Create writes the unit to the disk
func (unit *Unit) Create() {
	if unit.Name == "" {
		log.Println("Error starting unit with no name")
		if unit.etcdName != "" {
			// Faulty unit
			unit.Name = unit.etcdName
			unit.Destroy()
		}
		return // can't create file without a name
	}
	unit.runContext, unit.runCancel = context.WithCancel(context.Background())

	thisUnitPath := unitPath + unit.Name

	fileContent := []byte(unit.UnitContent)

	FS.MkdirAll(unitPath, 0755)

	FS.Remove(thisUnitPath) // make sure the old unit is gone
	file, err := FS.Create(thisUnitPath)
	if err != nil {
		panic(err)
	}
	_, err = file.Write(fileContent)
	if err != nil {
		panic(err)
	}

	c := make(chan string)
	_, stopErr := DBusConnection.StopUnit(unit.Name, "fail", c) // stop unit to make sure new one is loaded
	if stopErr == nil {
		<-c // wait on completion
	}

	unit.onDisk = true
	_, dberr := DBusConnection.LinkUnitFiles([]string{thisUnitPath}, true, true)
	fmt.Println(dberr)
	DBusConnection.Reload()
}

// PutOnQueue places a specific unit on the queue
func (unit *Unit) PutOnQueue() {
	if unit.Global != "" {
		return
	}
	log.Println("Placing", unit.Name, "on queue")
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), unit.Name)
}

func (unit *Unit) normalizeName() {
	if !strings.HasSuffix(unit.Name, ".service") {
		unit.Name += ".service"
	}
}

// SaveOnEtcd saves the unit to etcd
func (unit *Unit) SaveOnEtcd() {
	unit.normalizeName()

	log.Println("Saving", unit.Name, "to etcd")

	unit.setKeyOnEtcd("name", unit.Name)
	unit.setKeyOnEtcd("unit", unit.UnitContent)
	unit.setKeyOnEtcd("template", unit.Template)
	unit.setKeyOnEtcd("desiredState", unit.DesiredState.String())

	portStrings := []string{}
	for _, port := range unit.Ports {
		portStrings = append(portStrings, strconv.FormatInt(port, 10))
	}
	unit.setKeyOnEtcd("ports", strings.Join(portStrings, ","))

	if unit.Global != "" {
		unit.setKeyOnEtcd("global", unit.Global)
		EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), unit.Name)
	}

	unit.onEtcd = true
}

// Destroy destroys the given unit
func (unit *Unit) Destroy() {
	log.Println("Destroying unit", unit.Name)
	if unit.runCancel != nil {
		unit.runCancel()
	}

	unit.Stop() // just making sure
	fmt.Println("unit stopped for destroy")

	FS.Remove(unitPath + unit.Name)
	unit.onDisk = false
	DBusConnection.Reload()

	if unit.Name == "" {
		fmt.Println("Destroying unit not on etcd")
		// oopsie
		unit.onEtcd = false //probably not
		return
	}

	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/units/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/machines/%s/units/%s", Config.Zone, unit.Machine, unit.Name), etcd.WithPrefix())
	if unit.Global != "" {
		EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	}
	unit.etcdCache = etcdcache.New() // clear out all old cache!
	unit.onEtcd = false
	fmt.Println("Destroy done")
}

// LoadAndWatch loads the unit to the system and follows the desired state
func (unit *Unit) LoadAndWatch() {
	unit.Create()
	unit.becomeDesiredState()
	go unit.Watch()
}

func (unit *Unit) becomeDesiredState() {
	fmt.Println("desiredState:", unit.DesiredState)
	if unit.DesiredState == state.Active {
		unit.Start()
	} else if unit.DesiredState == state.Dead {
		unit.Stop()
	} else if unit.DesiredState == state.Destroy {
		fmt.Println("desiredState: found to be destroy")
		unit.Destroy()
	}
}

// Watch creates and etcd watcher for the desired state of a specific unit
func (unit *Unit) Watch() {
	chans := EtcdAPI.Watch(unit.runContext, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name))
	for resp := range chans {
		for _, ev := range resp.Events {
			if ev.IsModify() || ev.IsCreate() {
				unit.DesiredState = state.ForString(string(ev.Kv.Value))
				unit.becomeDesiredState()
			}
			if ev.Type == mvccpb.DELETE {
				break
			}
		}
	}
	fmt.Println("End watch for ", unit.Name)
}

// WaitOnDestroy waits for the unit to enter a destroyed state
func (unit *Unit) WaitOnDestroy() {
	unit.disableCache = true
	for unit.getKeyFromEtcd("name") != "" {
		time.Sleep(100 * time.Millisecond)
	}
}

// WaitOnState waits for the unit to enter a specific state
func (unit *Unit) WaitOnState(s state.State) {
	stateInfo, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name))
	if err != nil || stateInfo.Count == 0 { // hmmm
		return
	}
	if string(stateInfo.Kvs[0].Value) == s.String() {
		return
	}

	chans := EtcdAPI.Watch(context.Background(), fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name))
	for resp := range chans {
		for _, ev := range resp.Events {
			if string(ev.Kv.Value) == s.String() {
				break
			}
		}
	}
}

// SetState sets the state of the unit and updates etcd
func (unit *Unit) SetState(s state.State) {
	if unit.Global != "" || !unit.isHealthy() {
		return
	}
	unit.State = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name), s.String())
}

// SetDesiredState sets the desiredstate of the unit and updates etcd
func (unit *Unit) SetDesiredState(s state.State) {
	unit.DesiredState = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name), s.String())
}

func (unit *Unit) isHealthy() bool {
	unit.disableCache = true
	name := unit.getKeyFromEtcd("name")
	if name == "" {
		return false
	}
	return true
}

func (unit *Unit) getKeyFromEtcd(key string) string {
	if unit.etcdCache != nil {
		if kv, err := unit.etcdCache.Get(key); err == nil {
			return string(kv.Value)
		}
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	response, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key))
	if err != nil || response.Count == 0 {
		return ""
	}
	return string(response.Kvs[0].Value)
}

func (unit *Unit) setKeyOnEtcd(key, content string) {
	if unit.etcdCache != nil {
		unit.etcdCache.Invalidate(key)
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key), content)
}

func setUpDBus() {
	if DBusConnection == nil {
		var err error
		DBusConnection, err = dbus.NewSystemdConnection()
		if err != nil {
			panic(err)
		}
	}
}

// KillAllOldUnits makes sure all old Dispatch spawned unit files on the system are deleted
func KillAllOldUnits() {
	setUpDBus()

	FS.MkdirAll(unitPath, 0755) // maybe we have a first run
	files, err := afero.Afero{Fs: FS}.ReadDir(unitPath)
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		log.Println("Stopping unit", file.Name())
		DBusConnection.KillUnit(file.Name(), 9) // do we care at this point?
		FS.Remove(unitPath + file.Name())
	}

	DBusConnection.Reload()
}
<|endoftext|>


<commit_before>// Package activitystreams provides all the basic ActivityStreams
// implementation needed for Write.as.
package activitystreams

import (
	"time"
)

const (
	Namespace = "https://www.w3.org/ns/activitystreams"
	toPublic  = "https://www.w3.org/ns/activitystreams#Public"
)

// Activity describes actions that have either already occurred, are in the
// process of occurring, or may occur in the future.
type Activity struct {
	BaseObject
	Actor     string    `json:"actor"`
	Published time.Time `json:"published,omitempty"`
	To        []string  `json:"to,omitempty"`
	CC        []string  `json:"cc,omitempty"`
	Object    *Object   `json:"object"`
}

// NewCreateActivity builds a basic Create activity that includes the given
// Object and the Object's AttributedTo property as the Actor.
func NewCreateActivity(o *Object) *Activity {
	a := Activity{
		BaseObject: BaseObject{
			Context: []interface{}{
				Namespace,
			},
			ID:   o.ID,
			Type: "Create",
		},
		Actor:  o.AttributedTo,
		Object: o,
	}
	return &a
}

// NewDeleteActivity builds a basic Delete activity that includes the given
// Object and the Object's AttributedTo property as the Actor.
func NewDeleteActivity(o *Object) *Activity {
	a := Activity{
		BaseObject: BaseObject{
			Context: []interface{}{
				Namespace,
			},
			ID:   o.ID,
			Type: "Delete",
		},
		Actor:  o.AttributedTo,
		Object: o,
	}
	return &a
}

// Object is the primary base type for the Activity Streams vocabulary.
type Object struct {
	BaseObject
	Published    time.Time         `json:"published"`
	Summary      *string           `json:"summary,omitempty"`
	InReplyTo    *string           `json:"inReplyTo"`
	URL          string            `json:"url"`
	AttributedTo string            `json:"attributedTo"`
	To           []string          `json:"to"`
	CC           []string          `json:"cc,omitempty"`
	Name         string            `json:"name,omitempty"`
	Content      string            `json:"content"`
	ContentMap   map[string]string `json:"contentMap,omitempty"`
}

// NewNoteObject creates a basic Note object that includes the public
// namespace in IRIs it's addressed to.
func NewNoteObject() *Object {
	o := Object{
		BaseObject: BaseObject{
			Type: "Note",
		},
		To: []string{
			toPublic,
		},
	}
	return &o
}

// NewArticleObject creates a basic Article object that includes the public
// namespace in IRIs it's addressed to.
func NewArticleObject() *Object {
	o := Object{
		BaseObject: BaseObject{
			Type: "Article",
		},
		To: []string{
			toPublic,
		},
	}
	return &o
}
<commit_msg>Use object's Published date in Create activity<commit_after>// Package activitystreams provides all the basic ActivityStreams
// implementation needed for Write.as.
package activitystreams

import (
	"time"
)

const (
	Namespace = "https://www.w3.org/ns/activitystreams"
	toPublic  = "https://www.w3.org/ns/activitystreams#Public"
)

// Activity describes actions that have either already occurred, are in the
// process of occurring, or may occur in the future.
type Activity struct {
	BaseObject
	Actor     string    `json:"actor"`
	Published time.Time `json:"published,omitempty"`
	To        []string  `json:"to,omitempty"`
	CC        []string  `json:"cc,omitempty"`
	Object    *Object   `json:"object"`
}

// NewCreateActivity builds a basic Create activity that includes the given
// Object and the Object's AttributedTo property as the Actor.
func NewCreateActivity(o *Object) *Activity {
	a := Activity{
		BaseObject: BaseObject{
			Context: []interface{}{
				Namespace,
			},
			ID:   o.ID,
			Type: "Create",
		},
		Actor:     o.AttributedTo,
		Object:    o,
		Published: o.Published,
	}
	return &a
}

// NewDeleteActivity builds a basic Delete activity that includes the given
// Object and the Object's AttributedTo property as the Actor.
func NewDeleteActivity(o *Object) *Activity {
	a := Activity{
		BaseObject: BaseObject{
			Context: []interface{}{
				Namespace,
			},
			ID:   o.ID,
			Type: "Delete",
		},
		Actor:  o.AttributedTo,
		Object: o,
	}
	return &a
}

// Object is the primary base type for the Activity Streams vocabulary.
type Object struct {
	BaseObject
	Published    time.Time         `json:"published"`
	Summary      *string           `json:"summary,omitempty"`
	InReplyTo    *string           `json:"inReplyTo"`
	URL          string            `json:"url"`
	AttributedTo string            `json:"attributedTo"`
	To           []string          `json:"to"`
	CC           []string          `json:"cc,omitempty"`
	Name         string            `json:"name,omitempty"`
	Content      string            `json:"content"`
	ContentMap   map[string]string `json:"contentMap,omitempty"`
}

// NewNoteObject creates a basic Note object that includes the public
// namespace in IRIs it's addressed to.
func NewNoteObject() *Object {
	o := Object{
		BaseObject: BaseObject{
			Type: "Note",
		},
		To: []string{
			toPublic,
		},
	}
	return &o
}

// NewArticleObject creates a basic Article object that includes the public
// namespace in IRIs it's addressed to.
func NewArticleObject() *Object {
	o := Object{
		BaseObject: BaseObject{
			Type: "Article",
		},
		To: []string{
			toPublic,
		},
	}
	return &o
}
<|endoftext|>


<commit_before>// Basic data structures for Go.
package main

import (
	"fmt"
)

// Letters for quick lookup.
// These are the letters marked on top of the board.  "I" is
// always left out, likely to avoid confusion with "1".
const BoardLetters = "ABCDEFGHJKLMNOPQRST"
const ActionChars = "cbw"

// Here it would be nice if the type system allows us to say
// "integer with one of these value: 0, 1, 2.
type Stone int

const (
	// Stone colors. Used for board state.
	Empty = 0 // No stone here.
	Black = 1
	White = 2

	// Moves. These correspond to board state making it easy to
	// translate between the two. (A capture will clear that position on
	// the board.)
	Capture    = 0
	PlaceBlack = 1
	PlaceWhite = 2
)

// Position where a stone can be placed.
type Pos struct {
	X int // Left to right, 0-18 (marked as 1-19 on board.)
	Y int // Bottom to top, 0-18 (ABCDEFGHJKLMNOPQRST.)
}

type BoardSize struct {
	Width  int
	Height int
}

// A move. (Place or capture stone.)
type Move struct {
	Action Stone // (Hmm, this looks a bit weird here.)
	Pos
}

func (move Move) String() string {
	// Todo: check if values are out of range?
	letter := BoardLetters[move.X]
	number := move.Y + 1
	stoneColor := ActionChars[move.Action]

	return fmt.Sprintf("%c%d%c", letter, number, stoneColor)
}

// The board is just a cache of moves. We can recompute the board from
// the moves.
type Game struct {
	BoardSize
	Board map[Pos]Stone
	Moves []Move
	// Here we should have a stack (or slice) of boards so that we have
	// a history of previous board states.

	// + We need to keep score somehow. (Number of captured stones for
	// each player.)
}

func (game *Game) AddMove(move Move) {
	game.Moves = append(game.Moves, move)
	game.Board[move.Pos] = move.Action
}

func (game *Game) AddMoves(moves []Move) {
	for _, move := range moves {
		game.AddMove(move)
	}
}

func (game *Game) PrintBoard() {
	stoneChars := ".bw"

	fmt.Println("   " + BoardLetters)

	for y := game.Height - 1; y >= 0; y-- {
		line := ""
		for x := 0; x < game.Width; x++ {
			line += fmt.Sprintf("%c", stoneChars[game.Board[Pos{x, y}]])
		}
		fmt.Printf("%2d %s\n", y+1, line)
	}
}

func NewGame(boardSize BoardSize) Game {
	return Game{
		BoardSize: boardSize,
		Board:     make(map[Pos]Stone),
	}
}

func main() {
	game := NewGame(BoardSize{19, 19})
	moves := []Move{{PlaceBlack, Pos{0, 0}}, {PlaceWhite, Pos{1, 6}}}
	fmt.Println()
	fmt.Println("Moves:", moves)
	fmt.Println()
	game.AddMoves(moves)
	game.PrintBoard()

	// fmt.Println(game.Board[Pos{}])
	// fmt.Println(game.Board[Pos{1, 1}])

	// fmt.Println(moves)
}
<commit_msg>Added ParseMove() (unfinished) and FormatMove().<commit_after>// Basic data structures for Go.
package main

import (
	"errors"
	"fmt"
)

// Letters for quick lookup.
// These are the letters marked on top of the board.  "I" is
// always left out, likely to avoid confusion with "1".
const BoardLetters = "ABCDEFGHJKLMNOPQRST"
const ActionChars = "cbw"

// Here it would be nice if the type system allows us to say
// "integer with one of these value: 0, 1, 2.
type Stone int

const (
	// Stone colors. Used for board state.
	Empty = 0 // No stone here.
	Black = 1
	White = 2

	// Moves. These correspond to board state making it easy to
	// translate between the two. (A capture will clear that position on
	// the board.)
	Capture    = 0
	PlaceBlack = 1
	PlaceWhite = 2
)

// Position where a stone can be placed.
type Pos struct {
	X int // Left to right, 0-18 (marked as 1-19 on board.)
	Y int // Bottom to top, 0-18 (ABCDEFGHJKLMNOPQRST.)
}

type BoardSize struct {
	Width  int
	Height int
}

// A move. (Place or capture stone.)
type Move struct {
	Action Stone // (Hmm, this looks a bit weird here.)
	Pos
}

func (move Move) String() string {
	return FormatMove(move)
}

func ParseMove(text string) (Move, error) {
	// letter := text[0]
	// number := text[1:len(text)-1]
	// action := text[len(text)-1]

	return Move{}, errors.New("Parser not written")
}

func FormatMove(move Move) string {
	// Todo: check if values are out of range?
	letter := BoardLetters[move.X]
	number := move.Y + 1
	stoneColor := ActionChars[move.Action]

	return fmt.Sprintf("%c%d%c", letter, number, stoneColor)
}

// The board is just a cache of moves. We can recompute the board from
// the moves.
type Game struct {
	BoardSize
	Board map[Pos]Stone
	Moves []Move
	// Here we should have a stack (or slice) of boards so that we have
	// a history of previous board states.

	// + We need to keep score somehow. (Number of captured stones for
	// each player.)
}

func (game *Game) AddMove(move Move) {
	game.Moves = append(game.Moves, move)
	game.Board[move.Pos] = move.Action
}

func (game *Game) AddMoves(moves []Move) {
	for _, move := range moves {
		game.AddMove(move)
	}
}

func (game *Game) PrintBoard() {
	stoneChars := ".bw"

	fmt.Println("   " + BoardLetters)

	for y := game.Height - 1; y >= 0; y-- {
		line := ""
		for x := 0; x < game.Width; x++ {
			line += fmt.Sprintf("%c", stoneChars[game.Board[Pos{x, y}]])
		}
		fmt.Printf("%2d %s\n", y+1, line)
	}
}

func NewGame(boardSize BoardSize) Game {
	return Game{
		BoardSize: boardSize,
		Board:     make(map[Pos]Stone),
	}
}

func main() {
	game := NewGame(BoardSize{19, 19})
	moves := []Move{{PlaceBlack, Pos{0, 0}}, {PlaceWhite, Pos{1, 6}}}
	fmt.Println()
	fmt.Println("Moves:", moves)
	fmt.Println()
	game.AddMoves(moves)
	game.PrintBoard()

	// fmt.Println(game.Board[Pos{}])
	// fmt.Println(game.Board[Pos{1, 1}])

	// fmt.Println(moves)
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2019 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package lib

import (
	"math/big"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestExecutionSegments(t *testing.T) {
	t.Parallel()

	var nilEs *ExecutionSegment
	fullEs, err := NewExecutionSegment(big.NewRat(0, 1), big.NewRat(1, 1))
	require.NoError(t, err)
	require.Equal(t, true, nilEs.Equal(fullEs))
	//TODO: more tests
}

func TestExecutionSegmentSplit(t *testing.T) {
	t.Parallel()

	var nilEs *ExecutionSegment
	_, err := nilEs.Split(-1)
	require.Error(t, err)

	_, err = nilEs.Split(0)
	require.Error(t, err)

	segments, err := nilEs.Split(1)
	require.NoError(t, err)
	require.Len(t, segments, 1)
	assert.Equal(t, "0:1", segments[0].String())

	segments, err = nilEs.Split(2)
	require.NoError(t, err)
	require.Len(t, segments, 2)
	assert.Equal(t, "0:1/2", segments[0].String())
	assert.Equal(t, "1/2:1", segments[1].String())

	segments, err = nilEs.Split(3)
	require.NoError(t, err)
	require.Len(t, segments, 3)
	assert.Equal(t, "0:1/3", segments[0].String())
	assert.Equal(t, "1/3:2/3", segments[1].String())
	assert.Equal(t, "2/3:1", segments[2].String())

	secondQuarter, err := NewExecutionSegment(big.NewRat(1, 4), big.NewRat(2, 4))
	require.NoError(t, err)

	segments, err = secondQuarter.Split(1)
	require.NoError(t, err)
	require.Len(t, segments, 1)
	assert.Equal(t, "1/4:1/2", segments[0].String())

	segments, err = secondQuarter.Split(2)
	require.NoError(t, err)
	require.Len(t, segments, 2)
	assert.Equal(t, "1/4:3/8", segments[0].String())
	assert.Equal(t, "3/8:1/2", segments[1].String())

	segments, err = secondQuarter.Split(3)
	require.NoError(t, err)
	require.Len(t, segments, 3)
	assert.Equal(t, "1/4:1/3", segments[0].String())
	assert.Equal(t, "1/3:5/12", segments[1].String())
	assert.Equal(t, "5/12:1/2", segments[2].String())

	segments, err = secondQuarter.Split(4)
	require.NoError(t, err)
	require.Len(t, segments, 4)
	assert.Equal(t, "1/4:5/16", segments[0].String())
	assert.Equal(t, "5/16:3/8", segments[1].String())
	assert.Equal(t, "3/8:7/16", segments[2].String())
	assert.Equal(t, "7/16:1/2", segments[3].String())
}

//TODO: a lot of other tests
<commit_msg>Test everything function of lib.ExecutionSegment<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2019 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package lib

import (
	"math/big"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func stringToES(t *testing.T, str string) *ExecutionSegment {
	var es = new(ExecutionSegment)
	require.NoError(t, es.UnmarshalText([]byte(str)))
	return es
}
func TestExecutionSegmentEquals(t *testing.T) {
	t.Parallel()

	t.Run("nil segment to full", func(t *testing.T) {
		var nilEs *ExecutionSegment
		fullEs := stringToES(t, "0:1")
		require.True(t, nilEs.Equal(fullEs))
		require.True(t, fullEs.Equal(nilEs))
	})

	t.Run("To it's self", func(t *testing.T) {
		var es = stringToES(t, "1/2:2/3")
		require.True(t, es.Equal(es))
	})
}

func TestExecutionSegmentNew(t *testing.T) {
	t.Parallel()
	t.Run("from is below zero", func(t *testing.T) {
		_, err := NewExecutionSegment(big.NewRat(-1, 1), big.NewRat(1, 1))
		require.Error(t, err)
	})
	t.Run("to is more than 1", func(t *testing.T) {
		_, err := NewExecutionSegment(big.NewRat(0, 1), big.NewRat(2, 1))
		require.Error(t, err)
	})
	t.Run("from is smaller than to", func(t *testing.T) {
		_, err := NewExecutionSegment(big.NewRat(1, 2), big.NewRat(1, 3))
		require.Error(t, err)
	})

	t.Run("from is equal to 'to'", func(t *testing.T) {
		_, err := NewExecutionSegment(big.NewRat(1, 2), big.NewRat(1, 2))
		require.Error(t, err)
	})
	t.Run("ok", func(t *testing.T) {
		_, err := NewExecutionSegment(big.NewRat(0, 1), big.NewRat(1, 1))
		require.NoError(t, err)
	})
}

func TestExecutionSegmentUnmarshalText(t *testing.T) {
	t.Parallel()
	var testCases = []struct {
		input  string
		output *ExecutionSegment
		isErr  bool
	}{
		{input: "0:1", output: &ExecutionSegment{from: zeroRat, to: oneRat}},
		{input: "0.5:0.75", output: &ExecutionSegment{from: big.NewRat(1, 2), to: big.NewRat(3, 4)}},
		{input: "1/2:3/4", output: &ExecutionSegment{from: big.NewRat(1, 2), to: big.NewRat(3, 4)}},
		{input: "50%:75%", output: &ExecutionSegment{from: big.NewRat(1, 2), to: big.NewRat(3, 4)}},
		{input: "2/4:75%", output: &ExecutionSegment{from: big.NewRat(1, 2), to: big.NewRat(3, 4)}},
		{input: "75%", output: &ExecutionSegment{from: zeroRat, to: big.NewRat(3, 4)}},
		{input: "125%", isErr: true},
		{input: "1a5%", isErr: true},
		{input: "1a5", isErr: true},
		{input: "1a5%:2/3", isErr: true},
		{input: "125%:250%", isErr: true},
		{input: "55%:50%", isErr: true},
		// TODO add more strange or not so strange cases
	}
	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.input, func(t *testing.T) {
			var es = new(ExecutionSegment)
			err := es.UnmarshalText([]byte(testCase.input))
			if testCase.isErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
			require.True(t, es.Equal(testCase.output))

			// see if unmarshalling a stringified segment gets you back the same segment
			err = es.UnmarshalText([]byte(es.String()))
			require.NoError(t, err)
			require.True(t, es.Equal(testCase.output))
		})
	}

	t.Run("Unmarshal nilSegment.String", func(t *testing.T) {
		var nilEs *ExecutionSegment
		var nilEsStr = nilEs.String()
		require.Equal(t, "0:1", nilEsStr)

		var es = new(ExecutionSegment)
		err := es.UnmarshalText([]byte(nilEsStr))
		require.NoError(t, err)
		require.True(t, es.Equal(nilEs))
	})
}

func TestExecutionSegmentSplit(t *testing.T) {
	t.Parallel()

	var nilEs *ExecutionSegment
	_, err := nilEs.Split(-1)
	require.Error(t, err)

	_, err = nilEs.Split(0)
	require.Error(t, err)

	segments, err := nilEs.Split(1)
	require.NoError(t, err)
	require.Len(t, segments, 1)
	assert.Equal(t, "0:1", segments[0].String())

	segments, err = nilEs.Split(2)
	require.NoError(t, err)
	require.Len(t, segments, 2)
	assert.Equal(t, "0:1/2", segments[0].String())
	assert.Equal(t, "1/2:1", segments[1].String())

	segments, err = nilEs.Split(3)
	require.NoError(t, err)
	require.Len(t, segments, 3)
	assert.Equal(t, "0:1/3", segments[0].String())
	assert.Equal(t, "1/3:2/3", segments[1].String())
	assert.Equal(t, "2/3:1", segments[2].String())

	secondQuarter, err := NewExecutionSegment(big.NewRat(1, 4), big.NewRat(2, 4))
	require.NoError(t, err)

	segments, err = secondQuarter.Split(1)
	require.NoError(t, err)
	require.Len(t, segments, 1)
	assert.Equal(t, "1/4:1/2", segments[0].String())

	segments, err = secondQuarter.Split(2)
	require.NoError(t, err)
	require.Len(t, segments, 2)
	assert.Equal(t, "1/4:3/8", segments[0].String())
	assert.Equal(t, "3/8:1/2", segments[1].String())

	segments, err = secondQuarter.Split(3)
	require.NoError(t, err)
	require.Len(t, segments, 3)
	assert.Equal(t, "1/4:1/3", segments[0].String())
	assert.Equal(t, "1/3:5/12", segments[1].String())
	assert.Equal(t, "5/12:1/2", segments[2].String())

	segments, err = secondQuarter.Split(4)
	require.NoError(t, err)
	require.Len(t, segments, 4)
	assert.Equal(t, "1/4:5/16", segments[0].String())
	assert.Equal(t, "5/16:3/8", segments[1].String())
	assert.Equal(t, "3/8:7/16", segments[2].String())
	assert.Equal(t, "7/16:1/2", segments[3].String())
}

func TestExecutionSegmentScale(t *testing.T) {
	t.Parallel()
	var es = new(ExecutionSegment)
	require.NoError(t, es.UnmarshalText([]byte("0.5")))
	require.Equal(t, int64(1), es.Scale(2))
	require.Equal(t, int64(2), es.Scale(3))

	require.NoError(t, es.UnmarshalText([]byte("0.5:1.0")))
	require.Equal(t, int64(1), es.Scale(2))
	require.Equal(t, int64(1), es.Scale(3))
}

func TestExecutionSegmentCopyScaleRat(t *testing.T) {
	t.Parallel()
	var es = new(ExecutionSegment)
	var twoRat = big.NewRat(2, 1)
	var threeRat = big.NewRat(3, 1)
	require.NoError(t, es.UnmarshalText([]byte("0.5")))
	require.Equal(t, oneRat, es.CopyScaleRat(twoRat))
	require.Equal(t, big.NewRat(3, 2), es.CopyScaleRat(threeRat))

	require.NoError(t, es.UnmarshalText([]byte("0.5:1.0")))
	require.Equal(t, oneRat, es.CopyScaleRat(twoRat))
	require.Equal(t, big.NewRat(3, 2), es.CopyScaleRat(threeRat))

	var nilEs *ExecutionSegment
	require.Equal(t, twoRat, nilEs.CopyScaleRat(twoRat))
	require.Equal(t, threeRat, nilEs.CopyScaleRat(threeRat))
}

func TestExecutionSegmentInPlaceScaleRat(t *testing.T) {
	t.Parallel()
	var es = new(ExecutionSegment)
	var twoRat = big.NewRat(2, 1)
	var threeRat = big.NewRat(3, 1)
	var threeSecondsRat = big.NewRat(3, 2)
	require.NoError(t, es.UnmarshalText([]byte("0.5")))
	require.Equal(t, oneRat, es.InPlaceScaleRat(twoRat))
	require.Equal(t, oneRat, twoRat)
	require.Equal(t, threeSecondsRat, es.InPlaceScaleRat(threeRat))
	require.Equal(t, threeSecondsRat, threeRat)

	es = stringToES(t, "0.5:1.0")
	twoRat = big.NewRat(2, 1)
	threeRat = big.NewRat(3, 1)
	require.Equal(t, oneRat, es.InPlaceScaleRat(twoRat))
	require.Equal(t, oneRat, twoRat)
	require.Equal(t, threeSecondsRat, es.InPlaceScaleRat(threeRat))
	require.Equal(t, threeSecondsRat, threeRat)

	var nilEs *ExecutionSegment
	twoRat = big.NewRat(2, 1)
	threeRat = big.NewRat(3, 1)
	require.Equal(t, big.NewRat(2, 1), nilEs.InPlaceScaleRat(twoRat))
	require.Equal(t, big.NewRat(2, 1), twoRat)
	require.Equal(t, big.NewRat(3, 1), nilEs.InPlaceScaleRat(threeRat))
	require.Equal(t, big.NewRat(3, 1), threeRat)
}

func TestExecutionSegmentSubSegment(t *testing.T) {
	t.Parallel()
	var testCases = []struct {
		name              string
		base, sub, result *ExecutionSegment
	}{
		// TODO add more strange or not so strange cases
		{
			name:   "nil base",
			base:   (*ExecutionSegment)(nil),
			sub:    stringToES(t, "0.2:0.3"),
			result: stringToES(t, "0.2:0.3"),
		},

		{
			name:   "nil sub",
			base:   stringToES(t, "0.2:0.3"),
			sub:    (*ExecutionSegment)(nil),
			result: stringToES(t, "0.2:0.3"),
		},
		{
			name:   "doc example",
			base:   stringToES(t, "1/2:1"),
			sub:    stringToES(t, "0:1/2"),
			result: stringToES(t, "1/2:3/4"),
		},
	}

	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.name, func(t *testing.T) {
			require.Equal(t, testCase.result, testCase.base.SubSegment(testCase.sub))
		})
	}
}

func TestSplitBadSegment(t *testing.T) {
	t.Parallel()
	var es = &ExecutionSegment{from: oneRat, to: zeroRat}
	_, err := es.Split(5)
	require.Error(t, err)
}

func TestSegmentExecutionFloatLength(t *testing.T) {
	t.Parallel()
	t.Run("nil has 1.0", func(t *testing.T) {
		var nilEs *ExecutionSegment
		require.Equal(t, 1.0, nilEs.FloatLength())
	})

	var testCases = []struct {
		es       *ExecutionSegment
		expected float64
	}{
		// TODO add more strange or not so strange cases
		{
			es:       stringToES(t, "1/2:1"),
			expected: 0.5,
		},
		{
			es:       stringToES(t, "1/3:1"),
			expected: 0.66666,
		},

		{
			es:       stringToES(t, "0:1/2"),
			expected: 0.5,
		},
	}

	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.es.String(), func(t *testing.T) {
			require.InEpsilon(t, testCase.expected, testCase.es.FloatLength(), 0.001)
		})
	}
}
<|endoftext|>


<commit_before>// Package dockergoal is a library to reach a set of container goals.
package dockergoal

import (
	"errors"
	"strings"

	"github.com/facebookgo/dockerutil"
	"github.com/facebookgo/errgroup"
	"github.com/facebookgo/stackerr"
	"github.com/samalba/dockerclient"
)

var (
	errNameMissing = errors.New("dockergoal: ContainerName is required")
)

// A Container defines a "desired" container state.
type Container struct {
	name                string
	containerConfig     *dockerclient.ContainerConfig
	hostConfig          *dockerclient.HostConfig
	removeExisting      bool
	forceRemoveExisting bool
	authConfig          *dockerclient.AuthConfig
}

// ContainerOption configure options for a container.
type ContainerOption func(c *Container) error

// NewContainer creates a new desired container state. This is only the desired
// state, it isn't applied until Apply is called.
func NewContainer(options ...ContainerOption) (*Container, error) {
	var c Container
	for _, o := range options {
		if err := o(&c); err != nil {
			return nil, err
		}
	}
	if c.name == "" {
		return nil, errNameMissing
	}
	return &c, nil
}

// ContainerName configures the name of the container.
func ContainerName(name string) ContainerOption {
	return func(c *Container) error {
		c.name = name
		return nil
	}
}

// ContainerRemoveExisting removes an existing container if necessary. If this
// is option is not specified, and an existing container with a different
// configuration exists, it will not be removed and an error will be returned.
// If this option is specified, the existing container will be removed and a
// new container will be created.
func ContainerRemoveExisting() ContainerOption {
	return func(c *Container) error {
		c.removeExisting = true
		return nil
	}
}

// ContainerForceRemoveExisting forces removing the existing container if one
// is found. It does so even if the existing container matches the desired
// state.
func ContainerForceRemoveExisting() ContainerOption {
	return func(c *Container) error {
		c.forceRemoveExisting = true
		return nil
	}
}

// ContainerConfig specifies the container configuration.
func ContainerConfig(config *dockerclient.ContainerConfig) ContainerOption {
	return func(c *Container) error {
		c.containerConfig = config
		return nil
	}
}

// ContainerHostConfig specifies the container host configuration.
func ContainerHostConfig(config *dockerclient.HostConfig) ContainerOption {
	return func(c *Container) error {
		c.hostConfig = config
		return nil
	}
}

// ContainerAuthConfig specifies the auth credentials used when pulling an
// image.
func ContainerAuthConfig(ac *dockerclient.AuthConfig) ContainerOption {
	return func(c *Container) error {
		c.authConfig = ac
		return nil
	}
}

// Apply creates the container, possibly removing it as necessary based on the
// container options that were set.
func (c *Container) Apply(docker dockerclient.Client) error {
	ci, err := docker.InspectContainer(c.name)

	// force remove existing
	if ci != nil && c.forceRemoveExisting {
		if err := docker.RemoveContainer(ci.Id, true, false); err != nil {
			return stackerr.Wrap(err)
		}
		// we just removed the running container and want to start a new one
		err = dockerclient.ErrNotFound
	}

	// already exists
	if err == nil {
		ok, err := c.checkRunning(docker, ci)
		if err != nil {
			return err
		}
		if ok {
			// existing container is good
			if ci.State.Running {
				// and it's already running, so we do nothing
				return nil
			}
		} else {
			// we just removed the running container and want to start a new one
			err = dockerclient.ErrNotFound
		}
	}

	// unknown error, bail
	if err != nil && err != dockerclient.ErrNotFound {
		return stackerr.Wrap(err)
	}

	// container does not exist, create it
	if err == dockerclient.ErrNotFound {
		_, err := dockerutil.CreateWithPull(docker, c.containerConfig, c.name, c.authConfig)
		if err != nil {
			return err
		}

		ci, err = docker.InspectContainer(c.name)
		if err != nil {
			return stackerr.Wrap(err)
		}
	}

	// start the container
	err = docker.StartContainer(ci.Id, c.hostConfig)
	if err != nil {
		return stackerr.Wrap(err)
	}

	return nil
}

func (c *Container) checkRunning(docker dockerclient.Client, current *dockerclient.ContainerInfo) (bool, error) {
	// image comparison is by ID, so we need to find the ID of our desired image
	desiredImageID, err := dockerutil.ImageID(docker, c.containerConfig.Image, nil)
	if err != nil {
		return false, err
	}

	if current.Image != desiredImageID {
		// if we aren't allowed to remove the existing container, consider this a failure
		if !c.removeExisting {
			return false, stackerr.Newf(
				"container %q running with image %q but desired image is %q with id %q",
				c.name,
				current.Image,
				c.containerConfig.Image,
				desiredImageID,
			)
		}

		// otherwise remove it since it isn't want we want
		if err := docker.RemoveContainer(current.Id, true, false); err != nil {
			return false, stackerr.Wrap(err)
		}

		// trigger starting a new container
		return false, nil
	}

	// we're running the correct image
	return true, nil
}

// ApplyGraph creates all the specified containers. It handles links making
// sure the dependencies are created in the right order.
func ApplyGraph(docker dockerclient.Client, containers []*Container) error {
	known := map[string]struct{}{}
	started := map[string]bool{}

	// we want to know all known names so we can detect malformed links
	for _, c := range containers {
		known[c.name] = struct{}{}
	}

	// TODO: parallel pull pass?

	// keep doing rounds of parallel starts until we're all done or error out
	pending := containers
	for len(pending) > 0 {
		var eg errgroup.Group
		var starting []string
		var nextRound []*Container

	pendingLoop:
		for _, c := range pending {
			if c.hostConfig != nil {
				// TODO: also include c.hostConfig.VolumesFrom
				for _, link := range c.hostConfig.Links {
					// only care about the name, not the alias
					parts := strings.Split(link, ":")

					// make sure the link is known
					if _, ok := known[parts[0]]; !ok {
						return stackerr.Newf("%q expects unknown link %q", c.name, link)
					}

					// we need to wait for a dependency, schedule for the next round
					if !started[parts[0]] {
						nextRound = append(nextRound, c)
						continue pendingLoop
					}
				}
			}

			starting = append(starting, c.name)
			eg.Add(1)
			go func(c *Container) {
				defer eg.Done()
				if err := c.Apply(docker); err != nil {
					eg.Error(err)
				}
			}(c)
		}

		// now wait for all the containers we started in parallel
		if err := eg.Wait(); err != nil {
			return err
		}

		// we successfully started them
		for _, n := range starting {
			started[n] = true
		}

		// move on to the next round
		pending = nextRound
	}

	return nil
}
<commit_msg>doc nit<commit_after>// Package dockergoal is a library to reach a set of container goals. It
// provides the ability to apply "deltas" in the form of recreating containers
// as necessary as the goal configuration changes.
package dockergoal

import (
	"errors"
	"strings"

	"github.com/facebookgo/dockerutil"
	"github.com/facebookgo/errgroup"
	"github.com/facebookgo/stackerr"
	"github.com/samalba/dockerclient"
)

var (
	errNameMissing = errors.New("dockergoal: ContainerName is required")
)

// A Container defines a "desired" container state.
type Container struct {
	name                string
	containerConfig     *dockerclient.ContainerConfig
	hostConfig          *dockerclient.HostConfig
	removeExisting      bool
	forceRemoveExisting bool
	authConfig          *dockerclient.AuthConfig
}

// ContainerOption configure options for a container.
type ContainerOption func(c *Container) error

// NewContainer creates a new desired container state. This is only the desired
// state, it isn't applied until Apply is called.
func NewContainer(options ...ContainerOption) (*Container, error) {
	var c Container
	for _, o := range options {
		if err := o(&c); err != nil {
			return nil, err
		}
	}
	if c.name == "" {
		return nil, errNameMissing
	}
	return &c, nil
}

// ContainerName configures the name of the container.
func ContainerName(name string) ContainerOption {
	return func(c *Container) error {
		c.name = name
		return nil
	}
}

// ContainerRemoveExisting removes an existing container if necessary. If this
// is option is not specified, and an existing container with a different
// configuration exists, it will not be removed and an error will be returned.
// If this option is specified, the existing container will be removed and a
// new container will be created.
func ContainerRemoveExisting() ContainerOption {
	return func(c *Container) error {
		c.removeExisting = true
		return nil
	}
}

// ContainerForceRemoveExisting forces removing the existing container if one
// is found. It does so even if the existing container matches the desired
// state.
func ContainerForceRemoveExisting() ContainerOption {
	return func(c *Container) error {
		c.forceRemoveExisting = true
		return nil
	}
}

// ContainerConfig specifies the container configuration.
func ContainerConfig(config *dockerclient.ContainerConfig) ContainerOption {
	return func(c *Container) error {
		c.containerConfig = config
		return nil
	}
}

// ContainerHostConfig specifies the container host configuration.
func ContainerHostConfig(config *dockerclient.HostConfig) ContainerOption {
	return func(c *Container) error {
		c.hostConfig = config
		return nil
	}
}

// ContainerAuthConfig specifies the auth credentials used when pulling an
// image.
func ContainerAuthConfig(ac *dockerclient.AuthConfig) ContainerOption {
	return func(c *Container) error {
		c.authConfig = ac
		return nil
	}
}

// Apply creates the container, possibly removing it as necessary based on the
// container options that were set.
func (c *Container) Apply(docker dockerclient.Client) error {
	ci, err := docker.InspectContainer(c.name)

	// force remove existing
	if ci != nil && c.forceRemoveExisting {
		if err := docker.RemoveContainer(ci.Id, true, false); err != nil {
			return stackerr.Wrap(err)
		}
		// we just removed the running container and want to start a new one
		err = dockerclient.ErrNotFound
	}

	// already exists
	if err == nil {
		ok, err := c.checkRunning(docker, ci)
		if err != nil {
			return err
		}
		if ok {
			// existing container is good
			if ci.State.Running {
				// and it's already running, so we do nothing
				return nil
			}
		} else {
			// we just removed the running container and want to start a new one
			err = dockerclient.ErrNotFound
		}
	}

	// unknown error, bail
	if err != nil && err != dockerclient.ErrNotFound {
		return stackerr.Wrap(err)
	}

	// container does not exist, create it
	if err == dockerclient.ErrNotFound {
		_, err := dockerutil.CreateWithPull(docker, c.containerConfig, c.name, c.authConfig)
		if err != nil {
			return err
		}

		ci, err = docker.InspectContainer(c.name)
		if err != nil {
			return stackerr.Wrap(err)
		}
	}

	// start the container
	err = docker.StartContainer(ci.Id, c.hostConfig)
	if err != nil {
		return stackerr.Wrap(err)
	}

	return nil
}

func (c *Container) checkRunning(docker dockerclient.Client, current *dockerclient.ContainerInfo) (bool, error) {
	// image comparison is by ID, so we need to find the ID of our desired image
	desiredImageID, err := dockerutil.ImageID(docker, c.containerConfig.Image, nil)
	if err != nil {
		return false, err
	}

	if current.Image != desiredImageID {
		// if we aren't allowed to remove the existing container, consider this a failure
		if !c.removeExisting {
			return false, stackerr.Newf(
				"container %q running with image %q but desired image is %q with id %q",
				c.name,
				current.Image,
				c.containerConfig.Image,
				desiredImageID,
			)
		}

		// otherwise remove it since it isn't want we want
		if err := docker.RemoveContainer(current.Id, true, false); err != nil {
			return false, stackerr.Wrap(err)
		}

		// trigger starting a new container
		return false, nil
	}

	// we're running the correct image
	return true, nil
}

// ApplyGraph creates all the specified containers. It handles links making
// sure the dependencies are created in the right order.
func ApplyGraph(docker dockerclient.Client, containers []*Container) error {
	known := map[string]struct{}{}
	started := map[string]bool{}

	// we want to know all known names so we can detect malformed links
	for _, c := range containers {
		known[c.name] = struct{}{}
	}

	// TODO: parallel pull pass?

	// keep doing rounds of parallel starts until we're all done or error out
	pending := containers
	for len(pending) > 0 {
		var eg errgroup.Group
		var starting []string
		var nextRound []*Container

	pendingLoop:
		for _, c := range pending {
			if c.hostConfig != nil {
				// TODO: also include c.hostConfig.VolumesFrom
				for _, link := range c.hostConfig.Links {
					// only care about the name, not the alias
					parts := strings.Split(link, ":")

					// make sure the link is known
					if _, ok := known[parts[0]]; !ok {
						return stackerr.Newf("%q expects unknown link %q", c.name, link)
					}

					// we need to wait for a dependency, schedule for the next round
					if !started[parts[0]] {
						nextRound = append(nextRound, c)
						continue pendingLoop
					}
				}
			}

			starting = append(starting, c.name)
			eg.Add(1)
			go func(c *Container) {
				defer eg.Done()
				if err := c.Apply(docker); err != nil {
					eg.Error(err)
				}
			}(c)
		}

		// now wait for all the containers we started in parallel
		if err := eg.Wait(); err != nil {
			return err
		}

		// we successfully started them
		for _, n := range starting {
			started[n] = true
		}

		// move on to the next round
		pending = nextRound
	}

	return nil
}
<|endoftext|>