<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 executor

import (
	"context"
	"fmt"
	"math/big"
	"time"

	"github.com/loadimpact/k6/ui/pb"

	"github.com/sirupsen/logrus"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/types"
)

func sumStagesDuration(stages []Stage) (result time.Duration) {
	for _, s := range stages {
		result += time.Duration(s.Duration.Duration)
	}
	return
}

func getStagesUnscaledMaxTarget(unscaledStartValue int64, stages []Stage) int64 {
	max := unscaledStartValue
	for _, s := range stages {
		if s.Target.Int64 > max {
			max = s.Target.Int64
		}
	}
	return max
}

// A helper function to avoid code duplication
func validateStages(stages []Stage) []error {
	var errors []error
	if len(stages) == 0 { //nolint:nestif
		errors = append(errors, fmt.Errorf("at least one stage has to be specified"))
	} else {
		for i, s := range stages {
			stageNum := i + 1
			if !s.Duration.Valid {
				errors = append(errors, fmt.Errorf("stage %d doesn't have a duration", stageNum))
			} else if s.Duration.Duration < 0 {
				errors = append(errors, fmt.Errorf("the duration for stage %d shouldn't be negative", stageNum))
			}
			if !s.Target.Valid {
				errors = append(errors, fmt.Errorf("stage %d doesn't have a target", stageNum))
			} else if s.Target.Int64 < 0 {
				errors = append(errors, fmt.Errorf("the target for stage %d shouldn't be negative", stageNum))
			}
		}
	}
	return errors
}

// getIterationRunner is a helper function that returns an iteration executor
// closure. It takes care of updating the execution state statistics and
// warning messages.
//
// TODO: emit the end-of-test iteration metrics here (https://github.com/loadimpact/k6/issues/1250)
func getIterationRunner(
	executionState *lib.ExecutionState, logger *logrus.Entry,
) func(context.Context, lib.ActiveVU) {
	return func(ctx context.Context, vu lib.ActiveVU) {
		err := vu.RunOnce()

		// TODO: track (non-ramp-down) errors from script iterations as a metric,
		// and have a default threshold that will abort the script when the error
		// rate exceeds a certain percentage

		select {
		case <-ctx.Done():
			// Don't log errors or emit iterations metrics from cancelled iterations
			executionState.AddInterruptedIterations(1)
		default:
			if err != nil {
				if s, ok := err.(fmt.Stringer); ok {
					logger.Error(s.String())
				} else {
					logger.Error(err.Error())
				}
				// TODO: investigate context cancelled errors
			}

			// TODO: move emission of end-of-iteration metrics here?
			executionState.AddFullIterations(1)
		}
	}
}

// getDurationContexts is used to create sub-contexts that can restrict an
// executor to only run for its allotted time.
//
// If the executor doesn't have a graceful stop period for iterations, then
// both returned sub-contexts will be the same one, with a timeout equal to
// the supplied regular executor duration.
//
// But if a graceful stop is enabled, then the first returned context (and the
// cancel func) will be for the "outer" sub-context. Its timeout will include
// both the regular duration and the specified graceful stop period. The second
// context will be a sub-context of the first one and its timeout will include
// only the regular duration.
//
// In either case, the usage of these contexts should be like this:
//  - As long as the regDurationCtx isn't done, new iterations can be started.
//  - After regDurationCtx is done, no new iterations should be started; every
//    VU that finishes an iteration from now on can be returned to the buffer
//    pool in the ExecutionState struct.
//  - After maxDurationCtx is done, any VUs with iterations will be
//    interrupted by the context's closing and will be returned to the buffer.
//  - If you want to interrupt the execution of all VUs prematurely (e.g. there
//    was an error or something like that), trigger maxDurationCancel().
//  - If the whole test is aborted, the parent context will be cancelled, so
//    that will also cancel these contexts, thus the "general abort" case is
//    handled transparently.
func getDurationContexts(parentCtx context.Context, regularDuration, gracefulStop time.Duration) (
	startTime time.Time, maxDurationCtx, regDurationCtx context.Context, maxDurationCancel func(),
) {
	startTime = time.Now()
	maxEndTime := startTime.Add(regularDuration + gracefulStop)

	maxDurationCtx, maxDurationCancel = context.WithDeadline(parentCtx, maxEndTime)
	if gracefulStop == 0 {
		return startTime, maxDurationCtx, maxDurationCtx, maxDurationCancel
	}
	regDurationCtx, _ = context.WithDeadline(maxDurationCtx, startTime.Add(regularDuration)) //nolint:govet
	return startTime, maxDurationCtx, regDurationCtx, maxDurationCancel
}

// trackProgress is a helper function that monitors certain end-events in an
// executor and updates its progressbar accordingly.
func trackProgress(
	parentCtx, maxDurationCtx, regDurationCtx context.Context,
	exec lib.Executor, snapshot func() (float64, []string),
) {
	progressBar := exec.GetProgress()
	logger := exec.GetLogger()

	<-regDurationCtx.Done() // Wait for the regular context to be over
	gracefulStop := exec.GetConfig().GetGracefulStop()
	if parentCtx.Err() == nil && gracefulStop > 0 {
		p, right := snapshot()
		logger.WithField("gracefulStop", gracefulStop).Debug(
			"Regular duration is done, waiting for iterations to gracefully finish",
		)
		progressBar.Modify(
			pb.WithStatus(pb.Stopping),
			pb.WithConstProgress(p, right...),
		)
	}

	<-maxDurationCtx.Done()
	p, right := snapshot()
	constProg := pb.WithConstProgress(p, right...)
	select {
	case <-parentCtx.Done():
		progressBar.Modify(pb.WithStatus(pb.Interrupted), constProg)
	default:
		status := pb.WithStatus(pb.Done)
		if p < 1 {
			status = pb.WithStatus(pb.Interrupted)
		}
		progressBar.Modify(status, constProg)
	}
}

// getScaledArrivalRate returns a rational number containing the scaled value of
// the given rate over the given period. This should generally be the first
// function that's called, before we do any calculations with the users-supplied
// rates in the arrival-rate executors.
func getScaledArrivalRate(es *lib.ExecutionSegment, rate int64, period time.Duration) *big.Rat {
	return es.InPlaceScaleRat(big.NewRat(rate, int64(period)))
}

// getTickerPeriod is just a helper function that returns the ticker interval
// we need for given arrival-rate parameters.
//
// It's possible for this function to return a zero duration (i.e. valid=false)
// and 0 isn't a valid ticker period. This happens so we don't divide by 0 when
// the arrival-rate period is 0. This case has to be handled separately.
func getTickerPeriod(scaledArrivalRate *big.Rat) types.NullDuration {
	if scaledArrivalRate.Sign() == 0 {
		return types.NewNullDuration(0, false)
	}
	// Basically, the ticker rate is time.Duration(1/arrivalRate). Considering
	// that time.Duration is represented as int64 nanoseconds, no meaningful
	// precision is likely to be lost here...
	result, _ := new(big.Rat).Inv(scaledArrivalRate).Float64()
	return types.NewNullDuration(time.Duration(result), true)
}

// getArrivalRatePerSec returns the iterations per second rate.
func getArrivalRatePerSec(scaledArrivalRate *big.Rat) *big.Rat {
	perSecRate := big.NewRat(int64(time.Second), 1)
	return perSecRate.Mul(perSecRate, scaledArrivalRate)
}

func getVUActivationParams(
	ctx context.Context, conf BaseConfig, deactivateCallback func(lib.InitializedVU),
) *lib.VUActivationParams {
	return &lib.VUActivationParams{
		RunContext:         ctx,
		Exec:               conf.GetExec(),
		Env:                conf.GetEnv(),
		Tags:               conf.GetTags(),
		DeactivateCallback: deactivateCallback,
	}
}
<commit_msg>Remove an unnecesary else statement (#1461)<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 executor

import (
	"context"
	"fmt"
	"math/big"
	"time"

	"github.com/loadimpact/k6/ui/pb"

	"github.com/sirupsen/logrus"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/types"
)

func sumStagesDuration(stages []Stage) (result time.Duration) {
	for _, s := range stages {
		result += time.Duration(s.Duration.Duration)
	}
	return
}

func getStagesUnscaledMaxTarget(unscaledStartValue int64, stages []Stage) int64 {
	max := unscaledStartValue
	for _, s := range stages {
		if s.Target.Int64 > max {
			max = s.Target.Int64
		}
	}
	return max
}

// A helper function to avoid code duplication
func validateStages(stages []Stage) []error {
	var errors []error
	if len(stages) == 0 {
		errors = append(errors, fmt.Errorf("at least one stage has to be specified"))
		return errors
	}

	for i, s := range stages {
		stageNum := i + 1
		if !s.Duration.Valid {
			errors = append(errors, fmt.Errorf("stage %d doesn't have a duration", stageNum))
		} else if s.Duration.Duration < 0 {
			errors = append(errors, fmt.Errorf("the duration for stage %d shouldn't be negative", stageNum))
		}
		if !s.Target.Valid {
			errors = append(errors, fmt.Errorf("stage %d doesn't have a target", stageNum))
		} else if s.Target.Int64 < 0 {
			errors = append(errors, fmt.Errorf("the target for stage %d shouldn't be negative", stageNum))
		}
	}
	return errors
}

// getIterationRunner is a helper function that returns an iteration executor
// closure. It takes care of updating the execution state statistics and
// warning messages.
//
// TODO: emit the end-of-test iteration metrics here (https://github.com/loadimpact/k6/issues/1250)
func getIterationRunner(
	executionState *lib.ExecutionState, logger *logrus.Entry,
) func(context.Context, lib.ActiveVU) {
	return func(ctx context.Context, vu lib.ActiveVU) {
		err := vu.RunOnce()

		// TODO: track (non-ramp-down) errors from script iterations as a metric,
		// and have a default threshold that will abort the script when the error
		// rate exceeds a certain percentage

		select {
		case <-ctx.Done():
			// Don't log errors or emit iterations metrics from cancelled iterations
			executionState.AddInterruptedIterations(1)
		default:
			if err != nil {
				if s, ok := err.(fmt.Stringer); ok {
					logger.Error(s.String())
				} else {
					logger.Error(err.Error())
				}
				// TODO: investigate context cancelled errors
			}

			// TODO: move emission of end-of-iteration metrics here?
			executionState.AddFullIterations(1)
		}
	}
}

// getDurationContexts is used to create sub-contexts that can restrict an
// executor to only run for its allotted time.
//
// If the executor doesn't have a graceful stop period for iterations, then
// both returned sub-contexts will be the same one, with a timeout equal to
// the supplied regular executor duration.
//
// But if a graceful stop is enabled, then the first returned context (and the
// cancel func) will be for the "outer" sub-context. Its timeout will include
// both the regular duration and the specified graceful stop period. The second
// context will be a sub-context of the first one and its timeout will include
// only the regular duration.
//
// In either case, the usage of these contexts should be like this:
//  - As long as the regDurationCtx isn't done, new iterations can be started.
//  - After regDurationCtx is done, no new iterations should be started; every
//    VU that finishes an iteration from now on can be returned to the buffer
//    pool in the ExecutionState struct.
//  - After maxDurationCtx is done, any VUs with iterations will be
//    interrupted by the context's closing and will be returned to the buffer.
//  - If you want to interrupt the execution of all VUs prematurely (e.g. there
//    was an error or something like that), trigger maxDurationCancel().
//  - If the whole test is aborted, the parent context will be cancelled, so
//    that will also cancel these contexts, thus the "general abort" case is
//    handled transparently.
func getDurationContexts(parentCtx context.Context, regularDuration, gracefulStop time.Duration) (
	startTime time.Time, maxDurationCtx, regDurationCtx context.Context, maxDurationCancel func(),
) {
	startTime = time.Now()
	maxEndTime := startTime.Add(regularDuration + gracefulStop)

	maxDurationCtx, maxDurationCancel = context.WithDeadline(parentCtx, maxEndTime)
	if gracefulStop == 0 {
		return startTime, maxDurationCtx, maxDurationCtx, maxDurationCancel
	}
	regDurationCtx, _ = context.WithDeadline(maxDurationCtx, startTime.Add(regularDuration)) //nolint:govet
	return startTime, maxDurationCtx, regDurationCtx, maxDurationCancel
}

// trackProgress is a helper function that monitors certain end-events in an
// executor and updates its progressbar accordingly.
func trackProgress(
	parentCtx, maxDurationCtx, regDurationCtx context.Context,
	exec lib.Executor, snapshot func() (float64, []string),
) {
	progressBar := exec.GetProgress()
	logger := exec.GetLogger()

	<-regDurationCtx.Done() // Wait for the regular context to be over
	gracefulStop := exec.GetConfig().GetGracefulStop()
	if parentCtx.Err() == nil && gracefulStop > 0 {
		p, right := snapshot()
		logger.WithField("gracefulStop", gracefulStop).Debug(
			"Regular duration is done, waiting for iterations to gracefully finish",
		)
		progressBar.Modify(
			pb.WithStatus(pb.Stopping),
			pb.WithConstProgress(p, right...),
		)
	}

	<-maxDurationCtx.Done()
	p, right := snapshot()
	constProg := pb.WithConstProgress(p, right...)
	select {
	case <-parentCtx.Done():
		progressBar.Modify(pb.WithStatus(pb.Interrupted), constProg)
	default:
		status := pb.WithStatus(pb.Done)
		if p < 1 {
			status = pb.WithStatus(pb.Interrupted)
		}
		progressBar.Modify(status, constProg)
	}
}

// getScaledArrivalRate returns a rational number containing the scaled value of
// the given rate over the given period. This should generally be the first
// function that's called, before we do any calculations with the users-supplied
// rates in the arrival-rate executors.
func getScaledArrivalRate(es *lib.ExecutionSegment, rate int64, period time.Duration) *big.Rat {
	return es.InPlaceScaleRat(big.NewRat(rate, int64(period)))
}

// getTickerPeriod is just a helper function that returns the ticker interval
// we need for given arrival-rate parameters.
//
// It's possible for this function to return a zero duration (i.e. valid=false)
// and 0 isn't a valid ticker period. This happens so we don't divide by 0 when
// the arrival-rate period is 0. This case has to be handled separately.
func getTickerPeriod(scaledArrivalRate *big.Rat) types.NullDuration {
	if scaledArrivalRate.Sign() == 0 {
		return types.NewNullDuration(0, false)
	}
	// Basically, the ticker rate is time.Duration(1/arrivalRate). Considering
	// that time.Duration is represented as int64 nanoseconds, no meaningful
	// precision is likely to be lost here...
	result, _ := new(big.Rat).Inv(scaledArrivalRate).Float64()
	return types.NewNullDuration(time.Duration(result), true)
}

// getArrivalRatePerSec returns the iterations per second rate.
func getArrivalRatePerSec(scaledArrivalRate *big.Rat) *big.Rat {
	perSecRate := big.NewRat(int64(time.Second), 1)
	return perSecRate.Mul(perSecRate, scaledArrivalRate)
}

func getVUActivationParams(
	ctx context.Context, conf BaseConfig, deactivateCallback func(lib.InitializedVU),
) *lib.VUActivationParams {
	return &lib.VUActivationParams{
		RunContext:         ctx,
		Exec:               conf.GetExec(),
		Env:                conf.GetEnv(),
		Tags:               conf.GetTags(),
		DeactivateCallback: deactivateCallback,
	}
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"io"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/nomad/api"
	"github.com/mitchellh/cli"
)

type MonitorCommand struct {
	Meta
}

func (c *MonitorCommand) Help() string {
	helpText := `
Usage: nomad monitor [options]

	Stream log messages of a nomad agent. The monitor command lets you
	listen for log levels that may be filtered out of the Nomad agent. For
	example your agent may only be logging at INFO level, but with the monitor
	command you can set -log-level DEBUG

General Options:

	` + generalOptionsUsage() + `

Monitor Specific Options:

  -log-level <level>
    Sets the log level to monitor (default: INFO)

  -node-id <node-id>
    Sets the specific node to monitor

  -json
    Sets log output to JSON format
  `
	return strings.TrimSpace(helpText)
}

func (c *MonitorCommand) Synopsis() string {
	return "stream logs from a Nomad agent"
}

func (c *MonitorCommand) Name() string { return "monitor" }

func (c *MonitorCommand) Run(args []string) int {
	c.Ui = &cli.PrefixedUi{
		OutputPrefix: "    ",
		InfoPrefix:   "    ",
		ErrorPrefix:  "==> ",
		Ui:           c.Ui,
	}

	var logLevel string
	var nodeID string
	var serverID string
	var logJSON bool

	flags := c.Meta.FlagSet(c.Name(), FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.StringVar(&logLevel, "log-level", "", "")
	flags.StringVar(&nodeID, "node-id", "", "")
	flags.StringVar(&serverID, "server-id", "", "")
	flags.BoolVar(&logJSON, "json", false, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	args = flags.Args()
	if l := len(args); l != 0 {
		c.Ui.Error("This command takes no arguments")
		c.Ui.Error(commandErrorText(c))
		return 1
	}

	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		c.Ui.Error(commandErrorText(c))
		return 1
	}

	// Query the node info and lookup prefix
	if len(nodeID) == 1 {
		c.Ui.Error(fmt.Sprintf("Node identifier must contain at least two characters."))
		return 1
	}

	if nodeID != "" {
		nodeID = sanitizeUUIDPrefix(nodeID)
		nodes, _, err := client.Nodes().PrefixList(nodeID)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error querying node: %v", err))
			return 1
		}

		if len(nodes) > 1 {
			out := formatNodeStubList(nodes, false)
			c.Ui.Output(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", out))
			return 1
		}
		nodeID = nodes[0].ID
	}

	params := map[string]string{
		"log_level": logLevel,
		"node_id":   nodeID,
		"server_id": serverID,
		"log_json":  strconv.FormatBool(logJSON),
	}

	query := &api.QueryOptions{
		Params: params,
	}

	eventDoneCh := make(chan struct{})
	frames, errCh := client.Agent().Monitor(eventDoneCh, query)
	select {
	case err := <-errCh:
		c.Ui.Error(fmt.Sprintf("Error starting monitor: %s", err))
		c.Ui.Error(commandErrorText(c))
		return 1
	default:
	}

	// Create a reader
	var r io.ReadCloser
	frameReader := api.NewFrameReader(frames, errCh, eventDoneCh)
	frameReader.SetUnblockTime(500 * time.Millisecond)
	r = frameReader

	defer r.Close()

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-signalCh
		// End the streaming
		r.Close()
	}()

	_, err = io.Copy(os.Stdout, r)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("error monitoring logs: %s", err))
		return 1
	}

	return 0
}
<commit_msg>add server-id to -h output<commit_after>package command

import (
	"fmt"
	"io"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/nomad/api"
	"github.com/mitchellh/cli"
)

type MonitorCommand struct {
	Meta
}

func (c *MonitorCommand) Help() string {
	helpText := `
Usage: nomad monitor [options]

	Stream log messages of a nomad agent. The monitor command lets you
	listen for log levels that may be filtered out of the Nomad agent. For
	example your agent may only be logging at INFO level, but with the monitor
	command you can set -log-level DEBUG

General Options:

	` + generalOptionsUsage() + `

Monitor Specific Options:

  -log-level <level>
    Sets the log level to monitor (default: INFO)

  -node-id <node-id>
    Sets the specific node to monitor

  -server-id <server-id>
    Sets the specific server to monitor

  -json
    Sets log output to JSON format
  `
	return strings.TrimSpace(helpText)
}

func (c *MonitorCommand) Synopsis() string {
	return "stream logs from a Nomad agent"
}

func (c *MonitorCommand) Name() string { return "monitor" }

func (c *MonitorCommand) Run(args []string) int {
	c.Ui = &cli.PrefixedUi{
		OutputPrefix: "    ",
		InfoPrefix:   "    ",
		ErrorPrefix:  "==> ",
		Ui:           c.Ui,
	}

	var logLevel string
	var nodeID string
	var serverID string
	var logJSON bool

	flags := c.Meta.FlagSet(c.Name(), FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.StringVar(&logLevel, "log-level", "", "")
	flags.StringVar(&nodeID, "node-id", "", "")
	flags.StringVar(&serverID, "server-id", "", "")
	flags.BoolVar(&logJSON, "json", false, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	args = flags.Args()
	if l := len(args); l != 0 {
		c.Ui.Error("This command takes no arguments")
		c.Ui.Error(commandErrorText(c))
		return 1
	}

	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		c.Ui.Error(commandErrorText(c))
		return 1
	}

	// Query the node info and lookup prefix
	if len(nodeID) == 1 {
		c.Ui.Error(fmt.Sprintf("Node identifier must contain at least two characters."))
		return 1
	}

	if nodeID != "" {
		nodeID = sanitizeUUIDPrefix(nodeID)
		nodes, _, err := client.Nodes().PrefixList(nodeID)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error querying node: %v", err))
			return 1
		}

		if len(nodes) > 1 {
			out := formatNodeStubList(nodes, false)
			c.Ui.Output(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", out))
			return 1
		}
		nodeID = nodes[0].ID
	}

	params := map[string]string{
		"log_level": logLevel,
		"node_id":   nodeID,
		"server_id": serverID,
		"log_json":  strconv.FormatBool(logJSON),
	}

	query := &api.QueryOptions{
		Params: params,
	}

	eventDoneCh := make(chan struct{})
	frames, errCh := client.Agent().Monitor(eventDoneCh, query)
	select {
	case err := <-errCh:
		c.Ui.Error(fmt.Sprintf("Error starting monitor: %s", err))
		c.Ui.Error(commandErrorText(c))
		return 1
	default:
	}

	// Create a reader
	var r io.ReadCloser
	frameReader := api.NewFrameReader(frames, errCh, eventDoneCh)
	frameReader.SetUnblockTime(500 * time.Millisecond)
	r = frameReader

	defer r.Close()

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-signalCh
		// End the streaming
		r.Close()
	}()

	_, err = io.Copy(os.Stdout, r)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("error monitoring logs: %s", err))
		return 1
	}

	return 0
}
<|endoftext|>


<commit_before>package build

import (
	"fmt"
	"github.com/mitchellh/packer/packer"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"sync"
)

type Command byte

func (Command) Help() string {
	return strings.TrimSpace(helpText)
}

func (c Command) Run(env packer.Environment, args []string) int {
	if len(args) != 1 {
		env.Ui().Say(c.Help())
		return 1
	}

	// Read the file into a byte array so that we can parse the template
	log.Printf("Reading template: %s", args[0])
	tplData, err := ioutil.ReadFile(args[0])
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to read template file: %s", err))
		return 1
	}

	// Parse the template into a machine-usable format
	log.Println("Parsing template...")
	tpl, err := packer.ParseTemplate(tplData)
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to parse template: %s", err))
		return 1
	}

	// The component finder for our builds
	components := &packer.ComponentFinder{
		Builder:     env.Builder,
		Hook:        env.Hook,
		Provisioner: env.Provisioner,
	}

	// Go through each builder and compile the builds that we care about
	buildNames := tpl.BuildNames()
	builds := make([]packer.Build, 0, len(buildNames))
	for _, buildName := range buildNames {
		log.Printf("Creating build: %s", buildName)
		build, err := tpl.Build(buildName, components)
		if err != nil {
			env.Ui().Error(fmt.Sprintf("Failed to create build '%s': \n\n%s", buildName, err))
			return 1
		}

		builds = append(builds, build)
	}

	// Compile all the UIs for the builds
	colors := [5]packer.UiColor{
		packer.UiColorGreen,
		packer.UiColorYellow,
		packer.UiColorBlue,
		packer.UiColorMagenta,
		packer.UiColorCyan,
	}

	buildUis := make(map[string]packer.Ui)
	for i, b := range builds {
		var ui packer.Ui

		ui = &packer.ColoredUi{
			colors[i%len(colors)],
			env.Ui(),
		}

		ui = &packer.PrefixedUi{
			fmt.Sprintf("==> %s", b.Name()),
			fmt.Sprintf("    %s", b.Name()),
			ui,
		}

		buildUis[b.Name()] = ui
		ui.Say(fmt.Sprintf("%s output will be in this color.", b.Name()))
	}

	// Prepare all the builds
	for _, b := range builds {
		log.Printf("Preparing build: %s", b.Name())
		err := b.Prepare(buildUis[b.Name()])
		if err != nil {
			env.Ui().Error(err.Error())
			return 1
		}
	}

	// Run all the builds in parallel and wait for them to complete
	var wg sync.WaitGroup
	artifacts := make(map[string]packer.Artifact)
	for _, b := range builds {
		log.Printf("Starting build run: %s", b.Name())

		// Increment the waitgroup so we wait for this item to finish properly
		wg.Add(1)

		// Run the build in a goroutine
		go func() {
			defer wg.Done()
			artifacts[b.Name()] = b.Run(buildUis[b.Name()])
		}()
	}

	// Handle signals
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt)

	go func() {
		<-sigCh
		log.Println("Interrupted! Cancelling builds...")

		var wg sync.WaitGroup
		for _, b := range builds {
			wg.Add(1)

			go func() {
				defer wg.Done()

				log.Printf("Stopping build: %s", b.Name())
				b.Cancel()
			}()
		}

		wg.Wait()
	}()

	wg.Wait()

	// Output all the artifacts
	env.Ui().Say("\n==> The build completed! The artifacts created were:")
	for name, artifact := range artifacts {
		env.Ui().Say(fmt.Sprintf("--> %s:", name))
		env.Ui().Say(artifact.String())
	}

	return 0
}

func (Command) Synopsis() string {
	return "build image(s) from template"
}
<commit_msg>command/build: Cleanly exit after being interrupted<commit_after>package build

import (
	"fmt"
	"github.com/mitchellh/packer/packer"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"sync"
)

type Command byte

func (Command) Help() string {
	return strings.TrimSpace(helpText)
}

func (c Command) Run(env packer.Environment, args []string) int {
	if len(args) != 1 {
		env.Ui().Say(c.Help())
		return 1
	}

	// Read the file into a byte array so that we can parse the template
	log.Printf("Reading template: %s", args[0])
	tplData, err := ioutil.ReadFile(args[0])
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to read template file: %s", err))
		return 1
	}

	// Parse the template into a machine-usable format
	log.Println("Parsing template...")
	tpl, err := packer.ParseTemplate(tplData)
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to parse template: %s", err))
		return 1
	}

	// The component finder for our builds
	components := &packer.ComponentFinder{
		Builder:     env.Builder,
		Hook:        env.Hook,
		Provisioner: env.Provisioner,
	}

	// Go through each builder and compile the builds that we care about
	buildNames := tpl.BuildNames()
	builds := make([]packer.Build, 0, len(buildNames))
	for _, buildName := range buildNames {
		log.Printf("Creating build: %s", buildName)
		build, err := tpl.Build(buildName, components)
		if err != nil {
			env.Ui().Error(fmt.Sprintf("Failed to create build '%s': \n\n%s", buildName, err))
			return 1
		}

		builds = append(builds, build)
	}

	// Compile all the UIs for the builds
	colors := [5]packer.UiColor{
		packer.UiColorGreen,
		packer.UiColorYellow,
		packer.UiColorBlue,
		packer.UiColorMagenta,
		packer.UiColorCyan,
	}

	buildUis := make(map[string]packer.Ui)
	for i, b := range builds {
		var ui packer.Ui

		ui = &packer.ColoredUi{
			colors[i%len(colors)],
			env.Ui(),
		}

		ui = &packer.PrefixedUi{
			fmt.Sprintf("==> %s", b.Name()),
			fmt.Sprintf("    %s", b.Name()),
			ui,
		}

		buildUis[b.Name()] = ui
		ui.Say(fmt.Sprintf("%s output will be in this color.", b.Name()))
	}

	// Prepare all the builds
	for _, b := range builds {
		log.Printf("Preparing build: %s", b.Name())
		err := b.Prepare(buildUis[b.Name()])
		if err != nil {
			env.Ui().Error(err.Error())
			return 1
		}
	}

	// Run all the builds in parallel and wait for them to complete
	var wg sync.WaitGroup
	artifacts := make(map[string]packer.Artifact)
	for _, b := range builds {
		log.Printf("Starting build run: %s", b.Name())

		// Increment the waitgroup so we wait for this item to finish properly
		wg.Add(1)

		// Run the build in a goroutine
		go func() {
			defer wg.Done()
			artifacts[b.Name()] = b.Run(buildUis[b.Name()])
		}()
	}

	// Handle signals
	var interruptWg sync.WaitGroup
	interrupted := false
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt)

	go func() {
		<-sigCh
		interruptWg.Add(1)
		defer interruptWg.Done()
		interrupted = true

		log.Println("Interrupted! Cancelling builds...")

		var wg sync.WaitGroup
		for _, b := range builds {
			wg.Add(1)

			go func() {
				defer wg.Done()

				log.Printf("Stopping build: %s", b.Name())
				b.Cancel()
			}()
		}

		wg.Wait()
	}()

	// Wait for both the builds to complete and the interrupt handler,
	// if it is interrupted.
	wg.Wait()
	interruptWg.Wait()

	if interrupted {
		env.Ui().Say("Cleanly cancelled builds after being interrupted.")
		return 1
	}

	// Output all the artifacts
	env.Ui().Say("\n==> The build completed! The artifacts created were:")
	for name, artifact := range artifacts {
		env.Ui().Say(fmt.Sprintf("--> %s:", name))
		env.Ui().Say(artifact.String())
	}

	return 0
}

func (Command) Synopsis() string {
	return "build image(s) from template"
}
<|endoftext|>


<commit_before>package lib

import (
    "os"
    "fmt"
    "github.com/fsouza/go-dockerclient"
)

var (
    containerRe, _ = regexp.Compile(`[^/]+`)
)

// TODO (boldfield) :: This is kind of stupid... we shouldn't be loading all running containers
// TODO (boldfield) :: every time we check for the existance of a single container
func ContainerRunning(container string) (bool) {
    endpoint := "unix:///var/run/docker.sock"
    client, _ := docker.NewClient(endpoint)
    if containers, err := client.ListContainers(docker.ListContainersOptions{ All: false }); err != nil {
        fmt.Println(PrintRed("An error occoured while determining if docker container is running!"))
        fmt.Println(err)
        os.Exit(1)
    } else {
        for _, c := range containers {
            for _, n := range c.Names {
                clean = containerRe.FindString(n)
                if clean == container {
                    return true
                }
            }
        }
    }
    return false
}

// TODO (boldfield) :: Using this method as a data source for the above method would be nice...
func ListRunningContainers() ([]string) {
    runningContainers := make([]string, 0)
    endpoint := "unix:///var/run/docker.sock"
    client, _ := docker.NewClient(endpoint)
    if containers, err := client.ListContainers(docker.ListContainersOptions{ All: false }); err != nil {
        fmt.Println(PrintRed("An error occoured while determining if docker container is running!"))
        fmt.Println(err)
        os.Exit(1)
    } else {
        for _, c := range containers {
            for _, n := range c.Names {
                clean = containerRe.FindString(n)
                runningContainers = append(runningContainers, clean)
            }
        }
    }
    return runningContainers
}
<commit_msg>Adding missing import<commit_after>package lib

import (
    "os"
    "fmt"
    "regexp"
    "github.com/fsouza/go-dockerclient"
)

var (
    containerRe, _ = regexp.Compile(`[^/]+`)
)

// TODO (boldfield) :: This is kind of stupid... we shouldn't be loading all running containers
// TODO (boldfield) :: every time we check for the existance of a single container
func ContainerRunning(container string) (bool) {
    endpoint := "unix:///var/run/docker.sock"
    client, _ := docker.NewClient(endpoint)
    if containers, err := client.ListContainers(docker.ListContainersOptions{ All: false }); err != nil {
        fmt.Println(PrintRed("An error occoured while determining if docker container is running!"))
        fmt.Println(err)
        os.Exit(1)
    } else {
        for _, c := range containers {
            for _, n := range c.Names {
                clean = containerRe.FindString(n)
                if clean == container {
                    return true
                }
            }
        }
    }
    return false
}

// TODO (boldfield) :: Using this method as a data source for the above method would be nice...
func ListRunningContainers() ([]string) {
    runningContainers := make([]string, 0)
    endpoint := "unix:///var/run/docker.sock"
    client, _ := docker.NewClient(endpoint)
    if containers, err := client.ListContainers(docker.ListContainersOptions{ All: false }); err != nil {
        fmt.Println(PrintRed("An error occoured while determining if docker container is running!"))
        fmt.Println(err)
        os.Exit(1)
    } else {
        for _, c := range containers {
            for _, n := range c.Names {
                clean = containerRe.FindString(n)
                runningContainers = append(runningContainers, clean)
            }
        }
    }
    return runningContainers
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/github/hub/git"
	"github.com/github/hub/github"
	"github.com/github/hub/utils"
	"github.com/octokit/go-octokit/octokit"
)

var cmdPullRequest = &Command{
	Run: pullRequest,
	Usage: `
pull-request [-fo] [-b <BASE>] [-h <HEAD>] [-a <USER>] [-M <MILESTONE>] [-l <LABELS>]
pull-request -m <MESSAGE>
pull-request -F <FILE>
pull-request -i <ISSUE>
`,
	Long: `Create a GitHub pull request.

## Options:
	-f, --force
		Skip the check for unpushed commits.

	-m, --message <MESSAGE>
		Use the first line of <MESSAGE> as pull request title, and the rest as pull
		request description.

	-F, --file <FILE>
		Read the pull request title and description from <FILE>.

	-i, --issue <ISSUE>, <ISSUE-URL>
		(Deprecated) Convert <ISSUE> to a pull request.

	-o, --browse
		Open the new pull request in a web browser.

	-b, --base <BASE>
		The base branch in "[OWNER:]BRANCH" format. Defaults to the default branch
		(usually "master").

	-h, --head <HEAD>
		The base branch in "[OWNER:]BRANCH" format. Defaults to the current branch.

	-a, --assign <USER>
		Assign GitHub <USER> to this pull request.

	-M, --milestone <ID>
		Add this pull request to a GitHub milestone with id <ID>.

	-l, --labels <LABELS>
		Add a comma-separated list of labels to this pull request.

## See also:

hub(1), hub-merge(1), hub-checkout(1)
`,
}

var (
	flagPullRequestBase,
	flagPullRequestHead,
	flagPullRequestIssue,
	flagPullRequestMessage,
	flagPullRequestAssignee,
	flagPullRequestLabels,
	flagPullRequestFile string
	flagPullRequestBrowse,
	flagPullRequestForce bool
	flagPullRequestMilestone uint64
)

func init() {
	cmdPullRequest.Flag.StringVarP(&flagPullRequestBase, "base", "b", "", "BASE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestHead, "head", "h", "", "HEAD")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestIssue, "issue", "i", "", "ISSUE")
	cmdPullRequest.Flag.BoolVarP(&flagPullRequestBrowse, "browse", "o", false, "BROWSE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestMessage, "message", "m", "", "MESSAGE")
	cmdPullRequest.Flag.BoolVarP(&flagPullRequestForce, "force", "f", false, "FORCE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestFile, "file", "F", "", "FILE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestAssignee, "assign", "a", "", "USER")
	cmdPullRequest.Flag.Uint64VarP(&flagPullRequestMilestone, "milestone", "M", 0, "MILESTONE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestLabels, "labels", "l", "", "LABELS")

	CmdRunner.Use(cmdPullRequest)
}

func pullRequest(cmd *Command, args *Args) {
	localRepo, err := github.LocalRepo()
	utils.Check(err)

	currentBranch, err := localRepo.CurrentBranch()
	utils.Check(err)

	baseProject, err := localRepo.MainProject()
	utils.Check(err)

	host, err := github.CurrentConfig().PromptForHost(baseProject.Host)
	if err != nil {
		utils.Check(github.FormatError("creating pull request", err))
	}

	trackedBranch, headProject, err := localRepo.RemoteBranchAndProject(host.User, false)
	utils.Check(err)

	var (
		base, head string
		force      bool
	)

	force = flagPullRequestForce

	if flagPullRequestBase != "" {
		baseProject, base = parsePullRequestProject(baseProject, flagPullRequestBase)
	}

	if flagPullRequestHead != "" {
		headProject, head = parsePullRequestProject(headProject, flagPullRequestHead)
	}

	if args.ParamsSize() == 1 {
		arg := args.RemoveParam(0)
		flagPullRequestIssue = parsePullRequestIssueNumber(arg)
	}

	if base == "" {
		masterBranch := localRepo.MasterBranch()
		base = masterBranch.ShortName()
	}

	if head == "" && trackedBranch != nil {
		if !trackedBranch.IsRemote() {
			// the current branch tracking another branch
			// pretend there's no upstream at all
			trackedBranch = nil
		} else {
			if baseProject.SameAs(headProject) && base == trackedBranch.ShortName() {
				e := fmt.Errorf(`Aborted: head branch is the same as base ("%s")`, base)
				e = fmt.Errorf("%s\n(use `-h <branch>` to specify an explicit pull request head)", e)
				utils.Check(e)
			}
		}
	}

	if head == "" {
		if trackedBranch == nil {
			head = currentBranch.ShortName()
		} else {
			head = trackedBranch.ShortName()
		}
	}

	title, body, err := getTitleAndBodyFromFlags(flagPullRequestMessage, flagPullRequestFile)
	utils.Check(err)

	fullBase := fmt.Sprintf("%s:%s", baseProject.Owner, base)
	fullHead := fmt.Sprintf("%s:%s", headProject.Owner, head)

	if !force && trackedBranch != nil {
		remoteCommits, _ := git.RefList(trackedBranch.LongName(), "")
		if len(remoteCommits) > 0 {
			err = fmt.Errorf("Aborted: %d commits are not yet pushed to %s", len(remoteCommits), trackedBranch.LongName())
			err = fmt.Errorf("%s\n(use `-f` to force submit a pull request anyway)", err)
			utils.Check(err)
		}
	}

	var editor *github.Editor
	if title == "" && flagPullRequestIssue == "" {
		baseTracking := base
		headTracking := head

		remote := gitRemoteForProject(baseProject)
		if remote != nil {
			baseTracking = fmt.Sprintf("%s/%s", remote.Name, base)
		}
		if remote == nil || !baseProject.SameAs(headProject) {
			remote = gitRemoteForProject(headProject)
		}
		if remote != nil {
			headTracking = fmt.Sprintf("%s/%s", remote.Name, head)
		}

		message, err := pullRequestChangesMessage(baseTracking, headTracking, fullBase, fullHead)
		utils.Check(err)

		editor, err = github.NewEditor("PULLREQ", "pull request", message)
		utils.Check(err)

		title, body, err = editor.EditTitleAndBody()
		utils.Check(err)
	}

	if title == "" && flagPullRequestIssue == "" {
		utils.Check(fmt.Errorf("Aborting due to empty pull request title"))
	}

	var pullRequestURL string
	if args.Noop {
		args.Before(fmt.Sprintf("Would request a pull request to %s from %s", fullBase, fullHead), "")
		pullRequestURL = "PULL_REQUEST_URL"
	} else {
		var (
			pr  *octokit.PullRequest
			err error
		)

		client := github.NewClientWithHost(host)
		if title != "" {
			pr, err = client.CreatePullRequest(baseProject, base, fullHead, title, body)
		} else if flagPullRequestIssue != "" {
			pr, err = client.CreatePullRequestForIssue(baseProject, base, fullHead, flagPullRequestIssue)
		}

		if err == nil && editor != nil {
			defer editor.DeleteFile()
		}

		utils.Check(err)

		pullRequestURL = pr.HTMLURL

		if flagPullRequestAssignee != "" || flagPullRequestMilestone > 0 ||
			flagPullRequestLabels != "" {

			labels := []string{}
			for _, label := range strings.Split(flagPullRequestLabels, ",") {
				if label != "" {
					labels = append(labels, label)
				}
			}

			params := octokit.IssueParams{
				Assignee:  flagPullRequestAssignee,
				Milestone: flagPullRequestMilestone,
				Labels:    labels,
			}

			err = client.UpdateIssue(baseProject, pr.Number, params)
			utils.Check(err)
		}
	}

	if flagPullRequestBrowse {
		launcher, err := utils.BrowserLauncher()
		utils.Check(err)
		args.Replace(launcher[0], "", launcher[1:]...)
		args.AppendParams(pullRequestURL)
	} else {
		args.Replace("echo", "", pullRequestURL)
	}

	if flagPullRequestIssue != "" {
		args.After("echo", "Warning: Issue to pull request conversion is deprecated and might not work in the future.")
	}
}

func pullRequestChangesMessage(base, head, fullBase, fullHead string) (string, error) {
	var (
		defaultMsg string
		commitLogs string
		err        error
	)

	commits, _ := git.RefList(base, head)
	if len(commits) == 1 {
		defaultMsg, err = git.Show(commits[0])
		if err != nil {
			return "", err
		}
	} else if len(commits) > 1 {
		commitLogs, err = git.Log(base, head)
		if err != nil {
			return "", err
		}
	}

	cs := git.CommentChar()

	return renderPullRequestTpl(defaultMsg, cs, fullBase, fullHead, commitLogs)
}

func parsePullRequestProject(context *github.Project, s string) (p *github.Project, ref string) {
	p = context
	ref = s

	if strings.Contains(s, ":") {
		split := strings.SplitN(s, ":", 2)
		ref = split[1]
		var name string
		if !strings.Contains(split[0], "/") {
			name = context.Name
		}
		p = github.NewProject(split[0], name, context.Host)
	}

	return
}

func parsePullRequestIssueNumber(url string) string {
	u, e := github.ParseURL(url)
	if e != nil {
		return ""
	}

	r := regexp.MustCompile(`^issues\/(\d+)`)
	p := u.ProjectPath()
	if r.MatchString(p) {
		return r.FindStringSubmatch(p)[1]
	}

	return ""
}
<commit_msg>Fix copy-pasta<commit_after>package commands

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/github/hub/git"
	"github.com/github/hub/github"
	"github.com/github/hub/utils"
	"github.com/octokit/go-octokit/octokit"
)

var cmdPullRequest = &Command{
	Run: pullRequest,
	Usage: `
pull-request [-fo] [-b <BASE>] [-h <HEAD>] [-a <USER>] [-M <MILESTONE>] [-l <LABELS>]
pull-request -m <MESSAGE>
pull-request -F <FILE>
pull-request -i <ISSUE>
`,
	Long: `Create a GitHub pull request.

## Options:
	-f, --force
		Skip the check for unpushed commits.

	-m, --message <MESSAGE>
		Use the first line of <MESSAGE> as pull request title, and the rest as pull
		request description.

	-F, --file <FILE>
		Read the pull request title and description from <FILE>.

	-i, --issue <ISSUE>, <ISSUE-URL>
		(Deprecated) Convert <ISSUE> to a pull request.

	-o, --browse
		Open the new pull request in a web browser.

	-b, --base <BASE>
		The base branch in "[OWNER:]BRANCH" format. Defaults to the default branch
		(usually "master").

	-h, --head <HEAD>
		The head branch in "[OWNER:]BRANCH" format. Defaults to the current branch.

	-a, --assign <USER>
		Assign GitHub <USER> to this pull request.

	-M, --milestone <ID>
		Add this pull request to a GitHub milestone with id <ID>.

	-l, --labels <LABELS>
		Add a comma-separated list of labels to this pull request.

## See also:

hub(1), hub-merge(1), hub-checkout(1)
`,
}

var (
	flagPullRequestBase,
	flagPullRequestHead,
	flagPullRequestIssue,
	flagPullRequestMessage,
	flagPullRequestAssignee,
	flagPullRequestLabels,
	flagPullRequestFile string
	flagPullRequestBrowse,
	flagPullRequestForce bool
	flagPullRequestMilestone uint64
)

func init() {
	cmdPullRequest.Flag.StringVarP(&flagPullRequestBase, "base", "b", "", "BASE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestHead, "head", "h", "", "HEAD")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestIssue, "issue", "i", "", "ISSUE")
	cmdPullRequest.Flag.BoolVarP(&flagPullRequestBrowse, "browse", "o", false, "BROWSE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestMessage, "message", "m", "", "MESSAGE")
	cmdPullRequest.Flag.BoolVarP(&flagPullRequestForce, "force", "f", false, "FORCE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestFile, "file", "F", "", "FILE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestAssignee, "assign", "a", "", "USER")
	cmdPullRequest.Flag.Uint64VarP(&flagPullRequestMilestone, "milestone", "M", 0, "MILESTONE")
	cmdPullRequest.Flag.StringVarP(&flagPullRequestLabels, "labels", "l", "", "LABELS")

	CmdRunner.Use(cmdPullRequest)
}

func pullRequest(cmd *Command, args *Args) {
	localRepo, err := github.LocalRepo()
	utils.Check(err)

	currentBranch, err := localRepo.CurrentBranch()
	utils.Check(err)

	baseProject, err := localRepo.MainProject()
	utils.Check(err)

	host, err := github.CurrentConfig().PromptForHost(baseProject.Host)
	if err != nil {
		utils.Check(github.FormatError("creating pull request", err))
	}

	trackedBranch, headProject, err := localRepo.RemoteBranchAndProject(host.User, false)
	utils.Check(err)

	var (
		base, head string
		force      bool
	)

	force = flagPullRequestForce

	if flagPullRequestBase != "" {
		baseProject, base = parsePullRequestProject(baseProject, flagPullRequestBase)
	}

	if flagPullRequestHead != "" {
		headProject, head = parsePullRequestProject(headProject, flagPullRequestHead)
	}

	if args.ParamsSize() == 1 {
		arg := args.RemoveParam(0)
		flagPullRequestIssue = parsePullRequestIssueNumber(arg)
	}

	if base == "" {
		masterBranch := localRepo.MasterBranch()
		base = masterBranch.ShortName()
	}

	if head == "" && trackedBranch != nil {
		if !trackedBranch.IsRemote() {
			// the current branch tracking another branch
			// pretend there's no upstream at all
			trackedBranch = nil
		} else {
			if baseProject.SameAs(headProject) && base == trackedBranch.ShortName() {
				e := fmt.Errorf(`Aborted: head branch is the same as base ("%s")`, base)
				e = fmt.Errorf("%s\n(use `-h <branch>` to specify an explicit pull request head)", e)
				utils.Check(e)
			}
		}
	}

	if head == "" {
		if trackedBranch == nil {
			head = currentBranch.ShortName()
		} else {
			head = trackedBranch.ShortName()
		}
	}

	title, body, err := getTitleAndBodyFromFlags(flagPullRequestMessage, flagPullRequestFile)
	utils.Check(err)

	fullBase := fmt.Sprintf("%s:%s", baseProject.Owner, base)
	fullHead := fmt.Sprintf("%s:%s", headProject.Owner, head)

	if !force && trackedBranch != nil {
		remoteCommits, _ := git.RefList(trackedBranch.LongName(), "")
		if len(remoteCommits) > 0 {
			err = fmt.Errorf("Aborted: %d commits are not yet pushed to %s", len(remoteCommits), trackedBranch.LongName())
			err = fmt.Errorf("%s\n(use `-f` to force submit a pull request anyway)", err)
			utils.Check(err)
		}
	}

	var editor *github.Editor
	if title == "" && flagPullRequestIssue == "" {
		baseTracking := base
		headTracking := head

		remote := gitRemoteForProject(baseProject)
		if remote != nil {
			baseTracking = fmt.Sprintf("%s/%s", remote.Name, base)
		}
		if remote == nil || !baseProject.SameAs(headProject) {
			remote = gitRemoteForProject(headProject)
		}
		if remote != nil {
			headTracking = fmt.Sprintf("%s/%s", remote.Name, head)
		}

		message, err := pullRequestChangesMessage(baseTracking, headTracking, fullBase, fullHead)
		utils.Check(err)

		editor, err = github.NewEditor("PULLREQ", "pull request", message)
		utils.Check(err)

		title, body, err = editor.EditTitleAndBody()
		utils.Check(err)
	}

	if title == "" && flagPullRequestIssue == "" {
		utils.Check(fmt.Errorf("Aborting due to empty pull request title"))
	}

	var pullRequestURL string
	if args.Noop {
		args.Before(fmt.Sprintf("Would request a pull request to %s from %s", fullBase, fullHead), "")
		pullRequestURL = "PULL_REQUEST_URL"
	} else {
		var (
			pr  *octokit.PullRequest
			err error
		)

		client := github.NewClientWithHost(host)
		if title != "" {
			pr, err = client.CreatePullRequest(baseProject, base, fullHead, title, body)
		} else if flagPullRequestIssue != "" {
			pr, err = client.CreatePullRequestForIssue(baseProject, base, fullHead, flagPullRequestIssue)
		}

		if err == nil && editor != nil {
			defer editor.DeleteFile()
		}

		utils.Check(err)

		pullRequestURL = pr.HTMLURL

		if flagPullRequestAssignee != "" || flagPullRequestMilestone > 0 ||
			flagPullRequestLabels != "" {

			labels := []string{}
			for _, label := range strings.Split(flagPullRequestLabels, ",") {
				if label != "" {
					labels = append(labels, label)
				}
			}

			params := octokit.IssueParams{
				Assignee:  flagPullRequestAssignee,
				Milestone: flagPullRequestMilestone,
				Labels:    labels,
			}

			err = client.UpdateIssue(baseProject, pr.Number, params)
			utils.Check(err)
		}
	}

	if flagPullRequestBrowse {
		launcher, err := utils.BrowserLauncher()
		utils.Check(err)
		args.Replace(launcher[0], "", launcher[1:]...)
		args.AppendParams(pullRequestURL)
	} else {
		args.Replace("echo", "", pullRequestURL)
	}

	if flagPullRequestIssue != "" {
		args.After("echo", "Warning: Issue to pull request conversion is deprecated and might not work in the future.")
	}
}

func pullRequestChangesMessage(base, head, fullBase, fullHead string) (string, error) {
	var (
		defaultMsg string
		commitLogs string
		err        error
	)

	commits, _ := git.RefList(base, head)
	if len(commits) == 1 {
		defaultMsg, err = git.Show(commits[0])
		if err != nil {
			return "", err
		}
	} else if len(commits) > 1 {
		commitLogs, err = git.Log(base, head)
		if err != nil {
			return "", err
		}
	}

	cs := git.CommentChar()

	return renderPullRequestTpl(defaultMsg, cs, fullBase, fullHead, commitLogs)
}

func parsePullRequestProject(context *github.Project, s string) (p *github.Project, ref string) {
	p = context
	ref = s

	if strings.Contains(s, ":") {
		split := strings.SplitN(s, ":", 2)
		ref = split[1]
		var name string
		if !strings.Contains(split[0], "/") {
			name = context.Name
		}
		p = github.NewProject(split[0], name, context.Host)
	}

	return
}

func parsePullRequestIssueNumber(url string) string {
	u, e := github.ParseURL(url)
	if e != nil {
		return ""
	}

	r := regexp.MustCompile(`^issues\/(\d+)`)
	p := u.ProjectPath()
	if r.MatchString(p) {
		return r.FindStringSubmatch(p)[1]
	}

	return ""
}
<|endoftext|>


<commit_before>package manager

import (
	"context"
	"testing"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/services/encryption/ossencryption"
	"github.com/grafana/grafana/pkg/services/secrets"
	"github.com/grafana/grafana/pkg/services/secrets/database"
	"github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gopkg.in/ini.v1"
)

func TestSecretsService_EnvelopeEncryption(t *testing.T) {
	store := database.ProvideSecretsStore(sqlstore.InitTestDB(t))
	svc := SetupTestService(t, store)
	ctx := context.Background()

	t.Run("encrypting with no entity_id should create DEK", func(t *testing.T) {
		plaintext := []byte("very secret string")

		encrypted, err := svc.Encrypt(context.Background(), plaintext, secrets.WithoutScope())
		require.NoError(t, err)

		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, plaintext, decrypted)

		keys, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, len(keys), 1)
	})
	t.Run("encrypting another secret with no entity_id should use the same DEK", func(t *testing.T) {
		plaintext := []byte("another very secret string")

		encrypted, err := svc.Encrypt(context.Background(), plaintext, secrets.WithoutScope())
		require.NoError(t, err)

		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, plaintext, decrypted)

		keys, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, len(keys), 1)
	})
	t.Run("encrypting with entity_id provided should create a new DEK", func(t *testing.T) {
		plaintext := []byte("some test data")

		encrypted, err := svc.Encrypt(context.Background(), plaintext, secrets.WithScope("user:100"))
		require.NoError(t, err)

		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, plaintext, decrypted)

		keys, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, len(keys), 2)
	})

	t.Run("decrypting empty payload should return error", func(t *testing.T) {
		_, err := svc.Decrypt(context.Background(), []byte(""))
		require.Error(t, err)

		assert.Equal(t, "unable to decrypt empty payload", err.Error())
	})

	t.Run("decrypting legacy secret encrypted with secret key from settings", func(t *testing.T) {
		expected := "grafana"
		encrypted := []byte{122, 56, 53, 113, 101, 117, 73, 89, 20, 254, 36, 112, 112, 16, 128, 232, 227, 52, 166, 108, 192, 5, 28, 125, 126, 42, 197, 190, 251, 36, 94}
		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, expected, string(decrypted))
	})
}

func TestSecretsService_DataKeys(t *testing.T) {
	store := database.ProvideSecretsStore(sqlstore.InitTestDB(t))
	ctx := context.Background()

	dataKey := secrets.DataKey{
		Active:        true,
		Name:          "test1",
		Provider:      "test",
		EncryptedData: []byte{0x62, 0xAF, 0xA1, 0x1A},
	}

	t.Run("querying for a DEK that does not exist", func(t *testing.T) {
		res, err := store.GetDataKey(ctx, dataKey.Name)
		assert.ErrorIs(t, secrets.ErrDataKeyNotFound, err)
		assert.Nil(t, res)
	})

	t.Run("creating an active DEK", func(t *testing.T) {
		err := store.CreateDataKey(ctx, dataKey)
		require.NoError(t, err)

		res, err := store.GetDataKey(ctx, dataKey.Name)
		require.NoError(t, err)
		assert.Equal(t, dataKey.EncryptedData, res.EncryptedData)
		assert.Equal(t, dataKey.Provider, res.Provider)
		assert.Equal(t, dataKey.Name, res.Name)
		assert.True(t, dataKey.Active)
	})

	t.Run("creating an inactive DEK", func(t *testing.T) {
		k := secrets.DataKey{
			Active:        false,
			Name:          "test2",
			Provider:      "test",
			EncryptedData: []byte{0x62, 0xAF, 0xA1, 0x1A},
		}

		err := store.CreateDataKey(ctx, k)
		require.Error(t, err)

		res, err := store.GetDataKey(ctx, k.Name)
		assert.Equal(t, secrets.ErrDataKeyNotFound, err)
		assert.Nil(t, res)
	})

	t.Run("deleting DEK when no name provided must fail", func(t *testing.T) {
		beforeDelete, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		err = store.DeleteDataKey(ctx, "")
		require.Error(t, err)

		afterDelete, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, beforeDelete, afterDelete)
	})

	t.Run("deleting a DEK", func(t *testing.T) {
		err := store.DeleteDataKey(ctx, dataKey.Name)
		require.NoError(t, err)

		res, err := store.GetDataKey(ctx, dataKey.Name)
		assert.Equal(t, secrets.ErrDataKeyNotFound, err)
		assert.Nil(t, res)
	})
}

func TestSecretsService_GetCurrentProvider(t *testing.T) {
	t.Run("When encryption_provider is not specified explicitly, should use 'secretKey' as a current provider", func(t *testing.T) {
		cfg := `[security]
			secret_key = sdDkslslld`

		raw, err := ini.Load([]byte(cfg))
		require.NoError(t, err)
		settings := &setting.OSSImpl{Cfg: &setting.Cfg{Raw: raw}}

		svc := ProvideSecretsService(
			database.ProvideSecretsStore(sqlstore.InitTestDB(t)),
			bus.New(),
			ossencryption.ProvideService(),
			settings,
		)

		assert.Equal(t, "secretKey", svc.currentProvider)
	})

	t.Run("When encryption_provider value is set, should use it as a current provider", func(t *testing.T) {
		cfg := `[security]
			secret_key = sdDkslslld
			encryption_provider = awskms.second_key`

		raw, err := ini.Load([]byte(cfg))
		require.NoError(t, err)
		settings := &setting.OSSImpl{Cfg: &setting.Cfg{Raw: raw}}

		svc := ProvideSecretsService(
			database.ProvideSecretsStore(sqlstore.InitTestDB(t)),
			bus.New(),
			ossencryption.ProvideService(),
			settings,
		)

		assert.Equal(t, "awskms.second_key", svc.currentProvider)
	})
}
<commit_msg>Chore: Add current provider test for secrets service (#41387)<commit_after>package manager

import (
	"context"
	"testing"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/services/encryption/ossencryption"
	"github.com/grafana/grafana/pkg/services/secrets"
	"github.com/grafana/grafana/pkg/services/secrets/database"
	"github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gopkg.in/ini.v1"
)

func TestSecretsService_EnvelopeEncryption(t *testing.T) {
	store := database.ProvideSecretsStore(sqlstore.InitTestDB(t))
	svc := SetupTestService(t, store)
	ctx := context.Background()

	t.Run("encrypting with no entity_id should create DEK", func(t *testing.T) {
		plaintext := []byte("very secret string")

		encrypted, err := svc.Encrypt(context.Background(), plaintext, secrets.WithoutScope())
		require.NoError(t, err)

		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, plaintext, decrypted)

		keys, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, len(keys), 1)
	})
	t.Run("encrypting another secret with no entity_id should use the same DEK", func(t *testing.T) {
		plaintext := []byte("another very secret string")

		encrypted, err := svc.Encrypt(context.Background(), plaintext, secrets.WithoutScope())
		require.NoError(t, err)

		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, plaintext, decrypted)

		keys, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, len(keys), 1)
	})
	t.Run("encrypting with entity_id provided should create a new DEK", func(t *testing.T) {
		plaintext := []byte("some test data")

		encrypted, err := svc.Encrypt(context.Background(), plaintext, secrets.WithScope("user:100"))
		require.NoError(t, err)

		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, plaintext, decrypted)

		keys, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, len(keys), 2)
	})

	t.Run("decrypting empty payload should return error", func(t *testing.T) {
		_, err := svc.Decrypt(context.Background(), []byte(""))
		require.Error(t, err)

		assert.Equal(t, "unable to decrypt empty payload", err.Error())
	})

	t.Run("decrypting legacy secret encrypted with secret key from settings", func(t *testing.T) {
		expected := "grafana"
		encrypted := []byte{122, 56, 53, 113, 101, 117, 73, 89, 20, 254, 36, 112, 112, 16, 128, 232, 227, 52, 166, 108, 192, 5, 28, 125, 126, 42, 197, 190, 251, 36, 94}
		decrypted, err := svc.Decrypt(context.Background(), encrypted)
		require.NoError(t, err)
		assert.Equal(t, expected, string(decrypted))
	})
}

func TestSecretsService_DataKeys(t *testing.T) {
	store := database.ProvideSecretsStore(sqlstore.InitTestDB(t))
	ctx := context.Background()

	dataKey := secrets.DataKey{
		Active:        true,
		Name:          "test1",
		Provider:      "test",
		EncryptedData: []byte{0x62, 0xAF, 0xA1, 0x1A},
	}

	t.Run("querying for a DEK that does not exist", func(t *testing.T) {
		res, err := store.GetDataKey(ctx, dataKey.Name)
		assert.ErrorIs(t, secrets.ErrDataKeyNotFound, err)
		assert.Nil(t, res)
	})

	t.Run("creating an active DEK", func(t *testing.T) {
		err := store.CreateDataKey(ctx, dataKey)
		require.NoError(t, err)

		res, err := store.GetDataKey(ctx, dataKey.Name)
		require.NoError(t, err)
		assert.Equal(t, dataKey.EncryptedData, res.EncryptedData)
		assert.Equal(t, dataKey.Provider, res.Provider)
		assert.Equal(t, dataKey.Name, res.Name)
		assert.True(t, dataKey.Active)
	})

	t.Run("creating an inactive DEK", func(t *testing.T) {
		k := secrets.DataKey{
			Active:        false,
			Name:          "test2",
			Provider:      "test",
			EncryptedData: []byte{0x62, 0xAF, 0xA1, 0x1A},
		}

		err := store.CreateDataKey(ctx, k)
		require.Error(t, err)

		res, err := store.GetDataKey(ctx, k.Name)
		assert.Equal(t, secrets.ErrDataKeyNotFound, err)
		assert.Nil(t, res)
	})

	t.Run("deleting DEK when no name provided must fail", func(t *testing.T) {
		beforeDelete, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		err = store.DeleteDataKey(ctx, "")
		require.Error(t, err)

		afterDelete, err := store.GetAllDataKeys(ctx)
		require.NoError(t, err)
		assert.Equal(t, beforeDelete, afterDelete)
	})

	t.Run("deleting a DEK", func(t *testing.T) {
		err := store.DeleteDataKey(ctx, dataKey.Name)
		require.NoError(t, err)

		res, err := store.GetDataKey(ctx, dataKey.Name)
		assert.Equal(t, secrets.ErrDataKeyNotFound, err)
		assert.Nil(t, res)
	})
}

func TestSecretsService_UseCurrentProvider(t *testing.T) {
	t.Run("When encryption_provider is not specified explicitly, should use 'secretKey' as a current provider", func(t *testing.T) {
		svc := SetupTestService(t, database.ProvideSecretsStore(sqlstore.InitTestDB(t)))
		assert.Equal(t, "secretKey", svc.currentProvider)
	})

	t.Run("When encryption_provider value is set, should use it as a current provider", func(t *testing.T) {
		rawCfg := `[security]
			secret_key = sdDkslslld
			encryption_provider = awskms.second_key`

		raw, err := ini.Load([]byte(rawCfg))
		require.NoError(t, err)

		cfg := &setting.Cfg{Raw: raw, FeatureToggles: map[string]bool{envelopeEncryptionFeatureToggle: true}}
		settings := &setting.OSSImpl{Cfg: cfg}

		svc := ProvideSecretsService(
			database.ProvideSecretsStore(sqlstore.InitTestDB(t)),
			bus.New(),
			ossencryption.ProvideService(),
			settings,
		)

		assert.Equal(t, "awskms.second_key", svc.currentProvider)
	})

	t.Run("Should use encrypt/decrypt methods of the current provider", func(t *testing.T) {
		rawCfg := `
		[security]
		secret_key = sdDkslslld
		encryption_provider = fake-provider.some-key

		[security.encryption.fake-provider.some-key]
		`

		raw, err := ini.Load([]byte(rawCfg))
		require.NoError(t, err)

		cfg := &setting.Cfg{Raw: raw, FeatureToggles: map[string]bool{envelopeEncryptionFeatureToggle: true}}
		settings := &setting.OSSImpl{Cfg: cfg}

		secretStore := database.ProvideSecretsStore(sqlstore.InitTestDB(t))
		fake := fakeProvider{}
		providerID := "fake-provider.some-key"

		svcEncrypt := ProvideSecretsService(
			secretStore,
			bus.New(),
			ossencryption.ProvideService(),
			settings,
		)

		svcEncrypt.RegisterProvider(providerID, &fake)
		require.NoError(t, err)
		assert.Equal(t, providerID, svcEncrypt.CurrentProviderID())
		assert.Equal(t, 2, len(svcEncrypt.GetProviders()))
		encrypted, _ := svcEncrypt.Encrypt(context.Background(), []byte{}, secrets.WithoutScope())
		assert.True(t, fake.encryptCalled)

		// secret service tries to find a DEK in a cache first before calling provider's decrypt
		// to bypass the cache, we set up one more secrets service to test decrypting
		svcDecrypt := ProvideSecretsService(
			secretStore,
			bus.New(),
			ossencryption.ProvideService(),
			settings,
		)
		svcDecrypt.RegisterProvider(providerID, &fake)
		_, _ = svcDecrypt.Decrypt(context.Background(), encrypted)
		assert.True(t, fake.decryptCalled, "fake provider's decrypt should be called")
	})
}

type fakeProvider struct {
	encryptCalled bool
	decryptCalled bool
}

func (p *fakeProvider) Encrypt(_ context.Context, _ []byte) ([]byte, error) {
	p.encryptCalled = true
	return []byte{}, nil
}

func (p *fakeProvider) Decrypt(_ context.Context, _ []byte) ([]byte, error) {
	p.decryptCalled = true
	return []byte{}, nil
}
<|endoftext|>


<commit_before>package isolation

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"k8s.io/apimachinery/pkg/types"

	v1 "kubevirt.io/client-go/api/v1"
	cmdclient "kubevirt.io/kubevirt/pkg/virt-handler/cmd-client"
)

var (
	mountInfoTestData = map[string]string{
		"overlay":      "/proc/1/root/var/lib/docker/overlay2/f15d9ce07df72e80d809aa99ab4a171f2f3636f65f0653e75db8ca0befd8ae02/merged",
		"devicemapper": "/proc/1/root/var/lib/docker/devicemapper/mnt/d0990551ba8254871a449b2ff0d9063061ae96a2c195d7a850b62f030eae1710/rootfs",
		"btrfs":        "/proc/1/root/var/lib/containers/storage/btrfs/subvolumes/e9a94e2cde75c54834378d4835d4eda6bebb56b02068b9254780de6f9344ad0e",
	}
)

var _ = Describe("Isolation", func() {

	Context("With an existing socket", func() {

		var socket net.Listener
		var tmpDir string
		var podsDir string
		var podUID string
		var finished chan struct{} = nil

		podUID = "pid-uid-1234"
		vm := v1.NewMinimalVMIWithNS("default", "testvm")
		vm.UID = "1234"
		vm.Status = v1.VirtualMachineInstanceStatus{
			ActivePods: map[types.UID]string{
				types.UID(podUID): "myhost",
			},
		}
		vm.Status.NodeName = "myhost"

		BeforeEach(func() {
			var err error
			tmpDir, err = ioutil.TempDir("", "kubevirt")
			podsDir, err = ioutil.TempDir("", "pods")

			cmdclient.SetLegacyBaseDir(tmpDir)
			cmdclient.SetPodsBaseDir(tmpDir)

			os.MkdirAll(tmpDir+"/sockets/", os.ModePerm)
			socketFile := cmdclient.SocketFilePathOnHost(podUID)
			os.MkdirAll(filepath.Dir(socketFile), os.ModePerm)
			socket, err = net.Listen("unix", socketFile)
			Expect(err).ToNot(HaveOccurred())
			finished = make(chan struct{})
			go func() {
				for {
					conn, err := socket.Accept()
					if err != nil {
						close(finished)
						// closes when socket listener is closed
						return
					}
					conn.Close()
				}
			}()

		})

		It("Should detect the PID of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.Pid()).To(Equal(os.Getpid()))
		})

		It("Should not detect any slice if there is no matching controller", func() {
			_, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"not_existing_slice"}).Detect(vm)
			Expect(err).To(HaveOccurred())
		})

		It("Should detect the 'devices' controller slice of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.Slice()).To(HavePrefix("/"))
		})

		It("Should detect the PID namespace of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.PIDNamespace()).To(Equal(fmt.Sprintf("/proc/%d/ns/pid", os.Getpid())))
		})

		It("Should detect the Mount root of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.MountRoot()).To(Equal(fmt.Sprintf("/proc/%d/root", os.Getpid())))
		})

		It("Should detect the Network namespace of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.NetNamespace()).To(Equal(fmt.Sprintf("/proc/%d/ns/net", os.Getpid())))
		})

		It("Should detect the root mount info of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			mountInfo, err := result.MountInfoRoot()
			Expect(err).ToNot(HaveOccurred())
			Expect(mountInfo.MountPoint).To(Equal("/"))
		})

		It("Should detect the full path of the mount on a node", func() {
			// Restore the overwritten function
			defer func(f func(int) string) { mountInfoFunc = f }(mountInfoFunc)

			for testCase, want := range mountInfoTestData {
				mountInfoFunc = func(pid int) string {
					Expect(pid).To(SatisfyAny(Equal(1), Equal(os.Getpid())))
					base := filepath.Join("testdata", "mountinfo")
					if pid == 1 {
						return filepath.Join(base, fmt.Sprintf("%s_host", testCase))
					}
					return filepath.Join(base, fmt.Sprintf("%s_launcher", testCase))
				}

				mounted, err := NodeIsolationResult().IsMounted("/")
				Expect(err).ToNot(HaveOccurred())
				Expect(mounted).To(BeTrue())
				mounted, err = NodeIsolationResult().IsMounted("???")
				Expect(err).ToNot(HaveOccurred())
				Expect(mounted).To(BeFalse())

				result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
				Expect(err).ToNot(HaveOccurred())
				mountInfo, err := result.MountInfoRoot()
				Expect(err).ToNot(HaveOccurred())
				fullPath, err := NodeIsolationResult().FullPath(mountInfo)
				Expect(err).ToNot(HaveOccurred())
				Expect(fullPath).To(Equal(want))
			}
		})

		AfterEach(func() {
			socket.Close()
			os.RemoveAll(tmpDir)
			os.RemoveAll(podsDir)
			if finished != nil {
				<-finished
			}
		})
	})
})

var _ = Describe("getMemlockSize", func() {
	vm := v1.NewMinimalVMIWithNS("default", "testvm")

	It("Should return correct number of bytes for memlock limit", func() {
		bytes_, err := getMemlockSize(vm)
		Expect(err).ToNot(HaveOccurred())
		// 1Gb (static part for vfio VMs) + 256Mb (estimated overhead) + 8 Mb (VM)
		Expect(int(bytes_)).To(Equal(1264389000))
	})
})
<commit_msg>pkg/virt-handler/isolation/isolation_test: fix ineffectual assignment to err<commit_after>package isolation

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"k8s.io/apimachinery/pkg/types"

	v1 "kubevirt.io/client-go/api/v1"
	cmdclient "kubevirt.io/kubevirt/pkg/virt-handler/cmd-client"
)

var (
	mountInfoTestData = map[string]string{
		"overlay":      "/proc/1/root/var/lib/docker/overlay2/f15d9ce07df72e80d809aa99ab4a171f2f3636f65f0653e75db8ca0befd8ae02/merged",
		"devicemapper": "/proc/1/root/var/lib/docker/devicemapper/mnt/d0990551ba8254871a449b2ff0d9063061ae96a2c195d7a850b62f030eae1710/rootfs",
		"btrfs":        "/proc/1/root/var/lib/containers/storage/btrfs/subvolumes/e9a94e2cde75c54834378d4835d4eda6bebb56b02068b9254780de6f9344ad0e",
	}
)

var _ = Describe("Isolation", func() {

	Context("With an existing socket", func() {

		var socket net.Listener
		var tmpDir string
		var podsDir string
		var podUID string
		var finished chan struct{} = nil

		podUID = "pid-uid-1234"
		vm := v1.NewMinimalVMIWithNS("default", "testvm")
		vm.UID = "1234"
		vm.Status = v1.VirtualMachineInstanceStatus{
			ActivePods: map[types.UID]string{
				types.UID(podUID): "myhost",
			},
		}
		vm.Status.NodeName = "myhost"

		BeforeEach(func() {
			var err error
			tmpDir, err = ioutil.TempDir("", "kubevirt")
			Expect(err).ToNot(HaveOccurred())
			podsDir, err = ioutil.TempDir("", "pods")
			Expect(err).ToNot(HaveOccurred())

			cmdclient.SetLegacyBaseDir(tmpDir)
			cmdclient.SetPodsBaseDir(tmpDir)

			os.MkdirAll(tmpDir+"/sockets/", os.ModePerm)
			socketFile := cmdclient.SocketFilePathOnHost(podUID)
			os.MkdirAll(filepath.Dir(socketFile), os.ModePerm)
			socket, err = net.Listen("unix", socketFile)
			Expect(err).ToNot(HaveOccurred())
			finished = make(chan struct{})
			go func() {
				for {
					conn, err := socket.Accept()
					if err != nil {
						close(finished)
						// closes when socket listener is closed
						return
					}
					conn.Close()
				}
			}()

		})

		It("Should detect the PID of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.Pid()).To(Equal(os.Getpid()))
		})

		It("Should not detect any slice if there is no matching controller", func() {
			_, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"not_existing_slice"}).Detect(vm)
			Expect(err).To(HaveOccurred())
		})

		It("Should detect the 'devices' controller slice of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.Slice()).To(HavePrefix("/"))
		})

		It("Should detect the PID namespace of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.PIDNamespace()).To(Equal(fmt.Sprintf("/proc/%d/ns/pid", os.Getpid())))
		})

		It("Should detect the Mount root of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.MountRoot()).To(Equal(fmt.Sprintf("/proc/%d/root", os.Getpid())))
		})

		It("Should detect the Network namespace of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			Expect(result.NetNamespace()).To(Equal(fmt.Sprintf("/proc/%d/ns/net", os.Getpid())))
		})

		It("Should detect the root mount info of the test suite", func() {
			result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
			Expect(err).ToNot(HaveOccurred())
			mountInfo, err := result.MountInfoRoot()
			Expect(err).ToNot(HaveOccurred())
			Expect(mountInfo.MountPoint).To(Equal("/"))
		})

		It("Should detect the full path of the mount on a node", func() {
			// Restore the overwritten function
			defer func(f func(int) string) { mountInfoFunc = f }(mountInfoFunc)

			for testCase, want := range mountInfoTestData {
				mountInfoFunc = func(pid int) string {
					Expect(pid).To(SatisfyAny(Equal(1), Equal(os.Getpid())))
					base := filepath.Join("testdata", "mountinfo")
					if pid == 1 {
						return filepath.Join(base, fmt.Sprintf("%s_host", testCase))
					}
					return filepath.Join(base, fmt.Sprintf("%s_launcher", testCase))
				}

				mounted, err := NodeIsolationResult().IsMounted("/")
				Expect(err).ToNot(HaveOccurred())
				Expect(mounted).To(BeTrue())
				mounted, err = NodeIsolationResult().IsMounted("???")
				Expect(err).ToNot(HaveOccurred())
				Expect(mounted).To(BeFalse())

				result, err := NewSocketBasedIsolationDetector(tmpDir).Whitelist([]string{"devices"}).Detect(vm)
				Expect(err).ToNot(HaveOccurred())
				mountInfo, err := result.MountInfoRoot()
				Expect(err).ToNot(HaveOccurred())
				fullPath, err := NodeIsolationResult().FullPath(mountInfo)
				Expect(err).ToNot(HaveOccurred())
				Expect(fullPath).To(Equal(want))
			}
		})

		AfterEach(func() {
			socket.Close()
			os.RemoveAll(tmpDir)
			os.RemoveAll(podsDir)
			if finished != nil {
				<-finished
			}
		})
	})
})

var _ = Describe("getMemlockSize", func() {
	vm := v1.NewMinimalVMIWithNS("default", "testvm")

	It("Should return correct number of bytes for memlock limit", func() {
		bytes_, err := getMemlockSize(vm)
		Expect(err).ToNot(HaveOccurred())
		// 1Gb (static part for vfio VMs) + 256Mb (estimated overhead) + 8 Mb (VM)
		Expect(int(bytes_)).To(Equal(1264389000))
	})
})
<|endoftext|>


<commit_before>package lib

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/stats"
	"github.com/robertkrimen/otto"
	"gopkg.in/guregu/null.v3"
	"strconv"
	"sync"
	"time"
)

const (
	TickRate          = 1 * time.Millisecond
	ThresholdTickRate = 2 * time.Second
	ShutdownTimeout   = 10 * time.Second
)

var (
	MetricVUs    = &stats.Metric{Name: "vus", Type: stats.Gauge}
	MetricVUsMax = &stats.Metric{Name: "vus_max", Type: stats.Gauge}
	MetricErrors = &stats.Metric{Name: "errors", Type: stats.Counter}

	ErrTooManyVUs = errors.New("More VUs than the maximum requested")
	ErrMaxTooLow  = errors.New("Can't lower max below current VU count")

	// Special error used to taint a test, without printing an error.
	ErrVUWantsTaint = errors.New("[ErrVUWantsTaint is never logged]")
)

type vuEntry struct {
	VU     VU
	Buffer []stats.Sample
	Cancel context.CancelFunc
}

type Engine struct {
	Runner    Runner
	Status    Status
	Stages    []Stage
	Collector stats.Collector
	Pause     sync.WaitGroup
	Metrics   map[*stats.Metric]stats.Sink

	Thresholds  map[string][]*Threshold
	thresholdVM *otto.Otto

	ctx    context.Context
	vus    []*vuEntry
	nextID int64

	vuMutex   sync.Mutex
	waitGroup sync.WaitGroup
}

func NewEngine(r Runner) (*Engine, error) {
	e := &Engine{
		Runner: r,
		Status: Status{
			Running: null.BoolFrom(false),
			Tainted: null.BoolFrom(false),
			VUs:     null.IntFrom(0),
			VUsMax:  null.IntFrom(0),
			AtTime:  null.IntFrom(0),
		},
		Metrics:     make(map[*stats.Metric]stats.Sink),
		Thresholds:  make(map[string][]*Threshold),
		thresholdVM: otto.New(),
	}
	e.Pause.Add(1)

	return e, nil
}

func (e *Engine) Apply(opts Options) error {
	if opts.Run.Valid {
		e.SetRunning(opts.Run.Bool)
	}
	if opts.VUsMax.Valid {
		if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.VUs.Valid {
		if !opts.VUsMax.Valid {
			if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
				return err
			}
		}
		if err := e.SetVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.Duration.Valid {
		duration, err := time.ParseDuration(opts.Duration.String)
		if err != nil {
			return err
		}
		e.Stages = []Stage{Stage{Duration: null.IntFrom(int64(duration))}}
	}

	if opts.Quit.Valid {
		e.Status.Quit = opts.Quit
	}
	if opts.QuitOnTaint.Valid {
		e.Status.QuitOnTaint = opts.QuitOnTaint
	}

	if opts.Thresholds != nil {
		e.Thresholds = opts.Thresholds

		// Make sure all scripts are compiled!
		for m, scripts := range e.Thresholds {
			for i, script := range scripts {
				if script.Script != nil {
					continue
				}

				s, err := e.thresholdVM.Compile(fmt.Sprintf("threshold$%s:%i", m, i), script.Source)
				if err != nil {
					return err
				}
				script.Script = s
				scripts[i] = script
			}
		}
	}

	return nil
}

func (e *Engine) Run(ctx context.Context, opts Options) error {
	subctx, cancel := context.WithCancel(context.Background())
	e.ctx = subctx
	e.nextID = 1

	e.Apply(opts)

	if e.Collector != nil {
		e.waitGroup.Add(1)
		go func() {
			e.Collector.Run(subctx)
			log.Debug("Engine: Collector shut down")
			e.waitGroup.Done()
		}()
	} else {
		log.Debug("Engine: No Collector")
	}

	e.waitGroup.Add(1)
	go func() {
		e.runThresholds(subctx)
		log.Debug("Engine: Thresholds shut down")
		e.waitGroup.Done()
	}()

	e.consumeEngineStats()

	ticker := time.NewTicker(TickRate)
	lastTick := time.Now()

loop:
	for {
		select {
		case now := <-ticker.C:
			timeDelta := now.Sub(lastTick)
			e.Status.AtTime.Int64 += int64(timeDelta)
			lastTick = now

			stage, left, ok := StageAt(e.Stages, time.Duration(e.Status.AtTime.Int64))
			if stage.StartVUs.Valid && stage.EndVUs.Valid {
				progress := (float64(stage.Duration.Int64-int64(left)) / float64(stage.Duration.Int64))
				vus := Lerp(stage.StartVUs.Int64, stage.EndVUs.Int64, progress)
				e.SetVUs(vus)
			}

			for _, vu := range e.vus {
				buffer := vu.Buffer
				if buffer == nil {
					continue
				}
				vu.Buffer = nil
				e.consumeBuffer(buffer)
			}

			if !ok {
				e.SetRunning(false)

				if e.Status.Quit.Bool {
					break loop
				} else {
					log.Info("Test finished, press Ctrl+C to exit")
					<-ctx.Done()
					break loop
				}
			}

			if e.Status.QuitOnTaint.Bool && e.Status.Tainted.Bool {
				log.Warn("Test tainted, ending early...")
				break loop
			}

			e.consumeEngineStats()
		case <-ctx.Done():
			break loop
		}
	}

	e.SetRunning(false)
	e.SetVUs(0)
	e.SetMaxVUs(0)
	e.consumeEngineStats()

	e.ctx = nil
	cancel()

	log.Debug("Engine: Waiting for subsystem shutdown...")

	done := make(chan interface{})
	go func() {
		e.waitGroup.Wait()
		close(done)
	}()
	timeout := time.After(ShutdownTimeout)
	select {
	case <-done:
	case <-timeout:
		log.Warn("VUs took too long to finish, shutting down anyways")
	}

	return nil
}

func (e *Engine) IsRunning() bool {
	return e.ctx != nil
}

func (e *Engine) TotalTime() (total time.Duration, finite bool) {
	for _, stage := range e.Stages {
		if stage.Duration.Valid {
			total += time.Duration(stage.Duration.Int64)
			finite = true
		}
	}
	return total, finite
}

func (e *Engine) SetRunning(running bool) {
	if running && !e.Status.Running.Bool {
		e.Pause.Done()
		log.Debug("Engine Unpaused")
	} else if !running && e.Status.Running.Bool {
		e.Pause.Add(1)
		log.Debug("Engine Paused")
	}
	e.Status.Running.Bool = running
}

func (e *Engine) SetVUs(v int64) error {
	if e.Status.VUs.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUs.Int64, "to": v}).Debug("Setting VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v > e.Status.VUsMax.Int64 {
		return ErrTooManyVUs
	}

	current := e.Status.VUs.Int64
	for i := current; i < v; i++ {
		entry := e.vus[i]
		if entry.Cancel != nil {
			panic(errors.New("ATTEMPTED TO RESCHEDULE RUNNING VU"))
		}

		ctx, cancel := context.WithCancel(e.ctx)
		entry.Cancel = cancel

		if err := entry.VU.Reconfigure(e.nextID); err != nil {
			return err
		}
		e.waitGroup.Add(1)
		go func() {
			id := e.nextID
			e.runVU(ctx, id, entry)
			log.WithField("id", id).Debug("Engine: VU terminated")
			e.waitGroup.Done()
		}()
		e.nextID++
	}
	for i := current - 1; i >= v; i-- {
		entry := e.vus[i]
		entry.Cancel()
		entry.Cancel = nil
	}

	e.Status.VUs.Int64 = v
	return nil
}

func (e *Engine) SetMaxVUs(v int64) error {
	if e.Status.VUsMax.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUsMax.Int64, "to": v}).Debug("Setting Max VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v < e.Status.VUs.Int64 {
		return ErrMaxTooLow
	}

	current := e.Status.VUsMax.Int64
	if v > current {
		vus := e.vus
		for i := current; i < v; i++ {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			entry := &vuEntry{VU: vu}
			vus = append(vus, entry)
		}
		e.vus = vus
	} else if v < current {
		e.vus = e.vus[:v]
	}

	e.Status.VUsMax.Int64 = v
	return nil
}

func (e *Engine) runVU(ctx context.Context, id int64, vu *vuEntry) {
	idString := strconv.FormatInt(id, 10)

waitForPause:
	e.Pause.Wait()

	for {
		samples, err := vu.VU.RunOnce(ctx)

		// If the context is cancelled, the iteration is likely to produce erroneous output
		// due to cancelled HTTP requests and whatnot. Discard output from such runs.
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err != nil {
			e.Status.Tainted.Bool = true

			if err != ErrVUWantsTaint {
				if s, ok := err.(fmt.Stringer); ok {
					log.Error(s.String())
				} else {
					log.WithError(err).Error("Runtime Error")
				}

				samples = append(samples, stats.Sample{
					Metric: MetricErrors,
					Time:   time.Now(),
					Tags:   map[string]string{"vu": idString, "error": err.Error()},
					Value:  float64(1),
				})
			}
		}

		if samples != nil {
			buffer := vu.Buffer
			if buffer == nil {
				buffer = samples
			} else {
				buffer = append(buffer, samples...)
			}
			vu.Buffer = buffer
		}

		if !e.Status.Running.Bool {
			goto waitForPause
		}
	}
}

func (e *Engine) runThresholds(ctx context.Context) {
	ticker := time.NewTicker(ThresholdTickRate)
	for {
		select {
		case <-ticker.C:
			for m, sink := range e.Metrics {
				scripts, ok := e.Thresholds[m.Name]
				if !ok {
					continue
				}

				sample := sink.Format()
				for key, value := range sample {
					if m.Contains == stats.Time {
						value = value / float64(time.Millisecond)
					}
					// log.WithFields(log.Fields{"k": key, "v": value}).Debug("setting threshold data")
					e.thresholdVM.Set(key, value)
				}

				taint := false
				for _, script := range scripts {
					v, err := e.thresholdVM.Run(script.Script)
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold Error")
						taint = true
						continue
					}
					// log.WithFields(log.Fields{"metric": m.Name, "v": v, "s": sample}).Debug("threshold tick")
					bV, err := v.ToBoolean()
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold result is invalid")
						taint = true
						continue
					}
					if !bV {
						taint = true
						script.Failed = true
					}
				}

				for key, _ := range sample {
					e.thresholdVM.Set(key, otto.UndefinedValue())
				}

				if taint {
					m.Tainted = true
					e.Status.Tainted.Bool = true
				}
			}
		case <-ctx.Done():
			return
		}
	}
}

// Returns a value sink for a metric, created from the type if unavailable.
func (e *Engine) getSink(m *stats.Metric) stats.Sink {
	s, ok := e.Metrics[m]
	if !ok {
		switch m.Type {
		case stats.Counter:
			s = &stats.CounterSink{}
		case stats.Gauge:
			s = &stats.GaugeSink{}
		case stats.Trend:
			s = &stats.TrendSink{}
		case stats.Rate:
			s = &stats.RateSink{}
		}
		e.Metrics[m] = s
	}
	return s
}

func (e *Engine) consumeBuffer(buffer []stats.Sample) {
	for _, sample := range buffer {
		e.getSink(sample.Metric).Add(sample)
	}
	if e.Collector != nil {
		e.Collector.Collect(buffer)
	}
}

func (e *Engine) consumeEngineStats() {
	t := time.Now()
	e.consumeBuffer([]stats.Sample{
		stats.Sample{Metric: MetricVUs, Time: t, Value: float64(e.Status.VUs.Int64)},
		stats.Sample{Metric: MetricVUsMax, Time: t, Value: float64(e.Status.VUsMax.Int64)},
	})
}
<commit_msg>[docs] Comments in the Engine run loop<commit_after>package lib

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/stats"
	"github.com/robertkrimen/otto"
	"gopkg.in/guregu/null.v3"
	"strconv"
	"sync"
	"time"
)

const (
	TickRate          = 1 * time.Millisecond
	ThresholdTickRate = 2 * time.Second
	ShutdownTimeout   = 10 * time.Second
)

var (
	MetricVUs    = &stats.Metric{Name: "vus", Type: stats.Gauge}
	MetricVUsMax = &stats.Metric{Name: "vus_max", Type: stats.Gauge}
	MetricErrors = &stats.Metric{Name: "errors", Type: stats.Counter}

	ErrTooManyVUs = errors.New("More VUs than the maximum requested")
	ErrMaxTooLow  = errors.New("Can't lower max below current VU count")

	// Special error used to taint a test, without printing an error.
	ErrVUWantsTaint = errors.New("[ErrVUWantsTaint is never logged]")
)

type vuEntry struct {
	VU     VU
	Buffer []stats.Sample
	Cancel context.CancelFunc
}

type Engine struct {
	Runner    Runner
	Status    Status
	Stages    []Stage
	Collector stats.Collector
	Pause     sync.WaitGroup
	Metrics   map[*stats.Metric]stats.Sink

	Thresholds  map[string][]*Threshold
	thresholdVM *otto.Otto

	ctx    context.Context
	vus    []*vuEntry
	nextID int64

	vuMutex   sync.Mutex
	waitGroup sync.WaitGroup
}

func NewEngine(r Runner) (*Engine, error) {
	e := &Engine{
		Runner: r,
		Status: Status{
			Running: null.BoolFrom(false),
			Tainted: null.BoolFrom(false),
			VUs:     null.IntFrom(0),
			VUsMax:  null.IntFrom(0),
			AtTime:  null.IntFrom(0),
		},
		Metrics:     make(map[*stats.Metric]stats.Sink),
		Thresholds:  make(map[string][]*Threshold),
		thresholdVM: otto.New(),
	}
	e.Pause.Add(1)

	return e, nil
}

func (e *Engine) Apply(opts Options) error {
	if opts.Run.Valid {
		e.SetRunning(opts.Run.Bool)
	}
	if opts.VUsMax.Valid {
		if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.VUs.Valid {
		if !opts.VUsMax.Valid {
			if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
				return err
			}
		}
		if err := e.SetVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.Duration.Valid {
		duration, err := time.ParseDuration(opts.Duration.String)
		if err != nil {
			return err
		}
		e.Stages = []Stage{Stage{Duration: null.IntFrom(int64(duration))}}
	}

	if opts.Quit.Valid {
		e.Status.Quit = opts.Quit
	}
	if opts.QuitOnTaint.Valid {
		e.Status.QuitOnTaint = opts.QuitOnTaint
	}

	if opts.Thresholds != nil {
		e.Thresholds = opts.Thresholds

		// Make sure all scripts are compiled!
		for m, scripts := range e.Thresholds {
			for i, script := range scripts {
				if script.Script != nil {
					continue
				}

				s, err := e.thresholdVM.Compile(fmt.Sprintf("threshold$%s:%i", m, i), script.Source)
				if err != nil {
					return err
				}
				script.Script = s
				scripts[i] = script
			}
		}
	}

	return nil
}

func (e *Engine) Run(ctx context.Context, opts Options) error {
	subctx, cancel := context.WithCancel(context.Background())
	e.ctx = subctx
	e.nextID = 1

	e.Apply(opts)

	if e.Collector != nil {
		e.waitGroup.Add(1)
		go func() {
			e.Collector.Run(subctx)
			log.Debug("Engine: Collector shut down")
			e.waitGroup.Done()
		}()
	} else {
		log.Debug("Engine: No Collector")
	}

	e.waitGroup.Add(1)
	go func() {
		e.runThresholds(subctx)
		log.Debug("Engine: Thresholds shut down")
		e.waitGroup.Done()
	}()

	e.consumeEngineStats()

	ticker := time.NewTicker(TickRate)
	lastTick := time.Now()

loop:
	for {
		select {
		case now := <-ticker.C:
			// Track time deltas to ensure smooth interpolation even in the face of latency.
			timeDelta := now.Sub(lastTick)
			e.Status.AtTime.Int64 += int64(timeDelta)
			lastTick = now

			// Handle stages and VU interpolation.
			stage, left, ok := StageAt(e.Stages, time.Duration(e.Status.AtTime.Int64))
			if stage.StartVUs.Valid && stage.EndVUs.Valid {
				progress := (float64(stage.Duration.Int64-int64(left)) / float64(stage.Duration.Int64))
				vus := Lerp(stage.StartVUs.Int64, stage.EndVUs.Int64, progress)
				e.SetVUs(vus)
			}

			// Consume sample buffers. We use copies to avoid a race condition with runVU();
			// concurrent append() calls on the same list will result in a crash.
			for _, vu := range e.vus {
				buffer := vu.Buffer
				if buffer == nil {
					continue
				}
				vu.Buffer = nil
				e.consumeBuffer(buffer)
			}

			// If the test has ended, either pause or shut down.
			if !ok {
				e.SetRunning(false)

				if e.Status.Quit.Bool {
					break loop
				} else {
					log.Info("Test finished, press Ctrl+C to exit")
					<-ctx.Done()
					break loop
				}
			}

			// If the test is tainted, and we've requested --quit-on-taint, shut down.
			if e.Status.QuitOnTaint.Bool && e.Status.Tainted.Bool {
				log.Warn("Test tainted, ending early...")
				break loop
			}

			// Update internal metrics.
			e.consumeEngineStats()
		case <-ctx.Done():
			break loop
		}
	}

	e.SetRunning(false)
	e.SetVUs(0)
	e.SetMaxVUs(0)
	e.consumeEngineStats()

	e.ctx = nil
	cancel()

	log.Debug("Engine: Waiting for subsystem shutdown...")

	done := make(chan interface{})
	go func() {
		e.waitGroup.Wait()
		close(done)
	}()
	timeout := time.After(ShutdownTimeout)
	select {
	case <-done:
	case <-timeout:
		log.Warn("VUs took too long to finish, shutting down anyways")
	}

	return nil
}

func (e *Engine) IsRunning() bool {
	return e.ctx != nil
}

func (e *Engine) TotalTime() (total time.Duration, finite bool) {
	for _, stage := range e.Stages {
		if stage.Duration.Valid {
			total += time.Duration(stage.Duration.Int64)
			finite = true
		}
	}
	return total, finite
}

func (e *Engine) SetRunning(running bool) {
	if running && !e.Status.Running.Bool {
		e.Pause.Done()
		log.Debug("Engine Unpaused")
	} else if !running && e.Status.Running.Bool {
		e.Pause.Add(1)
		log.Debug("Engine Paused")
	}
	e.Status.Running.Bool = running
}

func (e *Engine) SetVUs(v int64) error {
	if e.Status.VUs.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUs.Int64, "to": v}).Debug("Setting VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v > e.Status.VUsMax.Int64 {
		return ErrTooManyVUs
	}

	current := e.Status.VUs.Int64
	for i := current; i < v; i++ {
		entry := e.vus[i]
		if entry.Cancel != nil {
			panic(errors.New("ATTEMPTED TO RESCHEDULE RUNNING VU"))
		}

		ctx, cancel := context.WithCancel(e.ctx)
		entry.Cancel = cancel

		if err := entry.VU.Reconfigure(e.nextID); err != nil {
			return err
		}
		e.waitGroup.Add(1)
		go func() {
			id := e.nextID
			e.runVU(ctx, id, entry)
			log.WithField("id", id).Debug("Engine: VU terminated")
			e.waitGroup.Done()
		}()
		e.nextID++
	}
	for i := current - 1; i >= v; i-- {
		entry := e.vus[i]
		entry.Cancel()
		entry.Cancel = nil
	}

	e.Status.VUs.Int64 = v
	return nil
}

func (e *Engine) SetMaxVUs(v int64) error {
	if e.Status.VUsMax.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUsMax.Int64, "to": v}).Debug("Setting Max VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v < e.Status.VUs.Int64 {
		return ErrMaxTooLow
	}

	current := e.Status.VUsMax.Int64
	if v > current {
		vus := e.vus
		for i := current; i < v; i++ {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			entry := &vuEntry{VU: vu}
			vus = append(vus, entry)
		}
		e.vus = vus
	} else if v < current {
		e.vus = e.vus[:v]
	}

	e.Status.VUsMax.Int64 = v
	return nil
}

func (e *Engine) runVU(ctx context.Context, id int64, vu *vuEntry) {
	idString := strconv.FormatInt(id, 10)

waitForPause:
	e.Pause.Wait()

	for {
		samples, err := vu.VU.RunOnce(ctx)

		// If the context is cancelled, the iteration is likely to produce erroneous output
		// due to cancelled HTTP requests and whatnot. Discard output from such runs.
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err != nil {
			e.Status.Tainted.Bool = true

			if err != ErrVUWantsTaint {
				if s, ok := err.(fmt.Stringer); ok {
					log.Error(s.String())
				} else {
					log.WithError(err).Error("Runtime Error")
				}

				samples = append(samples, stats.Sample{
					Metric: MetricErrors,
					Time:   time.Now(),
					Tags:   map[string]string{"vu": idString, "error": err.Error()},
					Value:  float64(1),
				})
			}
		}

		if samples != nil {
			buffer := vu.Buffer
			if buffer == nil {
				buffer = samples
			} else {
				buffer = append(buffer, samples...)
			}
			vu.Buffer = buffer
		}

		if !e.Status.Running.Bool {
			goto waitForPause
		}
	}
}

func (e *Engine) runThresholds(ctx context.Context) {
	ticker := time.NewTicker(ThresholdTickRate)
	for {
		select {
		case <-ticker.C:
			for m, sink := range e.Metrics {
				scripts, ok := e.Thresholds[m.Name]
				if !ok {
					continue
				}

				sample := sink.Format()
				for key, value := range sample {
					if m.Contains == stats.Time {
						value = value / float64(time.Millisecond)
					}
					// log.WithFields(log.Fields{"k": key, "v": value}).Debug("setting threshold data")
					e.thresholdVM.Set(key, value)
				}

				taint := false
				for _, script := range scripts {
					v, err := e.thresholdVM.Run(script.Script)
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold Error")
						taint = true
						continue
					}
					// log.WithFields(log.Fields{"metric": m.Name, "v": v, "s": sample}).Debug("threshold tick")
					bV, err := v.ToBoolean()
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold result is invalid")
						taint = true
						continue
					}
					if !bV {
						taint = true
						script.Failed = true
					}
				}

				for key, _ := range sample {
					e.thresholdVM.Set(key, otto.UndefinedValue())
				}

				if taint {
					m.Tainted = true
					e.Status.Tainted.Bool = true
				}
			}
		case <-ctx.Done():
			return
		}
	}
}

// Returns a value sink for a metric, created from the type if unavailable.
func (e *Engine) getSink(m *stats.Metric) stats.Sink {
	s, ok := e.Metrics[m]
	if !ok {
		switch m.Type {
		case stats.Counter:
			s = &stats.CounterSink{}
		case stats.Gauge:
			s = &stats.GaugeSink{}
		case stats.Trend:
			s = &stats.TrendSink{}
		case stats.Rate:
			s = &stats.RateSink{}
		}
		e.Metrics[m] = s
	}
	return s
}

func (e *Engine) consumeBuffer(buffer []stats.Sample) {
	for _, sample := range buffer {
		e.getSink(sample.Metric).Add(sample)
	}
	if e.Collector != nil {
		e.Collector.Collect(buffer)
	}
}

func (e *Engine) consumeEngineStats() {
	t := time.Now()
	e.consumeBuffer([]stats.Sample{
		stats.Sample{Metric: MetricVUs, Time: t, Value: float64(e.Status.VUs.Int64)},
		stats.Sample{Metric: MetricVUsMax, Time: t, Value: float64(e.Status.VUsMax.Int64)},
	})
}
<|endoftext|>