<commit_before>// Copyright 2010 The Go 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 image

import (
	"strconv"
)

// A Point is an X, Y coordinate pair. The axes increase right and down.
type Point struct {
	X, Y int
}

// String returns a string representation of p like "(3,4)".
func (p Point) String() string {
	return "(" + strconv.Itoa(p.X) + "," + strconv.Itoa(p.Y) + ")"
}

// Add returns the vector p+q.
func (p Point) Add(q Point) Point {
	return Point{p.X + q.X, p.Y + q.Y}
}

// Sub returns the vector p-q.
func (p Point) Sub(q Point) Point {
	return Point{p.X - q.X, p.Y - q.Y}
}

// ZP is the zero Point.
var ZP Point

// Pt is shorthand for Point{X, Y}.
func Pt(X, Y int) Point {
	return Point{X, Y}
}

// A Rectangle contains the points with Min.X <= X < Max.X, Min.Y <= Y < Max.Y.
// It is well-formed if Min.X <= Max.X and likewise for Y. Points are always
// well-formed. A rectangle's methods always return well-formed outputs for
// well-formed inputs.
type Rectangle struct {
	Min, Max Point
}

// String returns a string representation of r like "(3,4)-(6,5)".
func (r Rectangle) String() string {
	return r.Min.String() + "-" + r.Max.String()
}

// Dx returns r's width.
func (r Rectangle) Dx() int {
	return r.Max.X - r.Min.X
}

// Dy returns r's height.
func (r Rectangle) Dy() int {
	return r.Max.Y - r.Min.Y
}

// Add returns the rectangle r translated by p.
func (r Rectangle) Add(p Point) Rectangle {
	return Rectangle{
		Point{r.Min.X + p.X, r.Min.Y + p.Y},
		Point{r.Max.X + p.X, r.Max.Y + p.Y},
	}
}

// Add returns the rectangle r translated by -p.
func (r Rectangle) Sub(p Point) Rectangle {
	return Rectangle{
		Point{r.Min.X - p.X, r.Min.Y - p.Y},
		Point{r.Max.X - p.X, r.Max.Y - p.Y},
	}
}

// Inset returns the rectangle r inset by n, which may be negative. If either
// of r's dimensions is less than 2*n then an empty rectangle near the center
// of r will be returned.
func (r Rectangle) Inset(n int) Rectangle {
	if r.Dx() < 2*n {
		r.Min.X = (r.Min.X + r.Max.X) / 2
		r.Max.X = r.Min.X
	} else {
		r.Min.X += n
		r.Max.X -= n
	}
	if r.Dy() < 2*n {
		r.Min.Y = (r.Min.Y + r.Max.Y) / 2
		r.Max.Y = r.Min.Y
	} else {
		r.Min.Y += n
		r.Max.Y -= n
	}
	return r
}

// Intersect returns the largest rectangle contained by both r and s. If the
// two rectangles do not overlap then the zero rectangle will be returned.
func (r Rectangle) Intersect(s Rectangle) Rectangle {
	if r.Min.X < s.Min.X {
		r.Min.X = s.Min.X
	}
	if r.Min.Y < s.Min.Y {
		r.Min.Y = s.Min.Y
	}
	if r.Max.X > s.Max.X {
		r.Max.X = s.Max.X
	}
	if r.Max.Y > s.Max.Y {
		r.Max.Y = s.Max.Y
	}
	if r.Min.X > r.Max.X || r.Min.Y > r.Max.Y {
		return ZR
	}
	return r
}

// Union returns the smallest rectangle that contains both r and s.
func (r Rectangle) Union(s Rectangle) Rectangle {
	if r.Min.X > s.Min.X {
		r.Min.X = s.Min.X
	}
	if r.Min.Y > s.Min.Y {
		r.Min.Y = s.Min.Y
	}
	if r.Max.X < s.Max.X {
		r.Max.X = s.Max.X
	}
	if r.Max.Y < s.Max.Y {
		r.Max.Y = s.Max.Y
	}
	return r
}

// Empty returns whether the rectangle contains no points.
func (r Rectangle) Empty() bool {
	return r.Min.X >= r.Max.X || r.Min.Y >= r.Max.Y
}

// Eq returns whether r and s are equal.
func (r Rectangle) Eq(s Rectangle) bool {
	return r.Min.X == s.Min.X && r.Min.Y == s.Min.Y &&
		r.Max.X == s.Max.X && r.Max.Y == s.Max.Y
}

// Overlaps returns whether r and s have a non-empty intersection.
func (r Rectangle) Overlaps(s Rectangle) bool {
	return r.Min.X < s.Max.X && s.Min.X < r.Max.X &&
		r.Min.Y < s.Max.Y && s.Min.Y < r.Max.Y
}

// Contains returns whether r contains p.
func (r Rectangle) Contains(p Point) bool {
	return p.X >= r.Min.X && p.X < r.Max.X &&
		p.Y >= r.Min.Y && p.Y < r.Max.Y
}

// Canon returns the canonical version of r. The returned rectangle has minimum
// and maximum coordinates swapped if necessary so that it is well-formed.
func (r Rectangle) Canon() Rectangle {
	if r.Max.X < r.Min.X {
		r.Min.X, r.Max.X = r.Max.X, r.Min.X
	}
	if r.Max.Y < r.Min.Y {
		r.Min.Y, r.Max.Y = r.Max.Y, r.Min.Y
	}
	return r
}

// ZR is the zero Rectangle.
var ZR Rectangle

// Rect is shorthand for Rectangle{Pt(x0, y0), Pt(x1, y1)}.
func Rect(x0, y0, x1, y1 int) Rectangle {
	if x0 > x1 {
		x0, x1 = x1, x0
	}
	if y0 > y1 {
		y0, y1 = y1, y0
	}
	return Rectangle{Point{x0, y0}, Point{x1, y1}}
}
<commit_msg>exp/draw: add Point.Eq, Point.Mul, Point.Div, Rectangle.Size methods<commit_after>// Copyright 2010 The Go 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 image

import (
	"strconv"
)

// A Point is an X, Y coordinate pair. The axes increase right and down.
type Point struct {
	X, Y int
}

// String returns a string representation of p like "(3,4)".
func (p Point) String() string {
	return "(" + strconv.Itoa(p.X) + "," + strconv.Itoa(p.Y) + ")"
}

// Add returns the vector p+q.
func (p Point) Add(q Point) Point {
	return Point{p.X + q.X, p.Y + q.Y}
}

// Sub returns the vector p-q.
func (p Point) Sub(q Point) Point {
	return Point{p.X - q.X, p.Y - q.Y}
}

// Mul returns the vector p*k.
func (p Point) Mul(k int) Point {
	return Point{p.X * k, p.Y * k}
}

// Div returns the vector p/k.
func (p Point) Div(k int) Point {
	return Point{p.X / k, p.Y / k}
}

// Eq returns whether p and q are equal.
func (p Point) Eq(q Point) bool {
	return p.X == q.X && p.Y == q.Y
}

// ZP is the zero Point.
var ZP Point

// Pt is shorthand for Point{X, Y}.
func Pt(X, Y int) Point {
	return Point{X, Y}
}

// A Rectangle contains the points with Min.X <= X < Max.X, Min.Y <= Y < Max.Y.
// It is well-formed if Min.X <= Max.X and likewise for Y. Points are always
// well-formed. A rectangle's methods always return well-formed outputs for
// well-formed inputs.
type Rectangle struct {
	Min, Max Point
}

// String returns a string representation of r like "(3,4)-(6,5)".
func (r Rectangle) String() string {
	return r.Min.String() + "-" + r.Max.String()
}

// Dx returns r's width.
func (r Rectangle) Dx() int {
	return r.Max.X - r.Min.X
}

// Dy returns r's height.
func (r Rectangle) Dy() int {
	return r.Max.Y - r.Min.Y
}

// Size returns r's width and height.
func (r Rectangle) Size() Point {
	return Point{
		r.Max.X - r.Min.X,
		r.Max.Y - r.Min.Y,
	}
}

// Add returns the rectangle r translated by p.
func (r Rectangle) Add(p Point) Rectangle {
	return Rectangle{
		Point{r.Min.X + p.X, r.Min.Y + p.Y},
		Point{r.Max.X + p.X, r.Max.Y + p.Y},
	}
}

// Add returns the rectangle r translated by -p.
func (r Rectangle) Sub(p Point) Rectangle {
	return Rectangle{
		Point{r.Min.X - p.X, r.Min.Y - p.Y},
		Point{r.Max.X - p.X, r.Max.Y - p.Y},
	}
}

// Inset returns the rectangle r inset by n, which may be negative. If either
// of r's dimensions is less than 2*n then an empty rectangle near the center
// of r will be returned.
func (r Rectangle) Inset(n int) Rectangle {
	if r.Dx() < 2*n {
		r.Min.X = (r.Min.X + r.Max.X) / 2
		r.Max.X = r.Min.X
	} else {
		r.Min.X += n
		r.Max.X -= n
	}
	if r.Dy() < 2*n {
		r.Min.Y = (r.Min.Y + r.Max.Y) / 2
		r.Max.Y = r.Min.Y
	} else {
		r.Min.Y += n
		r.Max.Y -= n
	}
	return r
}

// Intersect returns the largest rectangle contained by both r and s. If the
// two rectangles do not overlap then the zero rectangle will be returned.
func (r Rectangle) Intersect(s Rectangle) Rectangle {
	if r.Min.X < s.Min.X {
		r.Min.X = s.Min.X
	}
	if r.Min.Y < s.Min.Y {
		r.Min.Y = s.Min.Y
	}
	if r.Max.X > s.Max.X {
		r.Max.X = s.Max.X
	}
	if r.Max.Y > s.Max.Y {
		r.Max.Y = s.Max.Y
	}
	if r.Min.X > r.Max.X || r.Min.Y > r.Max.Y {
		return ZR
	}
	return r
}

// Union returns the smallest rectangle that contains both r and s.
func (r Rectangle) Union(s Rectangle) Rectangle {
	if r.Min.X > s.Min.X {
		r.Min.X = s.Min.X
	}
	if r.Min.Y > s.Min.Y {
		r.Min.Y = s.Min.Y
	}
	if r.Max.X < s.Max.X {
		r.Max.X = s.Max.X
	}
	if r.Max.Y < s.Max.Y {
		r.Max.Y = s.Max.Y
	}
	return r
}

// Empty returns whether the rectangle contains no points.
func (r Rectangle) Empty() bool {
	return r.Min.X >= r.Max.X || r.Min.Y >= r.Max.Y
}

// Eq returns whether r and s are equal.
func (r Rectangle) Eq(s Rectangle) bool {
	return r.Min.X == s.Min.X && r.Min.Y == s.Min.Y &&
		r.Max.X == s.Max.X && r.Max.Y == s.Max.Y
}

// Overlaps returns whether r and s have a non-empty intersection.
func (r Rectangle) Overlaps(s Rectangle) bool {
	return r.Min.X < s.Max.X && s.Min.X < r.Max.X &&
		r.Min.Y < s.Max.Y && s.Min.Y < r.Max.Y
}

// Contains returns whether r contains p.
func (r Rectangle) Contains(p Point) bool {
	return p.X >= r.Min.X && p.X < r.Max.X &&
		p.Y >= r.Min.Y && p.Y < r.Max.Y
}

// Canon returns the canonical version of r. The returned rectangle has minimum
// and maximum coordinates swapped if necessary so that it is well-formed.
func (r Rectangle) Canon() Rectangle {
	if r.Max.X < r.Min.X {
		r.Min.X, r.Max.X = r.Max.X, r.Min.X
	}
	if r.Max.Y < r.Min.Y {
		r.Min.Y, r.Max.Y = r.Max.Y, r.Min.Y
	}
	return r
}

// ZR is the zero Rectangle.
var ZR Rectangle

// Rect is shorthand for Rectangle{Pt(x0, y0), Pt(x1, y1)}.
func Rect(x0, y0, x1, y1 int) Rectangle {
	if x0 > x1 {
		x0, x1 = x1, x0
	}
	if y0 > y1 {
		y0, y1 = y1, y0
	}
	return Rectangle{Point{x0, y0}, Point{x1, y1}}
}
<|endoftext|>


<commit_before>package atomic_test

import (
	"errors"
	chaining "jecolasurdo/go-chaining"
	"jecolasurdo/go-chaining/internal/atomic"
	"testing"

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

func Test_ApplyUnaryIface_PreviousError_IgnoresAction(t *testing.T) {
	a := new(atomic.Context)
	c := new(chaining.Context)
	timesActionWasCalled := 0
	action := func(interface{}) (interface{}, error) {
		timesActionWasCalled++
		return nil, nil
	}

	c.LocalError = errors.New("test error")
	a.ApplyUnaryIface(action, &chaining.ActionArg{}, c)

	assert.Equal(t, 0, timesActionWasCalled)
}

func Test_ApplyUnaryIface_NoPreviousError_ExecutesAction(t *testing.T) {
	a := new(atomic.Context)
	c := new(chaining.Context)
	timesActionWasCalled := 0
	action := func(interface{}) (interface{}, error) {
		timesActionWasCalled++
		return nil, nil
	}

	a.ApplyUnaryIface(action, &chaining.ActionArg{}, c)

	assert.Equal(t, 1, timesActionWasCalled)
}

func Test_ApplyUnaryIface_NoPreviousError_BehaviorIsNotSpecified_InjectsPreviousValue(t *testing.T) {
	a := new(atomic.Context)
	c := new(chaining.Context)
	injectedValue := ""
	action := func(value interface{}) (interface{}, error) {
		injectedValue = value.(string)
		return nil, nil
	}
	argWithBehaviorNotSpecified := chaining.ActionArg{}
	simulatedValueOfPreviousActionInChain := "somevalue"
	c.PreviousActionResult = simulatedValueOfPreviousActionInChain

	a.ApplyUnaryIface(action, &argWithBehaviorNotSpecified, c)

	assert.Equal(t, simulatedValueOfPreviousActionInChain, injectedValue)
}

// func Test_ApplyUnaryIface_NoPreviousError_BehaviorIsUsePrevious_InjectsPreviousValue(t *testing.T) {
// 	d := new(chaining.Context)
// 	injectedValue := ""
// 	action := func(value interface{}) (interface{}, error) {
// 		injectedValue = value.(string)
// 		return nil, nil
// 	}
// 	argWithSpecifiedBehavior := chaining.ActionArg{
// 		Behavior: injectionbehavior.InjectPreviousResult,
// 	}
// 	simulatedValueOfPreviousActionInChain := "somevalue"
// 	d.PreviousActionResult = simulatedValueOfPreviousActionInChain

// 	d.ApplyUnaryIface(action, argWithSpecifiedBehavior)

// 	assert.Equal(t, simulatedValueOfPreviousActionInChain, injectedValue)
// }

// func Test_ApplyUnaryIface_NoPreviousError_BehaviorIsOverridePrevious_InjectsSuppliedValue(t *testing.T) {
// 	d := new(chaining.Context)
// 	injectedValue := ""
// 	action := func(value interface{}) (interface{}, error) {
// 		injectedValue = value.(string)
// 		return nil, nil
// 	}
// 	valueSubmittedThroughArg := "valueFromArg"
// 	argWithSpecifiedBehavior := chaining.ActionArg{
// 		Behavior: injectionbehavior.InjectSuppliedValue,
// 		Value:    valueSubmittedThroughArg,
// 	}
// 	simulatedValueOfPreviousActionInChain := "previousValue"
// 	d.PreviousActionResult = simulatedValueOfPreviousActionInChain

// 	d.ApplyUnaryIface(action, argWithSpecifiedBehavior)

// 	assert.Equal(t, valueSubmittedThroughArg, injectedValue)
// }

// func Test_ApplyUnaryIface_NoPreviousError_ForAnySpecifiedBehavior_SetsPreviousActionResult(t *testing.T) {
// 	d := new(chaining.Context)
// 	expectedReturnValue := "expectedReturnValue"
// 	action := func(value interface{}) (interface{}, error) { return expectedReturnValue, nil }
// 	arg := chaining.ActionArg{
// 		Value: "valueFromArg",
// 	}

// 	d.PreviousActionResult = nil
// 	arg.Behavior = injectionbehavior.InjectSuppliedValue
// 	d.ApplyUnaryIface(action, arg)
// 	assert.Equal(t, expectedReturnValue, d.PreviousActionResult)

// 	d.PreviousActionResult = nil
// 	arg.Behavior = injectionbehavior.InjectPreviousResult
// 	d.ApplyUnaryIface(action, arg)
// 	assert.Equal(t, expectedReturnValue, d.PreviousActionResult)

// 	d.PreviousActionResult = nil
// 	arg.Behavior = injectionbehavior.NotSpecified
// 	d.ApplyUnaryIface(action, arg)
// 	assert.Equal(t, expectedReturnValue, d.PreviousActionResult)
// }
<commit_msg>Passing fourth  AUI function test.<commit_after>package atomic_test

import (
	"errors"
	chaining "jecolasurdo/go-chaining"
	"jecolasurdo/go-chaining/injectionbehavior"
	"jecolasurdo/go-chaining/internal/atomic"
	"testing"

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

func Test_ApplyUnaryIface_PreviousError_IgnoresAction(t *testing.T) {
	a := new(atomic.Context)
	c := new(chaining.Context)
	timesActionWasCalled := 0
	action := func(interface{}) (interface{}, error) {
		timesActionWasCalled++
		return nil, nil
	}

	c.LocalError = errors.New("test error")
	a.ApplyUnaryIface(action, &chaining.ActionArg{}, c)

	assert.Equal(t, 0, timesActionWasCalled)
}

func Test_ApplyUnaryIface_NoPreviousError_ExecutesAction(t *testing.T) {
	a := new(atomic.Context)
	c := new(chaining.Context)
	timesActionWasCalled := 0
	action := func(interface{}) (interface{}, error) {
		timesActionWasCalled++
		return nil, nil
	}

	a.ApplyUnaryIface(action, &chaining.ActionArg{}, c)

	assert.Equal(t, 1, timesActionWasCalled)
}

func Test_ApplyUnaryIface_NoPreviousError_BehaviorIsNotSpecified_InjectsPreviousValue(t *testing.T) {
	a := new(atomic.Context)
	c := new(chaining.Context)
	injectedValue := ""
	action := func(value interface{}) (interface{}, error) {
		injectedValue = value.(string)
		return nil, nil
	}
	argWithBehaviorNotSpecified := chaining.ActionArg{}
	simulatedValueOfPreviousActionInChain := "somevalue"
	c.PreviousActionResult = simulatedValueOfPreviousActionInChain

	a.ApplyUnaryIface(action, &argWithBehaviorNotSpecified, c)

	assert.Equal(t, simulatedValueOfPreviousActionInChain, injectedValue)
}

func Test_ApplyUnaryIface_NoPreviousError_BehaviorIsUsePrevious_InjectsPreviousValue(t *testing.T) {
	a := new(atomic.Context)
	c := new(chaining.Context)
	injectedValue := ""
	action := func(value interface{}) (interface{}, error) {
		injectedValue = value.(string)
		return nil, nil
	}
	argWithSpecifiedBehavior := chaining.ActionArg{
		Behavior: injectionbehavior.InjectPreviousResult,
	}
	simulatedValueOfPreviousActionInChain := "somevalue"
	c.PreviousActionResult = simulatedValueOfPreviousActionInChain

	a.ApplyUnaryIface(action, &argWithSpecifiedBehavior, c)

	assert.Equal(t, simulatedValueOfPreviousActionInChain, injectedValue)
}

// func Test_ApplyUnaryIface_NoPreviousError_BehaviorIsOverridePrevious_InjectsSuppliedValue(t *testing.T) {
// 	d := new(chaining.Context)
// 	injectedValue := ""
// 	action := func(value interface{}) (interface{}, error) {
// 		injectedValue = value.(string)
// 		return nil, nil
// 	}
// 	valueSubmittedThroughArg := "valueFromArg"
// 	argWithSpecifiedBehavior := chaining.ActionArg{
// 		Behavior: injectionbehavior.InjectSuppliedValue,
// 		Value:    valueSubmittedThroughArg,
// 	}
// 	simulatedValueOfPreviousActionInChain := "previousValue"
// 	d.PreviousActionResult = simulatedValueOfPreviousActionInChain

// 	d.ApplyUnaryIface(action, argWithSpecifiedBehavior)

// 	assert.Equal(t, valueSubmittedThroughArg, injectedValue)
// }

// func Test_ApplyUnaryIface_NoPreviousError_ForAnySpecifiedBehavior_SetsPreviousActionResult(t *testing.T) {
// 	d := new(chaining.Context)
// 	expectedReturnValue := "expectedReturnValue"
// 	action := func(value interface{}) (interface{}, error) { return expectedReturnValue, nil }
// 	arg := chaining.ActionArg{
// 		Value: "valueFromArg",
// 	}

// 	d.PreviousActionResult = nil
// 	arg.Behavior = injectionbehavior.InjectSuppliedValue
// 	d.ApplyUnaryIface(action, arg)
// 	assert.Equal(t, expectedReturnValue, d.PreviousActionResult)

// 	d.PreviousActionResult = nil
// 	arg.Behavior = injectionbehavior.InjectPreviousResult
// 	d.ApplyUnaryIface(action, arg)
// 	assert.Equal(t, expectedReturnValue, d.PreviousActionResult)

// 	d.PreviousActionResult = nil
// 	arg.Behavior = injectionbehavior.NotSpecified
// 	d.ApplyUnaryIface(action, arg)
// 	assert.Equal(t, expectedReturnValue, d.PreviousActionResult)
// }
<|endoftext|>


<commit_before>// Copyright 2020 Buf Technologies Inc.
//
// 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 buf

import (
	"context"

	"github.com/bufbuild/buf/internal/pkg/app/appcmd"
	"github.com/bufbuild/buf/internal/pkg/app/appflag"
)

const version = "0.13.0-dev"

// Main is the main.
func Main(use string, options ...RootCommandOption) {
	appcmd.Main(context.Background(), newRootCommand(use, options...), version)
}

// RootCommandOption is an option for a root Command.
type RootCommandOption func(*appcmd.Command, appflag.Builder)
<commit_msg>Update to v0.12.1<commit_after>// Copyright 2020 Buf Technologies Inc.
//
// 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 buf

import (
	"context"

	"github.com/bufbuild/buf/internal/pkg/app/appcmd"
	"github.com/bufbuild/buf/internal/pkg/app/appflag"
)

const version = "0.12.1"

// Main is the main.
func Main(use string, options ...RootCommandOption) {
	appcmd.Main(context.Background(), newRootCommand(use, options...), version)
}

// RootCommandOption is an option for a root Command.
type RootCommandOption func(*appcmd.Command, appflag.Builder)
<|endoftext|>


<commit_before>package run

import (
	"fmt"
	"strings"

	"go.pedge.io/protolog"

	"github.com/pachyderm/pachyderm/src/pkg/graph"
	"github.com/pachyderm/pachyderm/src/pkg/timing"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/container"
	"github.com/pachyderm/pachyderm/src/pps/source"
	"github.com/pachyderm/pachyderm/src/pps/store"
	"github.com/satori/go.uuid"
)

type runner struct {
	sourcer         source.Sourcer
	grapher         graph.Grapher
	containerClient container.Client
	storeClient     store.Client
}

func newRunner(
	sourcer source.Sourcer,
	grapher graph.Grapher,
	containerClient container.Client,
	storeClient store.Client,
) *runner {
	return &runner{
		sourcer,
		grapher,
		containerClient,
		storeClient,
	}
}

func (r *runner) Start(pipelineSource *pps.PipelineSource) (string, error) {
	dirPath, pipeline, err := r.sourcer.GetDirPathAndPipeline(pipelineSource)
	if err != nil {
		return "", err
	}
	pipelineRunID := strings.Replace(uuid.NewV4().String(), "-", "", -1)
	pipelineRun := &pps.PipelineRun{
		Id:             pipelineRunID,
		Pipeline:       pipeline,
		PipelineSource: pipelineSource,
	}
	if err := r.storeClient.AddPipelineRun(pipelineRun); err != nil {
		return "", err
	}
	protolog.Info(
		&AddedPipelineRun{
			PipelineRun: pipelineRun,
		},
	)
	nameToNode := pps.GetNameToNode(pipeline)
	nameToDockerService := pps.GetNameToDockerService(pipeline)
	nameToNodeInfo, err := getNameToNodeInfo(nameToNode)
	if err != nil {
		return "", err
	}
	nameToNodeFunc := make(map[string]func() error)
	for name, node := range nameToNode {
		nodeFunc, err := r.getNodeFunc(
			pipelineRunID,
			name,
			node,
			nameToDockerService,
			dirPath,
			1,
		)
		if err != nil {
			return "", err
		}
		nameToNodeFunc[name] = nodeFunc
	}
	run, err := r.grapher.Build(
		nameToNodeInfo,
		nameToNodeFunc,
	)
	if err != nil {
		return "", err
	}
	if err := r.storeClient.AddPipelineRunStatus(pipelineRunID, pps.PipelineRunStatusType_PIPELINE_RUN_STATUS_TYPE_STARTED); err != nil {
		return "", err
	}
	go func() {
		if err := run.Do(); err != nil {
			if storeErr := r.storeClient.AddPipelineRunStatus(pipelineRunID, pps.PipelineRunStatusType_PIPELINE_RUN_STATUS_TYPE_ERROR); storeErr != nil {
				protolog.Errorln(storeErr.Error())
			}
		} else {
			if storeErr := r.storeClient.AddPipelineRunStatus(pipelineRunID, pps.PipelineRunStatusType_PIPELINE_RUN_STATUS_TYPE_SUCCESS); storeErr != nil {
				protolog.Errorln(storeErr.Error())
			}
		}
	}()
	return pipelineRunID, nil
}

func (r *runner) getNodeFunc(
	pipelineRunID string,
	name string,
	node *pps.Node,
	nameToDockerService map[string]*pps.DockerService,
	dirPath string,
	numContainers int,
) (func() error, error) {
	dockerService, ok := nameToDockerService[node.Service]
	if !ok {
		return nil, fmt.Errorf("no service for name %s", node.Service)
	}
	if dockerService.Build != "" || dockerService.Dockerfile != "" {
		return nil, fmt.Errorf("build/dockerfile not supported yet")
	}
	return func() (retErr error) {
		if err := r.containerClient.Pull(
			dockerService.Image,
			container.PullOptions{
			//OutputStream: newPipelineRunLogWriter(
			//pipelineRunID,
			//"",
			//name,
			//pps.OutputStream_OUTPUT_STREAM_NONE,
			//timing.NewSystemTimer(),
			//r.storeClient,
			//),
			},
		); err != nil {
			return err
		}
		containers, err := r.containerClient.Create(
			dockerService.Image,
			container.CreateOptions{
				Input:         node.Input,
				Output:        node.Output,
				HasCommand:    len(node.Run) > 0,
				NumContainers: numContainers,
			},
		)
		if err != nil {
			return err
		}
		defer func() {
			for _, containerID := range containers {
				_ = r.containerClient.Kill(containerID, container.KillOptions{})
				_ = r.containerClient.Remove(containerID, container.RemoveOptions{})
			}
		}()
		for _, containerID := range containers {
			if err := r.containerClient.Start(
				containerID,
				container.StartOptions{
					Commands: node.Run,
				},
			); err != nil {
				return err
			}
		}
		errC := make(chan error, len(containers))
		for _, containerID := range containers {
			containerID := containerID
			go func() {
				errC <- r.containerClient.Logs(
					containerID,
					container.LogsOptions{
						Stdout: newPipelineRunLogWriter(
							pipelineRunID,
							containerID,
							name,
							pps.OutputStream_OUTPUT_STREAM_STDOUT,
							timing.NewSystemTimer(),
							r.storeClient,
						),
						Stderr: newPipelineRunLogWriter(
							pipelineRunID,
							containerID,
							name,
							pps.OutputStream_OUTPUT_STREAM_STDOUT,
							timing.NewSystemTimer(),
							r.storeClient,
						),
					},
				)
			}()
		}
		for _, containerID := range containers {
			if err := r.containerClient.Wait(containerID, container.WaitOptions{}); err != nil {
				return err
			}
		}
		err = nil
		for _ = range containers {
			if logsErr := <-errC; logsErr != nil && err == nil {
				err = logsErr
			}
		}
		return err
	}, nil
}
<commit_msg>fix setting both outputstream values to STDOUT, now correctly setting STDERR<commit_after>package run

import (
	"fmt"
	"strings"

	"go.pedge.io/protolog"

	"github.com/pachyderm/pachyderm/src/pkg/graph"
	"github.com/pachyderm/pachyderm/src/pkg/timing"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/container"
	"github.com/pachyderm/pachyderm/src/pps/source"
	"github.com/pachyderm/pachyderm/src/pps/store"
	"github.com/satori/go.uuid"
)

type runner struct {
	sourcer         source.Sourcer
	grapher         graph.Grapher
	containerClient container.Client
	storeClient     store.Client
}

func newRunner(
	sourcer source.Sourcer,
	grapher graph.Grapher,
	containerClient container.Client,
	storeClient store.Client,
) *runner {
	return &runner{
		sourcer,
		grapher,
		containerClient,
		storeClient,
	}
}

func (r *runner) Start(pipelineSource *pps.PipelineSource) (string, error) {
	dirPath, pipeline, err := r.sourcer.GetDirPathAndPipeline(pipelineSource)
	if err != nil {
		return "", err
	}
	pipelineRunID := strings.Replace(uuid.NewV4().String(), "-", "", -1)
	pipelineRun := &pps.PipelineRun{
		Id:             pipelineRunID,
		Pipeline:       pipeline,
		PipelineSource: pipelineSource,
	}
	if err := r.storeClient.AddPipelineRun(pipelineRun); err != nil {
		return "", err
	}
	protolog.Info(
		&AddedPipelineRun{
			PipelineRun: pipelineRun,
		},
	)
	nameToNode := pps.GetNameToNode(pipeline)
	nameToDockerService := pps.GetNameToDockerService(pipeline)
	nameToNodeInfo, err := getNameToNodeInfo(nameToNode)
	if err != nil {
		return "", err
	}
	nameToNodeFunc := make(map[string]func() error)
	for name, node := range nameToNode {
		nodeFunc, err := r.getNodeFunc(
			pipelineRunID,
			name,
			node,
			nameToDockerService,
			dirPath,
			1,
		)
		if err != nil {
			return "", err
		}
		nameToNodeFunc[name] = nodeFunc
	}
	run, err := r.grapher.Build(
		nameToNodeInfo,
		nameToNodeFunc,
	)
	if err != nil {
		return "", err
	}
	if err := r.storeClient.AddPipelineRunStatus(pipelineRunID, pps.PipelineRunStatusType_PIPELINE_RUN_STATUS_TYPE_STARTED); err != nil {
		return "", err
	}
	go func() {
		if err := run.Do(); err != nil {
			if storeErr := r.storeClient.AddPipelineRunStatus(pipelineRunID, pps.PipelineRunStatusType_PIPELINE_RUN_STATUS_TYPE_ERROR); storeErr != nil {
				protolog.Errorln(storeErr.Error())
			}
		} else {
			if storeErr := r.storeClient.AddPipelineRunStatus(pipelineRunID, pps.PipelineRunStatusType_PIPELINE_RUN_STATUS_TYPE_SUCCESS); storeErr != nil {
				protolog.Errorln(storeErr.Error())
			}
		}
	}()
	return pipelineRunID, nil
}

func (r *runner) getNodeFunc(
	pipelineRunID string,
	name string,
	node *pps.Node,
	nameToDockerService map[string]*pps.DockerService,
	dirPath string,
	numContainers int,
) (func() error, error) {
	dockerService, ok := nameToDockerService[node.Service]
	if !ok {
		return nil, fmt.Errorf("no service for name %s", node.Service)
	}
	if dockerService.Build != "" || dockerService.Dockerfile != "" {
		return nil, fmt.Errorf("build/dockerfile not supported yet")
	}
	return func() (retErr error) {
		if err := r.containerClient.Pull(
			dockerService.Image,
			container.PullOptions{
			//OutputStream: newPipelineRunLogWriter(
			//pipelineRunID,
			//"",
			//name,
			//pps.OutputStream_OUTPUT_STREAM_NONE,
			//timing.NewSystemTimer(),
			//r.storeClient,
			//),
			},
		); err != nil {
			return err
		}
		containers, err := r.containerClient.Create(
			dockerService.Image,
			container.CreateOptions{
				Input:         node.Input,
				Output:        node.Output,
				HasCommand:    len(node.Run) > 0,
				NumContainers: numContainers,
			},
		)
		if err != nil {
			return err
		}
		defer func() {
			for _, containerID := range containers {
				_ = r.containerClient.Kill(containerID, container.KillOptions{})
				_ = r.containerClient.Remove(containerID, container.RemoveOptions{})
			}
		}()
		for _, containerID := range containers {
			if err := r.containerClient.Start(
				containerID,
				container.StartOptions{
					Commands: node.Run,
				},
			); err != nil {
				return err
			}
		}
		errC := make(chan error, len(containers))
		for _, containerID := range containers {
			containerID := containerID
			go func() {
				errC <- r.containerClient.Logs(
					containerID,
					container.LogsOptions{
						Stdout: newPipelineRunLogWriter(
							pipelineRunID,
							containerID,
							name,
							pps.OutputStream_OUTPUT_STREAM_STDOUT,
							timing.NewSystemTimer(),
							r.storeClient,
						),
						Stderr: newPipelineRunLogWriter(
							pipelineRunID,
							containerID,
							name,
							pps.OutputStream_OUTPUT_STREAM_STDERR,
							timing.NewSystemTimer(),
							r.storeClient,
						),
					},
				)
			}()
		}
		for _, containerID := range containers {
			if err := r.containerClient.Wait(containerID, container.WaitOptions{}); err != nil {
				return err
			}
		}
		err = nil
		for _ = range containers {
			if logsErr := <-errC; logsErr != nil && err == nil {
				err = logsErr
			}
		}
		return err
	}, nil
}
<|endoftext|>


<commit_before>package console

import (
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/cenkalti/rain/internal/jsonutil"
	"github.com/cenkalti/rain/rainrpc"
	"github.com/cenkalti/rain/torrent"
	"github.com/jroimartin/gocui"
)

type Console struct {
	client          *rainrpc.Client
	torrents        []rainrpc.Torrent
	errTorrents     error
	selectedID      uint64
	stats           torrent.Stats
	errStats        error
	m               sync.Mutex
	updateTorrentsC chan struct{}
	updateStatsC    chan struct{}
}

func New(clt *rainrpc.Client) *Console {
	return &Console{
		client:          clt,
		updateTorrentsC: make(chan struct{}),
		updateStatsC:    make(chan struct{}),
	}
}

func (c *Console) Run() error {
	g, err := gocui.NewGui(gocui.OutputNormal)
	if err != nil {
		return err
	}
	defer g.Close()

	g.SetManagerFunc(c.layout)

	g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, quit)
	g.SetKeybinding("", 'q', gocui.ModNone, quit)
	g.SetKeybinding("torrents", 'j', gocui.ModNone, c.cursorDown)
	g.SetKeybinding("torrents", 'k', gocui.ModNone, c.cursorUp)
	g.SetKeybinding("torrents", 'R', gocui.ModNone, c.removeTorrent)
	g.SetKeybinding("torrents", 's', gocui.ModNone, c.startTorrent)
	g.SetKeybinding("torrents", 'S', gocui.ModNone, c.stopTorrent)

	go c.updateLoop(g)

	err = g.MainLoop()
	if err == gocui.ErrQuit {
		err = nil
	}
	return err
}

func (c *Console) layout(g *gocui.Gui) error {
	err := c.drawTorrents(g)
	if err != nil {
		return err
	}
	err = c.drawStats(g)
	if err != nil {
		return err
	}
	_, err = g.SetCurrentView("torrents")
	return err
}

func (c *Console) drawTorrents(g *gocui.Gui) error {
	c.m.Lock()
	defer c.m.Unlock()

	maxX, maxY := g.Size()
	halfY := maxY / 2
	if v, err := g.SetView("torrents", -1, -1, maxX, halfY); err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		v.Highlight = true
		v.SelBgColor = gocui.ColorGreen
		v.SelFgColor = gocui.ColorBlack
		fmt.Fprintln(v, "loading torrents...")
	} else {
		v.Clear()
		if c.errTorrents != nil {
			fmt.Fprintln(v, "error:", c.errTorrents)
			c.selectedID = 0
		} else {
			for _, t := range c.torrents {
				fmt.Fprintf(v, "%5d %s %5d %s\n", t.ID, t.InfoHash, t.Port, t.Name)
			}
			_, cy := v.Cursor()
			_, oy := v.Origin()
			selectedRow := cy + oy
			if selectedRow < len(c.torrents) {
				c.setSelectedID(c.torrents[selectedRow].ID)
			}
		}
	}
	return nil
}

func (c *Console) drawStats(g *gocui.Gui) error {
	c.m.Lock()
	defer c.m.Unlock()

	maxX, maxY := g.Size()
	halfY := maxY / 2
	if v, err := g.SetView("stats", -1, halfY, maxX, maxY); err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		v.Wrap = true
		fmt.Fprintln(v, "loading stats...")
	} else {
		v.Clear()
		if c.errStats != nil {
			fmt.Fprintln(v, "error:", c.errStats)
		} else {
			b, err := jsonutil.MarshalCompactPretty(c.stats)
			if err != nil {
				fmt.Fprintln(v, "error:", c.errStats)
			} else {
				fmt.Fprintln(v, string(b))
			}
		}
	}
	return nil
}

func (c *Console) updateLoop(g *gocui.Gui) {
	c.updateTorrents(g)
	c.updateStats(g)

	ticker := time.NewTicker(time.Second)
	for {
		select {
		case <-ticker.C:
			c.updateTorrents(g)
			c.updateStats(g)
		case <-c.updateTorrentsC:
			c.updateTorrents(g)
		case <-c.updateStatsC:
			c.updateStats(g)
		}
	}
}

func (c *Console) updateTorrents(g *gocui.Gui) {
	resp, err := c.client.ListTorrents()

	sort.Slice(resp.Torrents, func(i, j int) bool { return resp.Torrents[i].ID < resp.Torrents[j].ID })

	c.m.Lock()
	c.torrents = resp.Torrents
	c.errTorrents = err
	if len(c.torrents) == 0 {
		c.setSelectedID(0)
	} else if c.selectedID == 0 {
		c.setSelectedID(c.torrents[0].ID)
	}
	c.m.Unlock()

	g.Update(c.drawTorrents)
}

func (c *Console) updateStats(g *gocui.Gui) {
	c.m.Lock()
	selectedID := c.selectedID
	c.m.Unlock()

	if selectedID != 0 {
		resp, err := c.client.GetTorrentStats(selectedID)
		c.m.Lock()
		c.stats = resp.Stats
		c.errStats = err
		c.m.Unlock()
	} else {
		c.m.Lock()
		c.errStats = errors.New("no torrent selected")
		c.m.Unlock()
	}

	g.Update(c.drawStats)
}

func quit(g *gocui.Gui, v *gocui.View) error {
	return gocui.ErrQuit
}

func (c *Console) cursorDown(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	defer c.m.Unlock()

	cx, cy := v.Cursor()
	ox, oy := v.Origin()
	if cy+oy >= len(c.torrents)-1 {
		return nil
	}
	if err := v.SetCursor(cx, cy+1); err != nil {
		if err := v.SetOrigin(ox, oy+1); err != nil {
			return err
		}
	}
	c.setSelectedID(c.torrents[cy+oy+1].ID)
	return nil
}

func (c *Console) cursorUp(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	defer c.m.Unlock()

	cx, cy := v.Cursor()
	ox, oy := v.Origin()
	if cy+oy <= 0 {
		return nil
	}
	if err := v.SetCursor(cx, cy-1); err != nil && oy > 0 {
		if err := v.SetOrigin(ox, oy-1); err != nil {
			return err
		}
	}
	c.setSelectedID(c.torrents[cy+oy-1].ID)
	return nil
}

func (c *Console) removeTorrent(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	id := c.selectedID
	c.m.Unlock()

	_, err := c.client.RemoveTorrent(id)
	if err != nil {
		return err
	}
	select {
	case c.updateTorrentsC <- struct{}{}:
	default:
	}
	return nil
}

func (c *Console) setSelectedID(id uint64) {
	changed := id != c.selectedID
	c.selectedID = id
	if changed {
		select {
		case c.updateStatsC <- struct{}{}:
		default:
		}
	}
}

func (c *Console) startTorrent(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	id := c.selectedID
	c.m.Unlock()

	_, err := c.client.StartTorrent(id)
	return err
}

func (c *Console) stopTorrent(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	id := c.selectedID
	c.m.Unlock()

	_, err := c.client.StopTorrent(id)
	return err
}
<commit_msg>extra guard for console<commit_after>package console

import (
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/cenkalti/rain/internal/jsonutil"
	"github.com/cenkalti/rain/rainrpc"
	"github.com/cenkalti/rain/torrent"
	"github.com/jroimartin/gocui"
)

type Console struct {
	client          *rainrpc.Client
	torrents        []rainrpc.Torrent
	errTorrents     error
	selectedID      uint64
	stats           torrent.Stats
	errStats        error
	m               sync.Mutex
	updateTorrentsC chan struct{}
	updateStatsC    chan struct{}
}

func New(clt *rainrpc.Client) *Console {
	return &Console{
		client:          clt,
		updateTorrentsC: make(chan struct{}),
		updateStatsC:    make(chan struct{}),
	}
}

func (c *Console) Run() error {
	g, err := gocui.NewGui(gocui.OutputNormal)
	if err != nil {
		return err
	}
	defer g.Close()

	g.SetManagerFunc(c.layout)

	g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, quit)
	g.SetKeybinding("", 'q', gocui.ModNone, quit)
	g.SetKeybinding("torrents", 'j', gocui.ModNone, c.cursorDown)
	g.SetKeybinding("torrents", 'k', gocui.ModNone, c.cursorUp)
	g.SetKeybinding("torrents", 'R', gocui.ModNone, c.removeTorrent)
	g.SetKeybinding("torrents", 's', gocui.ModNone, c.startTorrent)
	g.SetKeybinding("torrents", 'S', gocui.ModNone, c.stopTorrent)

	go c.updateLoop(g)

	err = g.MainLoop()
	if err == gocui.ErrQuit {
		err = nil
	}
	return err
}

func (c *Console) layout(g *gocui.Gui) error {
	err := c.drawTorrents(g)
	if err != nil {
		return err
	}
	err = c.drawStats(g)
	if err != nil {
		return err
	}
	_, err = g.SetCurrentView("torrents")
	return err
}

func (c *Console) drawTorrents(g *gocui.Gui) error {
	c.m.Lock()
	defer c.m.Unlock()

	maxX, maxY := g.Size()
	halfY := maxY / 2
	if v, err := g.SetView("torrents", -1, -1, maxX, halfY); err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		v.Highlight = true
		v.SelBgColor = gocui.ColorGreen
		v.SelFgColor = gocui.ColorBlack
		fmt.Fprintln(v, "loading torrents...")
	} else {
		v.Clear()
		if c.errTorrents != nil {
			fmt.Fprintln(v, "error:", c.errTorrents)
			c.selectedID = 0
		} else {
			for _, t := range c.torrents {
				fmt.Fprintf(v, "%5d %s %5d %s\n", t.ID, t.InfoHash, t.Port, t.Name)
			}
			_, cy := v.Cursor()
			_, oy := v.Origin()
			selectedRow := cy + oy
			if selectedRow < len(c.torrents) {
				c.setSelectedID(c.torrents[selectedRow].ID)
			}
		}
	}
	return nil
}

func (c *Console) drawStats(g *gocui.Gui) error {
	c.m.Lock()
	defer c.m.Unlock()

	maxX, maxY := g.Size()
	halfY := maxY / 2
	if v, err := g.SetView("stats", -1, halfY, maxX, maxY); err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		v.Wrap = true
		fmt.Fprintln(v, "loading stats...")
	} else {
		v.Clear()
		if c.errStats != nil {
			fmt.Fprintln(v, "error:", c.errStats)
		} else {
			b, err := jsonutil.MarshalCompactPretty(c.stats)
			if err != nil {
				fmt.Fprintln(v, "error:", c.errStats)
			} else {
				fmt.Fprintln(v, string(b))
			}
		}
	}
	return nil
}

func (c *Console) updateLoop(g *gocui.Gui) {
	c.updateTorrents(g)
	c.updateStats(g)

	ticker := time.NewTicker(time.Second)
	for {
		select {
		case <-ticker.C:
			c.updateTorrents(g)
			c.updateStats(g)
		case <-c.updateTorrentsC:
			c.updateTorrents(g)
		case <-c.updateStatsC:
			c.updateStats(g)
		}
	}
}

func (c *Console) updateTorrents(g *gocui.Gui) {
	resp, err := c.client.ListTorrents()

	sort.Slice(resp.Torrents, func(i, j int) bool { return resp.Torrents[i].ID < resp.Torrents[j].ID })

	c.m.Lock()
	c.torrents = resp.Torrents
	c.errTorrents = err
	if len(c.torrents) == 0 {
		c.setSelectedID(0)
	} else if c.selectedID == 0 {
		c.setSelectedID(c.torrents[0].ID)
	}
	c.m.Unlock()

	g.Update(c.drawTorrents)
}

func (c *Console) updateStats(g *gocui.Gui) {
	c.m.Lock()
	selectedID := c.selectedID
	c.m.Unlock()

	if selectedID != 0 {
		resp, err := c.client.GetTorrentStats(selectedID)
		c.m.Lock()
		c.stats = resp.Stats
		c.errStats = err
		c.m.Unlock()
	} else {
		c.m.Lock()
		c.errStats = errors.New("no torrent selected")
		c.m.Unlock()
	}

	g.Update(c.drawStats)
}

func quit(g *gocui.Gui, v *gocui.View) error {
	return gocui.ErrQuit
}

func (c *Console) cursorDown(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	defer c.m.Unlock()

	cx, cy := v.Cursor()
	ox, oy := v.Origin()
	if cy+oy >= len(c.torrents)-1 {
		return nil
	}
	if err := v.SetCursor(cx, cy+1); err != nil {
		if err := v.SetOrigin(ox, oy+1); err != nil {
			return err
		}
	}
	row := cy + oy + 1
	if row >= 0 || row < len(c.torrents) {
		c.setSelectedID(c.torrents[row].ID)
	}
	return nil
}

func (c *Console) cursorUp(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	defer c.m.Unlock()

	cx, cy := v.Cursor()
	ox, oy := v.Origin()
	if cy+oy <= 0 {
		return nil
	}
	if err := v.SetCursor(cx, cy-1); err != nil && oy > 0 {
		if err := v.SetOrigin(ox, oy-1); err != nil {
			return err
		}
	}
	row := cy + oy - 1
	if row >= 0 || row < len(c.torrents) {
		c.setSelectedID(c.torrents[row].ID)
	}
	return nil
}

func (c *Console) removeTorrent(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	id := c.selectedID
	c.m.Unlock()

	_, err := c.client.RemoveTorrent(id)
	if err != nil {
		return err
	}
	select {
	case c.updateTorrentsC <- struct{}{}:
	default:
	}
	return nil
}

func (c *Console) setSelectedID(id uint64) {
	changed := id != c.selectedID
	c.selectedID = id
	if changed {
		select {
		case c.updateStatsC <- struct{}{}:
		default:
		}
	}
}

func (c *Console) startTorrent(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	id := c.selectedID
	c.m.Unlock()

	_, err := c.client.StartTorrent(id)
	return err
}

func (c *Console) stopTorrent(g *gocui.Gui, v *gocui.View) error {
	c.m.Lock()
	id := c.selectedID
	c.m.Unlock()

	_, err := c.client.StopTorrent(id)
	return err
}
<|endoftext|>


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

//go:generate go run mkstdlib.go

// Package imports implements a Go pretty-printer (like package "go/format")
// that also adds or removes import statements as necessary.
package imports

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
)

// Options is golang.org/x/tools/imports.Options with extra internal-only options.
type Options struct {
	Env *ProcessEnv // The environment to use. Note: this contains the cached module and filesystem state.

	Fragment  bool // Accept fragment of a source file (no package statement)
	AllErrors bool // Report all errors (not just the first 10 on different lines)

	Comments  bool // Print comments (true if nil *Options provided)
	TabIndent bool // Use tabs for indent (true if nil *Options provided)
	TabWidth  int  // Tab width (8 if nil *Options provided)

	FormatOnly bool // Disable the insertion and deletion of imports
}

// Process implements golang.org/x/tools/imports.Process with explicit context in env.
func Process(filename string, src []byte, opt *Options) (formatted []byte, err error) {
	src, opt, err = initialize(filename, src, opt)
	if err != nil {
		return nil, err
	}

	fileSet := token.NewFileSet()
	file, adjust, err := parse(fileSet, filename, src, opt)
	if err != nil {
		return nil, err
	}

	if !opt.FormatOnly {
		if err := fixImports(fileSet, file, filename, opt.Env); err != nil {
			return nil, err
		}
	}
	return formatFile(fileSet, file, src, adjust, opt)
}

// FixImports returns a list of fixes to the imports that, when applied,
// will leave the imports in the same state as Process.
//
// Note that filename's directory influences which imports can be chosen,
// so it is important that filename be accurate.
func FixImports(filename string, src []byte, opt *Options) (fixes []*ImportFix, err error) {
	src, opt, err = initialize(filename, src, opt)
	if err != nil {
		return nil, err
	}

	fileSet := token.NewFileSet()
	file, _, err := parse(fileSet, filename, src, opt)
	if err != nil {
		return nil, err
	}

	return getFixes(fileSet, file, filename, opt.Env)
}

// ApplyFixes applies all of the fixes to the file and formats it. extraMode
// is added in when parsing the file.
func ApplyFixes(fixes []*ImportFix, filename string, src []byte, opt *Options, extraMode parser.Mode) (formatted []byte, err error) {
	src, opt, err = initialize(filename, src, opt)
	if err != nil {
		return nil, err
	}

	// Don't use parse() -- we don't care about fragments or statement lists
	// here, and we need to work with unparseable files.
	fileSet := token.NewFileSet()
	parserMode := parser.Mode(0)
	if opt.Comments {
		parserMode |= parser.ParseComments
	}
	if opt.AllErrors {
		parserMode |= parser.AllErrors
	}
	parserMode |= extraMode

	file, err := parser.ParseFile(fileSet, filename, src, parserMode)
	if file == nil {
		return nil, err
	}

	// Apply the fixes to the file.
	apply(fileSet, file, fixes)

	return formatFile(fileSet, file, src, nil, opt)
}

// GetAllCandidates gets all of the packages starting with prefix that can be
// imported by filename, sorted by import path.
func GetAllCandidates(ctx context.Context, callback func(ImportFix), searchPrefix, filename, filePkg string, opt *Options) error {
	_, opt, err := initialize(filename, []byte{}, opt)
	if err != nil {
		return err
	}
	return getAllCandidates(ctx, callback, searchPrefix, filename, filePkg, opt.Env)
}

// GetPackageExports returns all known packages with name pkg and their exports.
func GetPackageExports(ctx context.Context, callback func(PackageExport), searchPkg, filename, filePkg string, opt *Options) error {
	_, opt, err := initialize(filename, []byte{}, opt)
	if err != nil {
		return err
	}
	return getPackageExports(ctx, callback, searchPkg, filename, filePkg, opt.Env)
}

// initialize sets the values for opt and src.
// If they are provided, they are not changed. Otherwise opt is set to the
// default values and src is read from the file system.
func initialize(filename string, src []byte, opt *Options) ([]byte, *Options, error) {
	// Use defaults if opt is nil.
	if opt == nil {
		opt = &Options{Comments: true, TabIndent: true, TabWidth: 8}
	}

	// Set the env if the user has not provided it.
	if opt.Env == nil {
		opt.Env = &ProcessEnv{
			GOPATH:      build.Default.GOPATH,
			GOROOT:      build.Default.GOROOT,
			GOFLAGS:     os.Getenv("GOFLAGS"),
			GO111MODULE: os.Getenv("GO111MODULE"),
			GOPROXY:     os.Getenv("GOPROXY"),
			GOSUMDB:     os.Getenv("GOSUMDB"),
		}
	}

	// Set the logger if the user has not provided it.
	if opt.Env.Logf == nil {
		opt.Env.Logf = log.Printf
	}

	if src == nil {
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			return nil, nil, err
		}
		src = b
	}

	return src, opt, nil
}

func formatFile(fileSet *token.FileSet, file *ast.File, src []byte, adjust func(orig []byte, src []byte) []byte, opt *Options) ([]byte, error) {
	mergeImports(opt.Env, fileSet, file)
	sortImports(opt.Env, fileSet, file)
	imps := astutil.Imports(fileSet, file)
	var spacesBefore []string // import paths we need spaces before
	for _, impSection := range imps {
		// Within each block of contiguous imports, see if any
		// import lines are in different group numbers. If so,
		// we'll need to put a space between them so it's
		// compatible with gofmt.
		lastGroup := -1
		for _, importSpec := range impSection {
			importPath, _ := strconv.Unquote(importSpec.Path.Value)
			groupNum := importGroup(opt.Env, importPath)
			if groupNum != lastGroup && lastGroup != -1 {
				spacesBefore = append(spacesBefore, importPath)
			}
			lastGroup = groupNum
		}

	}

	printerMode := printer.UseSpaces
	if opt.TabIndent {
		printerMode |= printer.TabIndent
	}
	printConfig := &printer.Config{Mode: printerMode, Tabwidth: opt.TabWidth}

	var buf bytes.Buffer
	err := printConfig.Fprint(&buf, fileSet, file)
	if err != nil {
		return nil, err
	}
	out := buf.Bytes()
	if adjust != nil {
		out = adjust(src, out)
	}
	if len(spacesBefore) > 0 {
		out, err = addImportSpaces(bytes.NewReader(out), spacesBefore)
		if err != nil {
			return nil, err
		}
	}

	out, err = format.Source(out)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// parse parses src, which was read from filename,
// as a Go source file or statement list.
func parse(fset *token.FileSet, filename string, src []byte, opt *Options) (*ast.File, func(orig, src []byte) []byte, error) {
	parserMode := parser.Mode(0)
	if opt.Comments {
		parserMode |= parser.ParseComments
	}
	if opt.AllErrors {
		parserMode |= parser.AllErrors
	}

	// Try as whole source file.
	file, err := parser.ParseFile(fset, filename, src, parserMode)
	if err == nil {
		return file, nil, nil
	}
	// If the error is that the source file didn't begin with a
	// package line and we accept fragmented input, fall through to
	// try as a source fragment.  Stop and return on any other error.
	if !opt.Fragment || !strings.Contains(err.Error(), "expected 'package'") {
		return nil, nil, err
	}

	// If this is a declaration list, make it a source file
	// by inserting a package clause.
	// Insert using a ;, not a newline, so that parse errors are on
	// the correct line.
	const prefix = "package main;"
	psrc := append([]byte(prefix), src...)
	file, err = parser.ParseFile(fset, filename, psrc, parserMode)
	if err == nil {
		// Gofmt will turn the ; into a \n.
		// Do that ourselves now and update the file contents,
		// so that positions and line numbers are correct going forward.
		psrc[len(prefix)-1] = '\n'
		fset.File(file.Package).SetLinesForContent(psrc)

		// If a main function exists, we will assume this is a main
		// package and leave the file.
		if containsMainFunc(file) {
			return file, nil, nil
		}

		adjust := func(orig, src []byte) []byte {
			// Remove the package clause.
			src = src[len(prefix):]
			return matchSpace(orig, src)
		}
		return file, adjust, nil
	}
	// If the error is that the source file didn't begin with a
	// declaration, fall through to try as a statement list.
	// Stop and return on any other error.
	if !strings.Contains(err.Error(), "expected declaration") {
		return nil, nil, err
	}

	// If this is a statement list, make it a source file
	// by inserting a package clause and turning the list
	// into a function body.  This handles expressions too.
	// Insert using a ;, not a newline, so that the line numbers
	// in fsrc match the ones in src.
	fsrc := append(append([]byte("package p; func _() {"), src...), '}')
	file, err = parser.ParseFile(fset, filename, fsrc, parserMode)
	if err == nil {
		adjust := func(orig, src []byte) []byte {
			// Remove the wrapping.
			// Gofmt has turned the ; into a \n\n.
			src = src[len("package p\n\nfunc _() {"):]
			src = src[:len(src)-len("}\n")]
			// Gofmt has also indented the function body one level.
			// Remove that indent.
			src = bytes.Replace(src, []byte("\n\t"), []byte("\n"), -1)
			return matchSpace(orig, src)
		}
		return file, adjust, nil
	}

	// Failed, and out of options.
	return nil, nil, err
}

// containsMainFunc checks if a file contains a function declaration with the
// function signature 'func main()'
func containsMainFunc(file *ast.File) bool {
	for _, decl := range file.Decls {
		if f, ok := decl.(*ast.FuncDecl); ok {
			if f.Name.Name != "main" {
				continue
			}

			if len(f.Type.Params.List) != 0 {
				continue
			}

			if f.Type.Results != nil && len(f.Type.Results.List) != 0 {
				continue
			}

			return true
		}
	}

	return false
}

func cutSpace(b []byte) (before, middle, after []byte) {
	i := 0
	for i < len(b) && (b[i] == ' ' || b[i] == '\t' || b[i] == '\n') {
		i++
	}
	j := len(b)
	for j > 0 && (b[j-1] == ' ' || b[j-1] == '\t' || b[j-1] == '\n') {
		j--
	}
	if i <= j {
		return b[:i], b[i:j], b[j:]
	}
	return nil, nil, b[j:]
}

// matchSpace reformats src to use the same space context as orig.
// 1) If orig begins with blank lines, matchSpace inserts them at the beginning of src.
// 2) matchSpace copies the indentation of the first non-blank line in orig
//    to every non-blank line in src.
// 3) matchSpace copies the trailing space from orig and uses it in place
//   of src's trailing space.
func matchSpace(orig []byte, src []byte) []byte {
	before, _, after := cutSpace(orig)
	i := bytes.LastIndex(before, []byte{'\n'})
	before, indent := before[:i+1], before[i+1:]

	_, src, _ = cutSpace(src)

	var b bytes.Buffer
	b.Write(before)
	for len(src) > 0 {
		line := src
		if i := bytes.IndexByte(line, '\n'); i >= 0 {
			line, src = line[:i+1], line[i+1:]
		} else {
			src = nil
		}
		if len(line) > 0 && line[0] != '\n' { // not blank
			b.Write(indent)
		}
		b.Write(line)
	}
	b.Write(after)
	return b.Bytes()
}

var impLine = regexp.MustCompile(`^\s+(?:[\w\.]+\s+)?"(.+)"`)

func addImportSpaces(r io.Reader, breaks []string) ([]byte, error) {
	var out bytes.Buffer
	in := bufio.NewReader(r)
	inImports := false
	done := false
	for {
		s, err := in.ReadString('\n')
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		if !inImports && !done && strings.HasPrefix(s, "import") {
			inImports = true
		}
		if inImports && (strings.HasPrefix(s, "var") ||
			strings.HasPrefix(s, "func") ||
			strings.HasPrefix(s, "const") ||
			strings.HasPrefix(s, "type")) {
			done = true
			inImports = false
		}
		if inImports && len(breaks) > 0 {
			if m := impLine.FindStringSubmatch(s); m != nil {
				if m[1] == breaks[0] {
					out.WriteByte('\n')
					breaks = breaks[1:]
				}
			}
		}

		fmt.Fprint(&out, s)
	}
	return out.Bytes(), nil
}
<commit_msg>internal/imports: don't set a logger unless the user has provided it<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:generate go run mkstdlib.go

// Package imports implements a Go pretty-printer (like package "go/format")
// that also adds or removes import statements as necessary.
package imports

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
)

// Options is golang.org/x/tools/imports.Options with extra internal-only options.
type Options struct {
	Env *ProcessEnv // The environment to use. Note: this contains the cached module and filesystem state.

	Fragment  bool // Accept fragment of a source file (no package statement)
	AllErrors bool // Report all errors (not just the first 10 on different lines)

	Comments  bool // Print comments (true if nil *Options provided)
	TabIndent bool // Use tabs for indent (true if nil *Options provided)
	TabWidth  int  // Tab width (8 if nil *Options provided)

	FormatOnly bool // Disable the insertion and deletion of imports
}

// Process implements golang.org/x/tools/imports.Process with explicit context in env.
func Process(filename string, src []byte, opt *Options) (formatted []byte, err error) {
	src, opt, err = initialize(filename, src, opt)
	if err != nil {
		return nil, err
	}

	fileSet := token.NewFileSet()
	file, adjust, err := parse(fileSet, filename, src, opt)
	if err != nil {
		return nil, err
	}

	if !opt.FormatOnly {
		if err := fixImports(fileSet, file, filename, opt.Env); err != nil {
			return nil, err
		}
	}
	return formatFile(fileSet, file, src, adjust, opt)
}

// FixImports returns a list of fixes to the imports that, when applied,
// will leave the imports in the same state as Process.
//
// Note that filename's directory influences which imports can be chosen,
// so it is important that filename be accurate.
func FixImports(filename string, src []byte, opt *Options) (fixes []*ImportFix, err error) {
	src, opt, err = initialize(filename, src, opt)
	if err != nil {
		return nil, err
	}

	fileSet := token.NewFileSet()
	file, _, err := parse(fileSet, filename, src, opt)
	if err != nil {
		return nil, err
	}

	return getFixes(fileSet, file, filename, opt.Env)
}

// ApplyFixes applies all of the fixes to the file and formats it. extraMode
// is added in when parsing the file.
func ApplyFixes(fixes []*ImportFix, filename string, src []byte, opt *Options, extraMode parser.Mode) (formatted []byte, err error) {
	src, opt, err = initialize(filename, src, opt)
	if err != nil {
		return nil, err
	}

	// Don't use parse() -- we don't care about fragments or statement lists
	// here, and we need to work with unparseable files.
	fileSet := token.NewFileSet()
	parserMode := parser.Mode(0)
	if opt.Comments {
		parserMode |= parser.ParseComments
	}
	if opt.AllErrors {
		parserMode |= parser.AllErrors
	}
	parserMode |= extraMode

	file, err := parser.ParseFile(fileSet, filename, src, parserMode)
	if file == nil {
		return nil, err
	}

	// Apply the fixes to the file.
	apply(fileSet, file, fixes)

	return formatFile(fileSet, file, src, nil, opt)
}

// GetAllCandidates gets all of the packages starting with prefix that can be
// imported by filename, sorted by import path.
func GetAllCandidates(ctx context.Context, callback func(ImportFix), searchPrefix, filename, filePkg string, opt *Options) error {
	_, opt, err := initialize(filename, []byte{}, opt)
	if err != nil {
		return err
	}
	return getAllCandidates(ctx, callback, searchPrefix, filename, filePkg, opt.Env)
}

// GetPackageExports returns all known packages with name pkg and their exports.
func GetPackageExports(ctx context.Context, callback func(PackageExport), searchPkg, filename, filePkg string, opt *Options) error {
	_, opt, err := initialize(filename, []byte{}, opt)
	if err != nil {
		return err
	}
	return getPackageExports(ctx, callback, searchPkg, filename, filePkg, opt.Env)
}

// initialize sets the values for opt and src.
// If they are provided, they are not changed. Otherwise opt is set to the
// default values and src is read from the file system.
func initialize(filename string, src []byte, opt *Options) ([]byte, *Options, error) {
	// Use defaults if opt is nil.
	if opt == nil {
		opt = &Options{Comments: true, TabIndent: true, TabWidth: 8}
	}

	// Set the env if the user has not provided it.
	if opt.Env == nil {
		opt.Env = &ProcessEnv{
			GOPATH:      build.Default.GOPATH,
			GOROOT:      build.Default.GOROOT,
			GOFLAGS:     os.Getenv("GOFLAGS"),
			GO111MODULE: os.Getenv("GO111MODULE"),
			GOPROXY:     os.Getenv("GOPROXY"),
			GOSUMDB:     os.Getenv("GOSUMDB"),
		}
	}
	if src == nil {
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			return nil, nil, err
		}
		src = b
	}

	return src, opt, nil
}

func formatFile(fileSet *token.FileSet, file *ast.File, src []byte, adjust func(orig []byte, src []byte) []byte, opt *Options) ([]byte, error) {
	mergeImports(opt.Env, fileSet, file)
	sortImports(opt.Env, fileSet, file)
	imps := astutil.Imports(fileSet, file)
	var spacesBefore []string // import paths we need spaces before
	for _, impSection := range imps {
		// Within each block of contiguous imports, see if any
		// import lines are in different group numbers. If so,
		// we'll need to put a space between them so it's
		// compatible with gofmt.
		lastGroup := -1
		for _, importSpec := range impSection {
			importPath, _ := strconv.Unquote(importSpec.Path.Value)
			groupNum := importGroup(opt.Env, importPath)
			if groupNum != lastGroup && lastGroup != -1 {
				spacesBefore = append(spacesBefore, importPath)
			}
			lastGroup = groupNum
		}

	}

	printerMode := printer.UseSpaces
	if opt.TabIndent {
		printerMode |= printer.TabIndent
	}
	printConfig := &printer.Config{Mode: printerMode, Tabwidth: opt.TabWidth}

	var buf bytes.Buffer
	err := printConfig.Fprint(&buf, fileSet, file)
	if err != nil {
		return nil, err
	}
	out := buf.Bytes()
	if adjust != nil {
		out = adjust(src, out)
	}
	if len(spacesBefore) > 0 {
		out, err = addImportSpaces(bytes.NewReader(out), spacesBefore)
		if err != nil {
			return nil, err
		}
	}

	out, err = format.Source(out)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// parse parses src, which was read from filename,
// as a Go source file or statement list.
func parse(fset *token.FileSet, filename string, src []byte, opt *Options) (*ast.File, func(orig, src []byte) []byte, error) {
	parserMode := parser.Mode(0)
	if opt.Comments {
		parserMode |= parser.ParseComments
	}
	if opt.AllErrors {
		parserMode |= parser.AllErrors
	}

	// Try as whole source file.
	file, err := parser.ParseFile(fset, filename, src, parserMode)
	if err == nil {
		return file, nil, nil
	}
	// If the error is that the source file didn't begin with a
	// package line and we accept fragmented input, fall through to
	// try as a source fragment.  Stop and return on any other error.
	if !opt.Fragment || !strings.Contains(err.Error(), "expected 'package'") {
		return nil, nil, err
	}

	// If this is a declaration list, make it a source file
	// by inserting a package clause.
	// Insert using a ;, not a newline, so that parse errors are on
	// the correct line.
	const prefix = "package main;"
	psrc := append([]byte(prefix), src...)
	file, err = parser.ParseFile(fset, filename, psrc, parserMode)
	if err == nil {
		// Gofmt will turn the ; into a \n.
		// Do that ourselves now and update the file contents,
		// so that positions and line numbers are correct going forward.
		psrc[len(prefix)-1] = '\n'
		fset.File(file.Package).SetLinesForContent(psrc)

		// If a main function exists, we will assume this is a main
		// package and leave the file.
		if containsMainFunc(file) {
			return file, nil, nil
		}

		adjust := func(orig, src []byte) []byte {
			// Remove the package clause.
			src = src[len(prefix):]
			return matchSpace(orig, src)
		}
		return file, adjust, nil
	}
	// If the error is that the source file didn't begin with a
	// declaration, fall through to try as a statement list.
	// Stop and return on any other error.
	if !strings.Contains(err.Error(), "expected declaration") {
		return nil, nil, err
	}

	// If this is a statement list, make it a source file
	// by inserting a package clause and turning the list
	// into a function body.  This handles expressions too.
	// Insert using a ;, not a newline, so that the line numbers
	// in fsrc match the ones in src.
	fsrc := append(append([]byte("package p; func _() {"), src...), '}')
	file, err = parser.ParseFile(fset, filename, fsrc, parserMode)
	if err == nil {
		adjust := func(orig, src []byte) []byte {
			// Remove the wrapping.
			// Gofmt has turned the ; into a \n\n.
			src = src[len("package p\n\nfunc _() {"):]
			src = src[:len(src)-len("}\n")]
			// Gofmt has also indented the function body one level.
			// Remove that indent.
			src = bytes.Replace(src, []byte("\n\t"), []byte("\n"), -1)
			return matchSpace(orig, src)
		}
		return file, adjust, nil
	}

	// Failed, and out of options.
	return nil, nil, err
}

// containsMainFunc checks if a file contains a function declaration with the
// function signature 'func main()'
func containsMainFunc(file *ast.File) bool {
	for _, decl := range file.Decls {
		if f, ok := decl.(*ast.FuncDecl); ok {
			if f.Name.Name != "main" {
				continue
			}

			if len(f.Type.Params.List) != 0 {
				continue
			}

			if f.Type.Results != nil && len(f.Type.Results.List) != 0 {
				continue
			}

			return true
		}
	}

	return false
}

func cutSpace(b []byte) (before, middle, after []byte) {
	i := 0
	for i < len(b) && (b[i] == ' ' || b[i] == '\t' || b[i] == '\n') {
		i++
	}
	j := len(b)
	for j > 0 && (b[j-1] == ' ' || b[j-1] == '\t' || b[j-1] == '\n') {
		j--
	}
	if i <= j {
		return b[:i], b[i:j], b[j:]
	}
	return nil, nil, b[j:]
}

// matchSpace reformats src to use the same space context as orig.
// 1) If orig begins with blank lines, matchSpace inserts them at the beginning of src.
// 2) matchSpace copies the indentation of the first non-blank line in orig
//    to every non-blank line in src.
// 3) matchSpace copies the trailing space from orig and uses it in place
//   of src's trailing space.
func matchSpace(orig []byte, src []byte) []byte {
	before, _, after := cutSpace(orig)
	i := bytes.LastIndex(before, []byte{'\n'})
	before, indent := before[:i+1], before[i+1:]

	_, src, _ = cutSpace(src)

	var b bytes.Buffer
	b.Write(before)
	for len(src) > 0 {
		line := src
		if i := bytes.IndexByte(line, '\n'); i >= 0 {
			line, src = line[:i+1], line[i+1:]
		} else {
			src = nil
		}
		if len(line) > 0 && line[0] != '\n' { // not blank
			b.Write(indent)
		}
		b.Write(line)
	}
	b.Write(after)
	return b.Bytes()
}

var impLine = regexp.MustCompile(`^\s+(?:[\w\.]+\s+)?"(.+)"`)

func addImportSpaces(r io.Reader, breaks []string) ([]byte, error) {
	var out bytes.Buffer
	in := bufio.NewReader(r)
	inImports := false
	done := false
	for {
		s, err := in.ReadString('\n')
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		if !inImports && !done && strings.HasPrefix(s, "import") {
			inImports = true
		}
		if inImports && (strings.HasPrefix(s, "var") ||
			strings.HasPrefix(s, "func") ||
			strings.HasPrefix(s, "const") ||
			strings.HasPrefix(s, "type")) {
			done = true
			inImports = false
		}
		if inImports && len(breaks) > 0 {
			if m := impLine.FindStringSubmatch(s); m != nil {
				if m[1] == breaks[0] {
					out.WriteByte('\n')
					breaks = breaks[1:]
				}
			}
		}

		fmt.Fprint(&out, s)
	}
	return out.Bytes(), nil
}
<|endoftext|>


<commit_before>package manager

import (
	"fmt"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/camptocamp/bivac/internal/utils"
	"github.com/camptocamp/bivac/pkg/orchestrators"
	"github.com/camptocamp/bivac/pkg/volume"
)

// Orchestrators groups the parameters of all supported orchestrators in one structure
type Orchestrators struct {
	Docker     orchestrators.DockerConfig
	Cattle     orchestrators.CattleConfig
	Kubernetes orchestrators.KubernetesConfig
}

// Manager contains all informations used by the Bivac manager
type Manager struct {
	Orchestrator orchestrators.Orchestrator
	Volumes      []*volume.Volume
	Server       *Server
	Providers    *Providers
	TargetURL    string
	RetryCount   int
	LogServer    string
	BuildInfo    utils.BuildInfo
	AgentImage   string

	backupSlots chan map[string][]*volume.Volume
}

// Start starts a Bivac manager which handle backups management
func Start(buildInfo utils.BuildInfo, o orchestrators.Orchestrator, s Server, volumeFilters volume.Filters, providersFile, targetURL, logServer, agentImage string, retryCount int) (err error) {
	p, err := LoadProviders(providersFile)
	if err != nil {
		err = fmt.Errorf("failed to read providers file: %s", err)
		return
	}

	m := &Manager{
		Orchestrator: o,
		Server:       &s,
		Providers:    &p,
		TargetURL:    targetURL,
		RetryCount:   retryCount,
		LogServer:    logServer,
		BuildInfo:    buildInfo,
		AgentImage:   agentImage,

		backupSlots: make(chan map[string][]*volume.Volume),
	}

	// Manage volumes
	go func(m *Manager, volumeFilters volume.Filters) {
		var bs map[string][]*volume.Volume

		log.Debugf("Starting volume manager...")

		for {
			err = retrieveVolumes(m, volumeFilters)
			if err != nil {
				log.Errorf("failed to retrieve volumes: %s", err)
			}

			bs = make(map[string][]*volume.Volume)
			for _, v := range m.Volumes {
				if !isBackupNeeded(v) {
					continue
				}
				bs[v.HostBind] = append(bs[v.HostBind], v)
			}

			for node := range bs {
				if ok, _ := m.Orchestrator.IsNodeAvailable(node); !ok {
					log.WithFields(log.Fields{
						"node": node,
					}).Warning("Node unavailable.")
					delete(bs, node)
				}
			}

			if len(bs) > 0 {
				select {
				case m.backupSlots <- bs:
				default:
				}
			}

			time.Sleep(10 * time.Minute)
		}
	}(m, volumeFilters)

	// Manage backups
	go func(m *Manager) {
		var wg sync.WaitGroup

		log.Debugf("Starting backup manager...")
		for {
			bs := <-m.backupSlots

			for _, volumes := range bs {
				wg.Add(len(volumes))
				// Release goroutines if they take more than 12h to run
				go func() {
					timeout := time.After(12 * time.Hour)
					<-timeout
					for i := 0; i < len(volumes); i++ {
						wg.Done()
					}
				}()
				go func(volumes []*volume.Volume) {
					instanceSem := make(chan bool, 2)
					for _, v := range volumes {
						instanceSem <- true
						go func(v *volume.Volume) {
							var timedout bool
							tearDown := make(chan bool)

							log.WithFields(log.Fields{
								"volume":   v.Name,
								"hostname": v.Hostname,
							}).Debugf("Backing up volume.")
							defer func() {
								if !timedout {
									tearDown <- true
									<-instanceSem
									wg.Done()
								}
							}()

							// Workaround which avoid a stucked backup to block the whole backup process
							// If the backup process takes more than one hour,
							// the backup slot is released.
							go func() {
								timeout := time.After(1 * time.Hour)
								select {
								case <-tearDown:
									return
								case <-timeout:
									timedout = true
									<-instanceSem
									wg.Done()
								}
							}()

							err = nil
							for i := 0; i <= m.RetryCount; i++ {
								err = backupVolume(m, v, false)
								if err != nil {
									log.WithFields(log.Fields{
										"volume":   v.Name,
										"hostname": v.Hostname,
										"try":      i + 1,
									}).Errorf("failed to backup volume: %s", err)

									time.Sleep(2 * time.Second)
								} else {
									break
								}
							}
						}(v)
					}
				}(volumes)
			}
			wg.Wait()
		}
	}(m)

	// Manage API server
	m.StartServer()

	return
}

func isBackupNeeded(v *volume.Volume) bool {
	if v.BackingUp {
		return false
	}

	if v.LastBackupDate == "" {
		return true
	}

	lbd, err := time.Parse("2006-01-02 15:04:05", v.LastBackupDate)
	if err != nil {
		log.WithFields(log.Fields{
			"volume":   v.Name,
			"hostname": v.Hostname,
		}).Errorf("failed to parse backup date of volume `%s': %s", v.Name, err)
		return false
	}

	if lbd.Add(time.Hour * 23).Before(time.Now()) {
		return true
	}

	if lbd.Add(time.Hour).Before(time.Now()) && v.LastBackupStatus == "Failed" {
		return true
	}
	return false
}

// GetOrchestrator returns an orchestrator interface based on the name you specified or on the orchestrator Bivac is running on
func GetOrchestrator(name string, orchs Orchestrators) (o orchestrators.Orchestrator, err error) {
	if name != "" {
		log.Debugf("Choosing orchestrator based on configuration...")
		switch name {
		case "docker":
			o, err = orchestrators.NewDockerOrchestrator(&orchs.Docker)
		case "cattle":
			o, err = orchestrators.NewCattleOrchestrator(&orchs.Cattle)
		case "kubernetes":
			o, err = orchestrators.NewKubernetesOrchestrator(&orchs.Kubernetes)
		default:
			err = fmt.Errorf("'%s' is not a valid orchestrator", name)
			return
		}
	} else {
		log.Debugf("Trying to detect orchestrator based on environment...")

		if orchestrators.DetectCattle() {
			o, err = orchestrators.NewCattleOrchestrator(&orchs.Cattle)
		} else if orchestrators.DetectKubernetes() {
			o, err = orchestrators.NewKubernetesOrchestrator(&orchs.Kubernetes)
		} else if orchestrators.DetectDocker(&orchs.Docker) {
			o, err = orchestrators.NewDockerOrchestrator(&orchs.Docker)
		} else {
			err = fmt.Errorf("no orchestrator detected")
			return
		}
	}
	if err != nil {
		log.Infof("Using orchestrator: %s", o.GetName())
	}
	return
}

// BackupVolume does a backup of a volume
func (m *Manager) BackupVolume(volumeID string, force bool) (err error) {
	for _, v := range m.Volumes {
		if v.ID == volumeID {
			log.WithFields(log.Fields{
				"volume":   v.Name,
				"hostname": v.Hostname,
			}).Debug("Backup manually requested.")
			err = backupVolume(m, v, force)
			if err != nil {
				err = fmt.Errorf("failed to backup volume: %s", err)
				return
			}
		}
	}
	return
}

// GetInformations returns informations regarding the Bivac manager
func (m *Manager) GetInformations() (informations map[string]string) {
	informations = map[string]string{
		"version":        m.BuildInfo.Version,
		"build_date":     m.BuildInfo.Date,
		"build_commit":   m.BuildInfo.CommitSha1,
		"golang_version": m.BuildInfo.Runtime,
		"orchestrator":   m.Orchestrator.GetName(),
		"address":        m.Server.Address,
		"volumes_count":  fmt.Sprintf("%d", len(m.Volumes)),
	}
	return
}
<commit_msg>improve scheduler (independant volume backup goroutine<commit_after>package manager

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/camptocamp/bivac/internal/utils"
	"github.com/camptocamp/bivac/pkg/orchestrators"
	"github.com/camptocamp/bivac/pkg/volume"
)

// Orchestrators groups the parameters of all supported orchestrators in one structure
type Orchestrators struct {
	Docker     orchestrators.DockerConfig
	Cattle     orchestrators.CattleConfig
	Kubernetes orchestrators.KubernetesConfig
}

// Manager contains all informations used by the Bivac manager
type Manager struct {
	Orchestrator orchestrators.Orchestrator
	Volumes      []*volume.Volume
	Server       *Server
	Providers    *Providers
	TargetURL    string
	RetryCount   int
	LogServer    string
	BuildInfo    utils.BuildInfo
	AgentImage   string

	backupSlots chan *volume.Volume
}

// Start starts a Bivac manager which handle backups management
func Start(buildInfo utils.BuildInfo, o orchestrators.Orchestrator, s Server, volumeFilters volume.Filters, providersFile, targetURL, logServer, agentImage string, retryCount int) (err error) {
	p, err := LoadProviders(providersFile)
	if err != nil {
		err = fmt.Errorf("failed to read providers file: %s", err)
		return
	}

	m := &Manager{
		Orchestrator: o,
		Server:       &s,
		Providers:    &p,
		TargetURL:    targetURL,
		RetryCount:   retryCount,
		LogServer:    logServer,
		BuildInfo:    buildInfo,
		AgentImage:   agentImage,

		backupSlots: make(chan *volume.Volume, 100),
	}

	// Manage volumes
	go func(m *Manager, volumeFilters volume.Filters) {

		log.Debugf("Starting volume manager...")

		for {
			err = retrieveVolumes(m, volumeFilters)
			if err != nil {
				log.Errorf("failed to retrieve volumes: %s", err)
			}

			for _, v := range m.Volumes {
				if !isBackupNeeded(v) {
					continue
				}

				m.backupSlots <- v
			}
			time.Sleep(10 * time.Minute)
		}
	}(m, volumeFilters)

	// Manage backups
	go func(m *Manager) {
		slots := make(map[string](chan bool))

		log.Infof("Starting backup manager...")

		for {
			v := <-m.backupSlots
			if _, ok := slots[v.HostBind]; !ok {
				slots[v.HostBind] = make(chan bool, 2)
			}
			select {
			case slots[v.HostBind] <- true:
			default:
				continue
			}
			if ok, _ := m.Orchestrator.IsNodeAvailable(v.HostBind); !ok {
				log.WithFields(log.Fields{
					"node": v.HostBind,
				}).Warning("Node unavailable.")
				<-slots[v.HostBind]
				continue
			}

			go func(v *volume.Volume) {
				var timedout bool
				tearDown := make(chan bool)

				log.WithFields(log.Fields{
					"volume":   v.Name,
					"hostname": v.Hostname,
				}).Debugf("Backing up volume.")
				defer func() {
					if !timedout {
						tearDown <- true
						<-slots[v.HostBind]
					}
				}()

				// Workaround which avoid a stucked backup to block the whole backup process
				// If the backup process takes more than one hour,
				// the backup slot is released.
				go func() {
					timeout := time.After(1 * time.Hour)
					select {
					case <-tearDown:
						return
					case <-timeout:
						timedout = true
						<-slots[v.HostBind]
					}
				}()

				err = nil
				for i := 0; i <= m.RetryCount; i++ {
					err = backupVolume(m, v, false)
					if err != nil {
						log.WithFields(log.Fields{
							"volume":   v.Name,
							"hostname": v.Hostname,
							"try":      i + 1,
						}).Errorf("failed to backup volume: %s", err)

						time.Sleep(2 * time.Second)
					} else {
						break
					}
				}
			}(v)
		}

		/*
			var wg sync.WaitGroup

			log.Debugf("Starting backup manager...")
			for {
				bs := <-m.backupSlots

				for _, volumes := range bs {
					wg.Add(len(volumes))
					// Release goroutines if they take more than 12h to run
					go func() {
						timeout := time.After(12 * time.Hour)
						<-timeout
						for i := 0; i < len(volumes); i++ {
							wg.Done()
						}
					}()
					go func(volumes []*volume.Volume) {
						instanceSem := make(chan bool, 2)
						for _, v := range volumes {
							instanceSem <- true
							go func(v *volume.Volume) {
								var timedout bool
								tearDown := make(chan bool)

								log.WithFields(log.Fields{
									"volume":   v.Name,
									"hostname": v.Hostname,
								}).Debugf("Backing up volume.")
								defer func() {
									if !timedout {
										tearDown <- true
										<-instanceSem
										wg.Done()
									}
								}()

								// Workaround which avoid a stucked backup to block the whole backup process
								// If the backup process takes more than one hour,
								// the backup slot is released.
								go func() {
									timeout := time.After(1 * time.Hour)
									select {
									case <-tearDown:
										return
									case <-timeout:
										timedout = true
										<-instanceSem
										wg.Done()
									}
								}()

								err = nil
								for i := 0; i <= m.RetryCount; i++ {
									err = backupVolume(m, v, false)
									if err != nil {
										log.WithFields(log.Fields{
											"volume":   v.Name,
											"hostname": v.Hostname,
											"try":      i + 1,
										}).Errorf("failed to backup volume: %s", err)

										time.Sleep(2 * time.Second)
									} else {
										break
									}
								}
							}(v)
						}
					}(volumes)
				}
				wg.Wait()
			}
		*/
	}(m)

	// Manage API server
	m.StartServer()

	return
}

func isBackupNeeded(v *volume.Volume) bool {
	if v.BackingUp {
		return false
	}

	if v.LastBackupDate == "" {
		return true
	}

	lbd, err := time.Parse("2006-01-02 15:04:05", v.LastBackupDate)
	if err != nil {
		log.WithFields(log.Fields{
			"volume":   v.Name,
			"hostname": v.Hostname,
		}).Errorf("failed to parse backup date of volume `%s': %s", v.Name, err)
		return false
	}

	if lbd.Add(time.Hour * 23).Before(time.Now()) {
		return true
	}

	if lbd.Add(time.Hour).Before(time.Now()) && v.LastBackupStatus == "Failed" {
		return true
	}
	return false
}

// GetOrchestrator returns an orchestrator interface based on the name you specified or on the orchestrator Bivac is running on
func GetOrchestrator(name string, orchs Orchestrators) (o orchestrators.Orchestrator, err error) {
	if name != "" {
		log.Debugf("Choosing orchestrator based on configuration...")
		switch name {
		case "docker":
			o, err = orchestrators.NewDockerOrchestrator(&orchs.Docker)
		case "cattle":
			o, err = orchestrators.NewCattleOrchestrator(&orchs.Cattle)
		case "kubernetes":
			o, err = orchestrators.NewKubernetesOrchestrator(&orchs.Kubernetes)
		default:
			err = fmt.Errorf("'%s' is not a valid orchestrator", name)
			return
		}
	} else {
		log.Debugf("Trying to detect orchestrator based on environment...")

		if orchestrators.DetectCattle() {
			o, err = orchestrators.NewCattleOrchestrator(&orchs.Cattle)
		} else if orchestrators.DetectKubernetes() {
			o, err = orchestrators.NewKubernetesOrchestrator(&orchs.Kubernetes)
		} else if orchestrators.DetectDocker(&orchs.Docker) {
			o, err = orchestrators.NewDockerOrchestrator(&orchs.Docker)
		} else {
			err = fmt.Errorf("no orchestrator detected")
			return
		}
	}
	if err != nil {
		log.Infof("Using orchestrator: %s", o.GetName())
	}
	return
}

// BackupVolume does a backup of a volume
func (m *Manager) BackupVolume(volumeID string, force bool) (err error) {
	for _, v := range m.Volumes {
		if v.ID == volumeID {
			log.WithFields(log.Fields{
				"volume":   v.Name,
				"hostname": v.Hostname,
			}).Debug("Backup manually requested.")
			err = backupVolume(m, v, force)
			if err != nil {
				err = fmt.Errorf("failed to backup volume: %s", err)
				return
			}
		}
	}
	return
}

// GetInformations returns informations regarding the Bivac manager
func (m *Manager) GetInformations() (informations map[string]string) {
	informations = map[string]string{
		"version":        m.BuildInfo.Version,
		"build_date":     m.BuildInfo.Date,
		"build_commit":   m.BuildInfo.CommitSha1,
		"golang_version": m.BuildInfo.Runtime,
		"orchestrator":   m.Orchestrator.GetName(),
		"address":        m.Server.Address,
		"volumes_count":  fmt.Sprintf("%d", len(m.Volumes)),
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2021 The Go 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 metrics provides a service for reporting metrics to
// Stackdriver, or locally during development.
package metrics

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"path"
	"time"

	"cloud.google.com/go/compute/metadata"
	"contrib.go.opencensus.io/exporter/prometheus"
	"contrib.go.opencensus.io/exporter/stackdriver"
	"go.opencensus.io/stats/view"
	"google.golang.org/appengine"
	mrpb "google.golang.org/genproto/googleapis/api/monitoredres"
)

// NewService initializes a *Service.
//
// The Service returned is configured to send metric data to
// StackDriver. When not running on GCE, it will host metrics through
// a prometheus HTTP handler.
//
// views will be passed to view.Register for export to the metric
// service.
func NewService(resource *MonitoredResource, views []*view.View) (*Service, error) {
	err := view.Register(views...)
	if err != nil {
		return nil, err
	}

	if !metadata.OnGCE() {
		view.SetReportingPeriod(5 * time.Second)
		pe, err := prometheus.NewExporter(prometheus.Options{})
		if err != nil {
			return nil, fmt.Errorf("prometheus.NewExporter: %w", err)
		}
		view.RegisterExporter(pe)
		return &Service{pExporter: pe}, nil
	}

	projID, err := metadata.ProjectID()
	if err != nil {
		return nil, err
	}
	if resource == nil {
		return nil, errors.New("resource is required, got nil")
	}
	sde, err := stackdriver.NewExporter(stackdriver.Options{
		ProjectID:         projID,
		MonitoredResource: resource,
		ReportingInterval: time.Minute, // Minimum interval for Stackdriver is 1 minute.
	})
	if err != nil {
		return nil, err
	}

	// Minimum interval for Stackdriver is 1 minute.
	view.SetReportingPeriod(time.Minute)
	// Start the metrics exporter.
	if err := sde.StartMetricsExporter(); err != nil {
		return nil, err
	}

	return &Service{sdExporter: sde}, nil
}

// Service controls metric exporters.
type Service struct {
	sdExporter *stackdriver.Exporter
	pExporter  *prometheus.Exporter
}

func (m *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if m.pExporter != nil {
		m.pExporter.ServeHTTP(w, r)
		return
	}
	http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}

// Stop flushes metrics and stops exporting. Stop should be called
// before exiting.
func (m *Service) Stop() {
	if sde := m.sdExporter; sde != nil {
		// Flush any unsent data before exiting.
		sde.Flush()

		sde.StopMetricsExporter()
	}
}

// MonitoredResource wraps a *mrpb.MonitoredResource to implement the
// monitoredresource.MonitoredResource interface.
type MonitoredResource mrpb.MonitoredResource

func (r *MonitoredResource) MonitoredResource() (resType string, labels map[string]string) {
	return r.Type, r.Labels
}

// GCEResource populates a MonitoredResource with GCE Metadata.
//
// The returned MonitoredResource will have the type set to "generic_task".
func GCEResource(jobName string) (*MonitoredResource, error) {
	projID, err := metadata.ProjectID()
	if err != nil {
		return nil, err
	}
	zone, err := metadata.Zone()
	if err != nil {
		return nil, err
	}
	inst, err := metadata.InstanceName()
	if err != nil {
		return nil, err
	}
	group, err := instanceGroupName()
	if err != nil {
		return nil, err
	} else if group == "" {
		group = projID
	}

	return (*MonitoredResource)(&mrpb.MonitoredResource{
		Type: "generic_task", // See: https://cloud.google.com/monitoring/api/resources#tag_generic_task
		Labels: map[string]string{
			"project_id": projID,
			"location":   zone,
			"namespace":  group,
			"job":        jobName,
			"task_id":    inst,
		},
	}), nil
}

// GAEResource returns a *MonitoredResource with fields populated and
// for StackDriver.
//
// The resource will be in StackDrvier's gae_instance type.
func GAEResource(ctx context.Context) (*MonitoredResource, error) {
	// appengine.IsAppEngine is confusingly false as we're using a custom
	// container and building without the appenginevm build constraint.
	// Check metadata.OnGCE instead.
	if !metadata.OnGCE() {
		return nil, fmt.Errorf("not running on appengine")
	}
	projID, err := metadata.ProjectID()
	if err != nil {
		return nil, err
	}
	return (*MonitoredResource)(&mrpb.MonitoredResource{
		Type: "gae_instance",
		Labels: map[string]string{
			"project_id":  projID,
			"module_id":   appengine.ModuleName(ctx),
			"version_id":  appengine.VersionID(ctx),
			"instance_id": appengine.InstanceID(),
			"location":    appengine.Datacenter(ctx),
		},
	}), nil
}

// instanceGroupName fetches the instanceGroupName from the instance
// metadata.
//
// The instance group manager applies a custom "created-by" attribute
// to the instance, which is not part of the metadata package API, and
// must be queried separately.
//
// An empty string will be returned if a metadata.NotDefinedError is
// returned when fetching metadata. An error will be returned if other
// errors occur when fetching metadata.
func instanceGroupName() (string, error) {
	ig, err := metadata.InstanceAttributeValue("created-by")
	if errors.As(err, new(metadata.NotDefinedError)) {
		return "", nil
	} else if err != nil {
		return "", err
	}
	// "created-by" format: "projects/{{InstanceID}}/zones/{{Zone}}/instanceGroupManagers/{{Instance Group Name}}
	ig = path.Base(ig)
	return ig, nil
}
<commit_msg>internal/metrics: fix leaking stackdriver exporter<commit_after>// Copyright 2021 The Go 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 metrics provides a service for reporting metrics to
// Stackdriver, or locally during development.
package metrics

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"path"
	"time"

	"cloud.google.com/go/compute/metadata"
	"contrib.go.opencensus.io/exporter/prometheus"
	"contrib.go.opencensus.io/exporter/stackdriver"
	"go.opencensus.io/stats/view"
	"google.golang.org/appengine"
	mrpb "google.golang.org/genproto/googleapis/api/monitoredres"
)

// NewService initializes a *Service.
//
// The Service returned is configured to send metric data to
// StackDriver. When not running on GCE, it will host metrics through
// a prometheus HTTP handler.
//
// views will be passed to view.Register for export to the metric
// service.
func NewService(resource *MonitoredResource, views []*view.View) (*Service, error) {
	err := view.Register(views...)
	if err != nil {
		return nil, err
	}

	if !metadata.OnGCE() {
		view.SetReportingPeriod(5 * time.Second)
		pe, err := prometheus.NewExporter(prometheus.Options{})
		if err != nil {
			return nil, fmt.Errorf("prometheus.NewExporter: %w", err)
		}
		view.RegisterExporter(pe)
		return &Service{pExporter: pe}, nil
	}

	projID, err := metadata.ProjectID()
	if err != nil {
		return nil, err
	}
	if resource == nil {
		return nil, errors.New("resource is required, got nil")
	}
	sde, err := stackdriver.NewExporter(stackdriver.Options{
		ProjectID:         projID,
		MonitoredResource: resource,
		ReportingInterval: time.Minute, // Minimum interval for Stackdriver is 1 minute.
	})
	if err != nil {
		return nil, err
	}

	// Minimum interval for Stackdriver is 1 minute.
	view.SetReportingPeriod(time.Minute)
	// Start the metrics exporter.
	if err := sde.StartMetricsExporter(); err != nil {
		sde.Close()
		return nil, err
	}

	return &Service{sdExporter: sde}, nil
}

// Service controls metric exporters.
type Service struct {
	sdExporter *stackdriver.Exporter
	pExporter  *prometheus.Exporter
}

func (m *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if m.pExporter != nil {
		m.pExporter.ServeHTTP(w, r)
		return
	}
	http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}

// Stop flushes metrics and stops exporting. Stop should be called
// before exiting.
func (m *Service) Stop() {
	if sde := m.sdExporter; sde != nil {
		// Flush any unsent data before exiting.
		sde.Flush()

		sde.StopMetricsExporter()
	}
}

// MonitoredResource wraps a *mrpb.MonitoredResource to implement the
// monitoredresource.MonitoredResource interface.
type MonitoredResource mrpb.MonitoredResource

func (r *MonitoredResource) MonitoredResource() (resType string, labels map[string]string) {
	return r.Type, r.Labels
}

// GCEResource populates a MonitoredResource with GCE Metadata.
//
// The returned MonitoredResource will have the type set to "generic_task".
func GCEResource(jobName string) (*MonitoredResource, error) {
	projID, err := metadata.ProjectID()
	if err != nil {
		return nil, err
	}
	zone, err := metadata.Zone()
	if err != nil {
		return nil, err
	}
	inst, err := metadata.InstanceName()
	if err != nil {
		return nil, err
	}
	group, err := instanceGroupName()
	if err != nil {
		return nil, err
	} else if group == "" {
		group = projID
	}

	return (*MonitoredResource)(&mrpb.MonitoredResource{
		Type: "generic_task", // See: https://cloud.google.com/monitoring/api/resources#tag_generic_task
		Labels: map[string]string{
			"project_id": projID,
			"location":   zone,
			"namespace":  group,
			"job":        jobName,
			"task_id":    inst,
		},
	}), nil
}

// GAEResource returns a *MonitoredResource with fields populated and
// for StackDriver.
//
// The resource will be in StackDrvier's gae_instance type.
func GAEResource(ctx context.Context) (*MonitoredResource, error) {
	// appengine.IsAppEngine is confusingly false as we're using a custom
	// container and building without the appenginevm build constraint.
	// Check metadata.OnGCE instead.
	if !metadata.OnGCE() {
		return nil, fmt.Errorf("not running on appengine")
	}
	projID, err := metadata.ProjectID()
	if err != nil {
		return nil, err
	}
	return (*MonitoredResource)(&mrpb.MonitoredResource{
		Type: "gae_instance",
		Labels: map[string]string{
			"project_id":  projID,
			"module_id":   appengine.ModuleName(ctx),
			"version_id":  appengine.VersionID(ctx),
			"instance_id": appengine.InstanceID(),
			"location":    appengine.Datacenter(ctx),
		},
	}), nil
}

// instanceGroupName fetches the instanceGroupName from the instance
// metadata.
//
// The instance group manager applies a custom "created-by" attribute
// to the instance, which is not part of the metadata package API, and
// must be queried separately.
//
// An empty string will be returned if a metadata.NotDefinedError is
// returned when fetching metadata. An error will be returned if other
// errors occur when fetching metadata.
func instanceGroupName() (string, error) {
	ig, err := metadata.InstanceAttributeValue("created-by")
	if errors.As(err, new(metadata.NotDefinedError)) {
		return "", nil
	} else if err != nil {
		return "", err
	}
	// "created-by" format: "projects/{{InstanceID}}/zones/{{Zone}}/instanceGroupManagers/{{Instance Group Name}}
	ig = path.Base(ig)
	return ig, nil
}
<|endoftext|>