<commit_before>// Package storage implements services to persist data. The storage collection
// bundles storage instances to pass them around more easily.
package storage

import (
	"sync"

	"github.com/cenk/backoff"
	"github.com/the-anna-project/instrumentor"
	"github.com/the-anna-project/logger"

	"github.com/the-anna-project/storage/memory"
	"github.com/the-anna-project/storage/redis"
	"github.com/the-anna-project/storage/spec"
)

const (
	// KindMemory is the kind to be used to create a memory storage services.
	KindMemory = "memory"
	// KindRedis is the kind to be used to create a collection of redis storage
	// services.
	KindRedis = "redis"
)

// RedisConfig is the config applied to each redis instance. This is not
// relevant in case the memory kind is used.
type RedisConfig struct {
	Address string
	Prefix  string
}

// Redis is a config bundle of redis configs.
type Redis struct {
	Connection   RedisConfig
	Event        RedisConfig
	Feature      RedisConfig
	General      RedisConfig
	Index        RedisConfig
	Instrumentor RedisConfig
	Peer         RedisConfig
}

// CollectionConfig represents the configuration used to create a new storage
// collection.
type CollectionConfig struct {
	// Dependencies.
	BackoffFactory         func() spec.Backoff
	InstrumentorCollection *instrumentor.Collection
	LoggerService          logger.Service

	// Settings.
	Kind  string
	Redis *Redis
}

// DefaultCollectionConfig provides a default configuration to create a new
// storage collection by best effort.
func DefaultCollectionConfig() CollectionConfig {
	var err error

	var instrumentorCollection *instrumentor.Collection
	{
		instrumentorConfig := instrumentor.DefaultCollectionConfig()
		instrumentorCollection, err = instrumentor.NewCollection(instrumentorConfig)
		if err != nil {
			panic(err)
		}
	}

	var loggerService logger.Service
	{
		loggerConfig := logger.DefaultConfig()
		loggerService, err = logger.New(loggerConfig)
		if err != nil {
			panic(err)
		}
	}

	config := CollectionConfig{
		// Dependencies.
		BackoffFactory: func() spec.Backoff {
			return &backoff.StopBackOff{}
		},
		InstrumentorCollection: instrumentorCollection,
		LoggerService:          loggerService,

		// Settings.
		Kind:  KindMemory,
		Redis: nil,
	}

	return config
}

// NewCollection creates a new configured storage Collection.
func NewCollection(config CollectionConfig) (*Collection, error) {
	// Dependencies.
	if config.BackoffFactory == nil {
		return nil, maskAnyf(invalidConfigError, "backoff factory must not be empty")
	}
	if config.InstrumentorCollection == nil {
		return nil, maskAnyf(invalidConfigError, "instrumentor collection must not be empty")
	}
	if config.LoggerService == nil {
		return nil, maskAnyf(invalidConfigError, "logger service must not be empty")
	}

	// Settings.
	if config.Kind == "" {
		return nil, maskAnyf(invalidConfigError, "kind must not be empty")
	}
	if config.Kind != KindMemory && config.Kind != KindRedis {
		return nil, maskAnyf(invalidConfigError, "kind must be one of: %s, %s", KindMemory, KindRedis)
	}
	if config.Kind == KindRedis && config.Redis == nil {
		return nil, maskAnyf(invalidConfigError, "redis config must not be empty")
	}

	var err error

	var connectionService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			connectionConfig := memory.DefaultConfig()
			connectionService, err = memory.New(connectionConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			connectionConfig := redis.DefaultConfig()
			connectionConfig.Address = config.Redis.Connection.Address
			connectionConfig.BackoffFactory = config.BackoffFactory
			connectionConfig.InstrumentorCollection = config.InstrumentorCollection
			connectionConfig.LoggerService = config.LoggerService
			connectionConfig.Prefix = config.Redis.Connection.Prefix
			connectionService, err = redis.New(connectionConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var eventService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			eventConfig := memory.DefaultConfig()
			eventService, err = memory.New(eventConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			eventConfig := redis.DefaultConfig()
			eventConfig.Address = config.Redis.Event.Address
			eventConfig.BackoffFactory = config.BackoffFactory
			eventConfig.InstrumentorCollection = config.InstrumentorCollection
			eventConfig.LoggerService = config.LoggerService
			eventConfig.Prefix = config.Redis.Event.Prefix
			eventService, err = redis.New(eventConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var featureService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			featureConfig := memory.DefaultConfig()
			featureService, err = memory.New(featureConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			featureConfig := redis.DefaultConfig()
			featureConfig.Address = config.Redis.Feature.Address
			featureConfig.BackoffFactory = config.BackoffFactory
			featureConfig.InstrumentorCollection = config.InstrumentorCollection
			featureConfig.LoggerService = config.LoggerService
			featureConfig.Prefix = config.Redis.Feature.Prefix
			featureService, err = redis.New(featureConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var generalService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			generalConfig := memory.DefaultConfig()
			generalService, err = memory.New(generalConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			generalConfig := redis.DefaultConfig()
			generalConfig.Address = config.Redis.General.Address
			generalConfig.BackoffFactory = config.BackoffFactory
			generalConfig.InstrumentorCollection = config.InstrumentorCollection
			generalConfig.LoggerService = config.LoggerService
			generalConfig.Prefix = config.Redis.General.Prefix
			generalService, err = redis.New(generalConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var indexService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			indexConfig := memory.DefaultConfig()
			indexService, err = memory.New(indexConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			indexConfig := redis.DefaultConfig()
			indexConfig.Address = config.Redis.Index.Address
			indexConfig.BackoffFactory = config.BackoffFactory
			indexConfig.InstrumentorCollection = config.InstrumentorCollection
			indexConfig.LoggerService = config.LoggerService
			indexConfig.Prefix = config.Redis.Index.Prefix
			indexService, err = redis.New(indexConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var instrumentorService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			instrumentorConfig := memory.DefaultConfig()
			instrumentorService, err = memory.New(instrumentorConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			instrumentorConfig := redis.DefaultConfig()
			instrumentorConfig.Address = config.Redis.Instrumentor.Address
			instrumentorConfig.BackoffFactory = config.BackoffFactory
			instrumentorConfig.InstrumentorCollection = config.InstrumentorCollection
			instrumentorConfig.LoggerService = config.LoggerService
			instrumentorConfig.Prefix = config.Redis.Instrumentor.Prefix
			instrumentorService, err = redis.New(instrumentorConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var peerService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			peerConfig := memory.DefaultConfig()
			peerService, err = memory.New(peerConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			peerConfig := redis.DefaultConfig()
			peerConfig.Address = config.Redis.Peer.Address
			peerConfig.BackoffFactory = config.BackoffFactory
			peerConfig.InstrumentorCollection = config.InstrumentorCollection
			peerConfig.LoggerService = config.LoggerService
			peerConfig.Prefix = config.Redis.Peer.Prefix
			peerService, err = redis.New(peerConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	newCollection := &Collection{
		// Internals.
		bootOnce:     sync.Once{},
		shutdownOnce: sync.Once{},

		// Public.
		List: []spec.Service{
			connectionService,
			eventService,
			featureService,
			generalService,
			indexService,
			instrumentorService,
			peerService,
		},

		Connection:   connectionService,
		Event:        eventService,
		Feature:      featureService,
		General:      generalService,
		Index:        indexService,
		Instrumentor: instrumentorService,
		Peer:         peerService,
	}

	return newCollection, nil
}

// Collection is the object bundling all storages.
type Collection struct {
	// Internals.
	bootOnce     sync.Once
	shutdownOnce sync.Once

	// Public.
	List []spec.Service

	Connection   spec.Service
	Event        spec.Service
	Feature      spec.Service
	General      spec.Service
	Index        spec.Service
	Instrumentor spec.Service
	Peer         spec.Service
}

func (c *Collection) Boot() {
	c.bootOnce.Do(func() {
		var wg sync.WaitGroup

		for _, s := range c.List {
			wg.Add(1)
			go func() {
				s.Boot()
				wg.Done()
			}()
		}

		wg.Wait()
	})
}

func (c *Collection) Shutdown() {
	c.shutdownOnce.Do(func() {
		var wg sync.WaitGroup

		for _, s := range c.List {
			wg.Add(1)
			go func() {
				s.Shutdown()
				wg.Done()
			}()
		}

		wg.Wait()
	})
}
<commit_msg>added configuration storage to collection (#18)<commit_after>// Package storage implements services to persist data. The storage collection
// bundles storage instances to pass them around more easily.
package storage

import (
	"sync"

	"github.com/cenk/backoff"
	"github.com/the-anna-project/instrumentor"
	"github.com/the-anna-project/logger"

	"github.com/the-anna-project/storage/memory"
	"github.com/the-anna-project/storage/redis"
	"github.com/the-anna-project/storage/spec"
)

const (
	// KindMemory is the kind to be used to create a memory storage services.
	KindMemory = "memory"
	// KindRedis is the kind to be used to create a collection of redis storage
	// services.
	KindRedis = "redis"
)

// RedisConfig is the config applied to each redis instance. This is not
// relevant in case the memory kind is used.
type RedisConfig struct {
	Address string
	Prefix  string
}

// Redis is a config bundle of redis configs.
type Redis struct {
	Configuration RedisConfig
	Connection    RedisConfig
	Event         RedisConfig
	Feature       RedisConfig
	General       RedisConfig
	Index         RedisConfig
	Instrumentor  RedisConfig
	Peer          RedisConfig
}

// CollectionConfig represents the configuration used to create a new storage
// collection.
type CollectionConfig struct {
	// Dependencies.
	BackoffFactory         func() spec.Backoff
	InstrumentorCollection *instrumentor.Collection
	LoggerService          logger.Service

	// Settings.
	Kind  string
	Redis *Redis
}

// DefaultCollectionConfig provides a default configuration to create a new
// storage collection by best effort.
func DefaultCollectionConfig() CollectionConfig {
	var err error

	var instrumentorCollection *instrumentor.Collection
	{
		instrumentorConfig := instrumentor.DefaultCollectionConfig()
		instrumentorCollection, err = instrumentor.NewCollection(instrumentorConfig)
		if err != nil {
			panic(err)
		}
	}

	var loggerService logger.Service
	{
		loggerConfig := logger.DefaultConfig()
		loggerService, err = logger.New(loggerConfig)
		if err != nil {
			panic(err)
		}
	}

	config := CollectionConfig{
		// Dependencies.
		BackoffFactory: func() spec.Backoff {
			return &backoff.StopBackOff{}
		},
		InstrumentorCollection: instrumentorCollection,
		LoggerService:          loggerService,

		// Settings.
		Kind:  KindMemory,
		Redis: nil,
	}

	return config
}

// NewCollection creates a new configured storage Collection.
func NewCollection(config CollectionConfig) (*Collection, error) {
	// Dependencies.
	if config.BackoffFactory == nil {
		return nil, maskAnyf(invalidConfigError, "backoff factory must not be empty")
	}
	if config.InstrumentorCollection == nil {
		return nil, maskAnyf(invalidConfigError, "instrumentor collection must not be empty")
	}
	if config.LoggerService == nil {
		return nil, maskAnyf(invalidConfigError, "logger service must not be empty")
	}

	// Settings.
	if config.Kind == "" {
		return nil, maskAnyf(invalidConfigError, "kind must not be empty")
	}
	if config.Kind != KindMemory && config.Kind != KindRedis {
		return nil, maskAnyf(invalidConfigError, "kind must be one of: %s, %s", KindMemory, KindRedis)
	}
	if config.Kind == KindRedis && config.Redis == nil {
		return nil, maskAnyf(invalidConfigError, "redis config must not be empty")
	}

	var err error

	var configurationService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			configurationConfig := memory.DefaultConfig()
			configurationService, err = memory.New(configurationConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			configurationConfig := redis.DefaultConfig()
			configurationConfig.Address = config.Redis.Configuration.Address
			configurationConfig.BackoffFactory = config.BackoffFactory
			configurationConfig.InstrumentorCollection = config.InstrumentorCollection
			configurationConfig.LoggerService = config.LoggerService
			configurationConfig.Prefix = config.Redis.Configuration.Prefix
			configurationService, err = redis.New(configurationConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var connectionService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			connectionConfig := memory.DefaultConfig()
			connectionService, err = memory.New(connectionConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			connectionConfig := redis.DefaultConfig()
			connectionConfig.Address = config.Redis.Connection.Address
			connectionConfig.BackoffFactory = config.BackoffFactory
			connectionConfig.InstrumentorCollection = config.InstrumentorCollection
			connectionConfig.LoggerService = config.LoggerService
			connectionConfig.Prefix = config.Redis.Connection.Prefix
			connectionService, err = redis.New(connectionConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var eventService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			eventConfig := memory.DefaultConfig()
			eventService, err = memory.New(eventConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			eventConfig := redis.DefaultConfig()
			eventConfig.Address = config.Redis.Event.Address
			eventConfig.BackoffFactory = config.BackoffFactory
			eventConfig.InstrumentorCollection = config.InstrumentorCollection
			eventConfig.LoggerService = config.LoggerService
			eventConfig.Prefix = config.Redis.Event.Prefix
			eventService, err = redis.New(eventConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var featureService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			featureConfig := memory.DefaultConfig()
			featureService, err = memory.New(featureConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			featureConfig := redis.DefaultConfig()
			featureConfig.Address = config.Redis.Feature.Address
			featureConfig.BackoffFactory = config.BackoffFactory
			featureConfig.InstrumentorCollection = config.InstrumentorCollection
			featureConfig.LoggerService = config.LoggerService
			featureConfig.Prefix = config.Redis.Feature.Prefix
			featureService, err = redis.New(featureConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var generalService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			generalConfig := memory.DefaultConfig()
			generalService, err = memory.New(generalConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			generalConfig := redis.DefaultConfig()
			generalConfig.Address = config.Redis.General.Address
			generalConfig.BackoffFactory = config.BackoffFactory
			generalConfig.InstrumentorCollection = config.InstrumentorCollection
			generalConfig.LoggerService = config.LoggerService
			generalConfig.Prefix = config.Redis.General.Prefix
			generalService, err = redis.New(generalConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var indexService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			indexConfig := memory.DefaultConfig()
			indexService, err = memory.New(indexConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			indexConfig := redis.DefaultConfig()
			indexConfig.Address = config.Redis.Index.Address
			indexConfig.BackoffFactory = config.BackoffFactory
			indexConfig.InstrumentorCollection = config.InstrumentorCollection
			indexConfig.LoggerService = config.LoggerService
			indexConfig.Prefix = config.Redis.Index.Prefix
			indexService, err = redis.New(indexConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var instrumentorService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			instrumentorConfig := memory.DefaultConfig()
			instrumentorService, err = memory.New(instrumentorConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			instrumentorConfig := redis.DefaultConfig()
			instrumentorConfig.Address = config.Redis.Instrumentor.Address
			instrumentorConfig.BackoffFactory = config.BackoffFactory
			instrumentorConfig.InstrumentorCollection = config.InstrumentorCollection
			instrumentorConfig.LoggerService = config.LoggerService
			instrumentorConfig.Prefix = config.Redis.Instrumentor.Prefix
			instrumentorService, err = redis.New(instrumentorConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	var peerService spec.Service
	{
		switch config.Kind {
		case KindMemory:
			peerConfig := memory.DefaultConfig()
			peerService, err = memory.New(peerConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		case KindRedis:
			peerConfig := redis.DefaultConfig()
			peerConfig.Address = config.Redis.Peer.Address
			peerConfig.BackoffFactory = config.BackoffFactory
			peerConfig.InstrumentorCollection = config.InstrumentorCollection
			peerConfig.LoggerService = config.LoggerService
			peerConfig.Prefix = config.Redis.Peer.Prefix
			peerService, err = redis.New(peerConfig)
			if err != nil {
				return nil, maskAny(err)
			}
		}
	}

	newCollection := &Collection{
		// Internals.
		bootOnce:     sync.Once{},
		shutdownOnce: sync.Once{},

		// Public.
		List: []spec.Service{
			configurationService,
			connectionService,
			eventService,
			featureService,
			generalService,
			indexService,
			instrumentorService,
			peerService,
		},

		Configuration: configurationService,
		Connection:    connectionService,
		Event:         eventService,
		Feature:       featureService,
		General:       generalService,
		Index:         indexService,
		Instrumentor:  instrumentorService,
		Peer:          peerService,
	}

	return newCollection, nil
}

// Collection is the object bundling all storages.
type Collection struct {
	// Internals.
	bootOnce     sync.Once
	shutdownOnce sync.Once

	// Public.
	List []spec.Service

	Configuration spec.Service
	Connection    spec.Service
	Event         spec.Service
	Feature       spec.Service
	General       spec.Service
	Index         spec.Service
	Instrumentor  spec.Service
	Peer          spec.Service
}

func (c *Collection) Boot() {
	c.bootOnce.Do(func() {
		var wg sync.WaitGroup

		for _, s := range c.List {
			wg.Add(1)
			go func() {
				s.Boot()
				wg.Done()
			}()
		}

		wg.Wait()
	})
}

func (c *Collection) Shutdown() {
	c.shutdownOnce.Do(func() {
		var wg sync.WaitGroup

		for _, s := range c.List {
			wg.Add(1)
			go func() {
				s.Shutdown()
				wg.Done()
			}()
		}

		wg.Wait()
	})
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/buildkite/agent/api"
	"github.com/buildkite/agent/logger"
	"github.com/buildkite/agent/proctitle"
	"github.com/buildkite/agent/retry"
)

type AgentWorker struct {
	// The API Client used when this agent is communicating with the API
	APIClient *api.Client

	// The endpoint that should be used when communicating with the API
	Endpoint string

	// The registred agent API record
	Agent *api.Agent

	// The configuration of the agent from the CLI
	AgentConfiguration *AgentConfiguration

	// Whether or not the agent is running
	running bool

	// Used by the Start call to control the looping of the pings
	ticker *time.Ticker

	// Stop controls
	stop      chan struct{}
	stopping  bool
	stopMutex sync.Mutex

	// When this worker runs a job, we'll store an instance of the
	// JobRunner here
	jobRunner *JobRunner
}

// Creates the agent worker and initializes it's API Client
func (a AgentWorker) Create() AgentWorker {
	var endpoint string
	if a.Agent.Endpoint != "" {
		endpoint = a.Agent.Endpoint
	} else {
		endpoint = a.Endpoint
	}

	a.APIClient = APIClient{Endpoint: endpoint, Token: a.Agent.AccessToken}.Create()

	return a
}

// Starts the agent worker
func (a *AgentWorker) Start() error {
	// Mark the agent as running
	a.running = true

	// Create the intervals we'll be using
	pingInterval := time.Second * time.Duration(a.Agent.PingInterval)
	heartbeatInterval := time.Second * time.Duration(a.Agent.HearbeatInterval)

	// Setup and start the heartbeater
	go func() {
		// Keep the heartbeat running as long as the agent is
		for a.running {
			err := a.Heartbeat()
			if err != nil {
				logger.Error("Failed to heartbeat %s. Will try again in %s", err, heartbeatInterval)
			}

			time.Sleep(heartbeatInterval)
		}
	}()

	// Create the ticker and stop channels
	a.ticker = time.NewTicker(pingInterval)
	a.stop = make(chan struct{})

	// Continue this loop until the the ticker is stopped, and we received
	// a message on the stop channel.
	for {
		a.Ping()

		select {
		case <-a.ticker.C:
			continue
		case <-a.stop:
			a.ticker.Stop()
			return nil
		}
	}

	// Mark the agent as not running anymore
	a.running = false

	return nil
}

// Stops the agent from accepting new work and cancels any current work it's
// running
func (a *AgentWorker) Stop(graceful bool) {
	// Only allow one stop to run at a time (because we're playing with
	// channels)
	a.stopMutex.Lock()
	defer a.stopMutex.Unlock()

	if graceful {
		if a.stopping {
			logger.Warn("Agent is already gracefully stopping...")
		} else {
			// If we have a job, tell the user that we'll wait for
			// it to finish before disconnecting
			if a.jobRunner != nil {
				logger.Info("Gracefully stopping agent. Waiting for current job to finish before disconnecting...")
			} else {
				logger.Info("Gracefully stopping agent. Since there is no job running, the agent will disconnect immediately")
			}
		}
	} else {
		// If there's a job running, kill it, then disconnect
		if a.jobRunner != nil {
			logger.Info("Forcefully stopping agent. The current job will be canceled before disconnecting...")

			// Kill the current job. Doesn't do anything if the job
			// is already being killed, so it's safe to call
			// multiple times.
			a.jobRunner.Kill()
		} else {
			logger.Info("Forcefully stopping agent. Since there is no job running, the agent will disconnect immediately")
		}
	}

	// We don't need to do the below operations again since we've already
	// done them before
	if a.stopping {
		return
	}

	// Update the proc title
	a.UpdateProcTitle("stopping")

	// If we have a ticker, stop it, and send a signal to the stop channel,
	// which will cause the agent worker to stop looping immediatly.
	if a.ticker != nil {
		close(a.stop)
	}

	// Mark the agent as stopping
	a.stopping = true
}

// Connects the agent to the Buildkite Agent API, retrying up to 30 times if it
// fails.
func (a *AgentWorker) Connect() error {
	// Update the proc title
	a.UpdateProcTitle("connecting")

	return retry.Do(func(s *retry.Stats) error {
		_, err := a.APIClient.Agents.Connect()
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}

		return err
	}, &retry.Config{Maximum: 10, Interval: 5 * time.Second})
}

// Performs a heatbeat
func (a *AgentWorker) Heartbeat() error {
	var beat *api.Heartbeat
	var err error

	// Retry the heartbeat a few times
	err = retry.Do(func(s *retry.Stats) error {
		beat, _, err = a.APIClient.Heartbeats.Beat()
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}
		return err
	}, &retry.Config{Maximum: 5, Interval: 5 * time.Second})

	if err != nil {
		return err
	}

	logger.Debug("Heartbeat sent at %s and received at %s", beat.SentAt, beat.ReceivedAt)
	return nil
}

// Performs a ping, which returns what action the agent should take next.
func (a *AgentWorker) Ping() {
	// Update the proc title
	a.UpdateProcTitle("pinging")

	ping, _, err := a.APIClient.Pings.Get()
	if err != nil {
		// If a ping fails, we don't really care, because it'll
		// ping again after the interval.
		logger.Warn("Failed to ping: %s", err)
		return
	}

	// Should we switch endpoints?
	if ping.Endpoint != "" && ping.Endpoint != a.Agent.Endpoint {
		// Before switching to the new one, do a ping test to make sure it's
		// valid. If it is, switch and carry on, otherwise ignore the switch
		// for now.
		newAPIClient := APIClient{Endpoint: ping.Endpoint, Token: a.Agent.AccessToken}.Create()
		newPing, _, err := newAPIClient.Pings.Get()
		if err != nil {
			logger.Warn("Failed to ping the new endpoint %s - ignoring switch for now (%s)", ping.Endpoint, err)
		} else {
			// Replace the APIClient and process the new ping
			a.APIClient = newAPIClient
			a.Agent.Endpoint = ping.Endpoint
			ping = newPing
		}
	}

	// Is there a message that should be shown in the logs?
	if ping.Message != "" {
		logger.Info(ping.Message)
	}

	// Should the agent disconnect?
	if ping.Action == "disconnect" {
		a.Stop(false)
		return
	}

	// If we don't have a job, there's nothing to do!
	if ping.Job == nil {
		// Update the proc title
		a.UpdateProcTitle("idle")

		return
	}

	// Update the proc title
	a.UpdateProcTitle(fmt.Sprintf("job %s", strings.Split(ping.Job.ID, "-")[0]))

	logger.Info("Assigned job %s. Accepting...", ping.Job.ID)

	// Accept the job. We'll retry on connection related issues, but if
	// Buildkite returns a 422 or 500 for example, we'll just bail out,
	// re-ping, and try the whole process again.
	var accepted *api.Job
	retry.Do(func(s *retry.Stats) error {
		accepted, _, err = a.APIClient.Jobs.Accept(ping.Job)

		if err != nil {
			if api.IsRetryableError(err) {
				logger.Warn("%s (%s)", err, s)
			} else {
				logger.Warn("Buildkite rejected the call to accept the job (%s)", err)
				s.Break()
			}
		}

		return err
	}, &retry.Config{Maximum: 30, Interval: 1 * time.Second})

	// If `accepted` is nil, then the job was never accepted
	if accepted == nil {
		logger.Error("Failed to accept job")
		return
	}

	// Now that the job has been accepted, we can start it.
	a.jobRunner, err = JobRunner{
		Endpoint:           accepted.Endpoint,
		Agent:              a.Agent,
		AgentConfiguration: a.AgentConfiguration,
		Job:                accepted,
	}.Create()

	// Was there an error creating the job runner?
	if err != nil {
		logger.Error("Failed to initialize job: %s", err)
		return
	}

	// Start running the job
	if err = a.jobRunner.Run(); err != nil {
		logger.Error("Failed to run job: %s", err)
	}

	// No more job, no more runner.
	a.jobRunner = nil
}

// Disconnects the agent from the Buildkite Agent API, doesn't bother retrying
// because we want to disconnect as fast as possible.
func (a *AgentWorker) Disconnect() error {
	// Update the proc title
	a.UpdateProcTitle("disconnecting")

	_, err := a.APIClient.Agents.Disconnect()
	if err != nil {
		logger.Warn("There was an error sending the disconnect API call to Buildkite. If this agent still appears online, you may have to manually stop it (%s)", err)
	}

	return err
}

func (a *AgentWorker) UpdateProcTitle(action string) {
	proctitle.Replace(fmt.Sprintf("buildkite-agent v%s [%s]", Version(), action))
}
<commit_msg>Increase ping retry to every 5 seconds<commit_after>package agent

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/buildkite/agent/api"
	"github.com/buildkite/agent/logger"
	"github.com/buildkite/agent/proctitle"
	"github.com/buildkite/agent/retry"
)

type AgentWorker struct {
	// The API Client used when this agent is communicating with the API
	APIClient *api.Client

	// The endpoint that should be used when communicating with the API
	Endpoint string

	// The registred agent API record
	Agent *api.Agent

	// The configuration of the agent from the CLI
	AgentConfiguration *AgentConfiguration

	// Whether or not the agent is running
	running bool

	// Used by the Start call to control the looping of the pings
	ticker *time.Ticker

	// Stop controls
	stop      chan struct{}
	stopping  bool
	stopMutex sync.Mutex

	// When this worker runs a job, we'll store an instance of the
	// JobRunner here
	jobRunner *JobRunner
}

// Creates the agent worker and initializes it's API Client
func (a AgentWorker) Create() AgentWorker {
	var endpoint string
	if a.Agent.Endpoint != "" {
		endpoint = a.Agent.Endpoint
	} else {
		endpoint = a.Endpoint
	}

	a.APIClient = APIClient{Endpoint: endpoint, Token: a.Agent.AccessToken}.Create()

	return a
}

// Starts the agent worker
func (a *AgentWorker) Start() error {
	// Mark the agent as running
	a.running = true

	// Create the intervals we'll be using
	pingInterval := time.Second * time.Duration(a.Agent.PingInterval)
	heartbeatInterval := time.Second * time.Duration(a.Agent.HearbeatInterval)

	// Setup and start the heartbeater
	go func() {
		// Keep the heartbeat running as long as the agent is
		for a.running {
			err := a.Heartbeat()
			if err != nil {
				logger.Error("Failed to heartbeat %s. Will try again in %s", err, heartbeatInterval)
			}

			time.Sleep(heartbeatInterval)
		}
	}()

	// Create the ticker and stop channels
	a.ticker = time.NewTicker(pingInterval)
	a.stop = make(chan struct{})

	// Continue this loop until the the ticker is stopped, and we received
	// a message on the stop channel.
	for {
		a.Ping()

		select {
		case <-a.ticker.C:
			continue
		case <-a.stop:
			a.ticker.Stop()
			return nil
		}
	}

	// Mark the agent as not running anymore
	a.running = false

	return nil
}

// Stops the agent from accepting new work and cancels any current work it's
// running
func (a *AgentWorker) Stop(graceful bool) {
	// Only allow one stop to run at a time (because we're playing with
	// channels)
	a.stopMutex.Lock()
	defer a.stopMutex.Unlock()

	if graceful {
		if a.stopping {
			logger.Warn("Agent is already gracefully stopping...")
		} else {
			// If we have a job, tell the user that we'll wait for
			// it to finish before disconnecting
			if a.jobRunner != nil {
				logger.Info("Gracefully stopping agent. Waiting for current job to finish before disconnecting...")
			} else {
				logger.Info("Gracefully stopping agent. Since there is no job running, the agent will disconnect immediately")
			}
		}
	} else {
		// If there's a job running, kill it, then disconnect
		if a.jobRunner != nil {
			logger.Info("Forcefully stopping agent. The current job will be canceled before disconnecting...")

			// Kill the current job. Doesn't do anything if the job
			// is already being killed, so it's safe to call
			// multiple times.
			a.jobRunner.Kill()
		} else {
			logger.Info("Forcefully stopping agent. Since there is no job running, the agent will disconnect immediately")
		}
	}

	// We don't need to do the below operations again since we've already
	// done them before
	if a.stopping {
		return
	}

	// Update the proc title
	a.UpdateProcTitle("stopping")

	// If we have a ticker, stop it, and send a signal to the stop channel,
	// which will cause the agent worker to stop looping immediatly.
	if a.ticker != nil {
		close(a.stop)
	}

	// Mark the agent as stopping
	a.stopping = true
}

// Connects the agent to the Buildkite Agent API, retrying up to 30 times if it
// fails.
func (a *AgentWorker) Connect() error {
	// Update the proc title
	a.UpdateProcTitle("connecting")

	return retry.Do(func(s *retry.Stats) error {
		_, err := a.APIClient.Agents.Connect()
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}

		return err
	}, &retry.Config{Maximum: 10, Interval: 5 * time.Second})
}

// Performs a heatbeat
func (a *AgentWorker) Heartbeat() error {
	var beat *api.Heartbeat
	var err error

	// Retry the heartbeat a few times
	err = retry.Do(func(s *retry.Stats) error {
		beat, _, err = a.APIClient.Heartbeats.Beat()
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}
		return err
	}, &retry.Config{Maximum: 5, Interval: 5 * time.Second})

	if err != nil {
		return err
	}

	logger.Debug("Heartbeat sent at %s and received at %s", beat.SentAt, beat.ReceivedAt)
	return nil
}

// Performs a ping, which returns what action the agent should take next.
func (a *AgentWorker) Ping() {
	// Update the proc title
	a.UpdateProcTitle("pinging")

	ping, _, err := a.APIClient.Pings.Get()
	if err != nil {
		// If a ping fails, we don't really care, because it'll
		// ping again after the interval.
		logger.Warn("Failed to ping: %s", err)
		return
	}

	// Should we switch endpoints?
	if ping.Endpoint != "" && ping.Endpoint != a.Agent.Endpoint {
		// Before switching to the new one, do a ping test to make sure it's
		// valid. If it is, switch and carry on, otherwise ignore the switch
		// for now.
		newAPIClient := APIClient{Endpoint: ping.Endpoint, Token: a.Agent.AccessToken}.Create()
		newPing, _, err := newAPIClient.Pings.Get()
		if err != nil {
			logger.Warn("Failed to ping the new endpoint %s - ignoring switch for now (%s)", ping.Endpoint, err)
		} else {
			// Replace the APIClient and process the new ping
			a.APIClient = newAPIClient
			a.Agent.Endpoint = ping.Endpoint
			ping = newPing
		}
	}

	// Is there a message that should be shown in the logs?
	if ping.Message != "" {
		logger.Info(ping.Message)
	}

	// Should the agent disconnect?
	if ping.Action == "disconnect" {
		a.Stop(false)
		return
	}

	// If we don't have a job, there's nothing to do!
	if ping.Job == nil {
		// Update the proc title
		a.UpdateProcTitle("idle")

		return
	}

	// Update the proc title
	a.UpdateProcTitle(fmt.Sprintf("job %s", strings.Split(ping.Job.ID, "-")[0]))

	logger.Info("Assigned job %s. Accepting...", ping.Job.ID)

	// Accept the job. We'll retry on connection related issues, but if
	// Buildkite returns a 422 or 500 for example, we'll just bail out,
	// re-ping, and try the whole process again.
	var accepted *api.Job
	retry.Do(func(s *retry.Stats) error {
		accepted, _, err = a.APIClient.Jobs.Accept(ping.Job)

		if err != nil {
			if api.IsRetryableError(err) {
				logger.Warn("%s (%s)", err, s)
			} else {
				logger.Warn("Buildkite rejected the call to accept the job (%s)", err)
				s.Break()
			}
		}

		return err
	}, &retry.Config{Maximum: 30, Interval: 5 * time.Second})

	// If `accepted` is nil, then the job was never accepted
	if accepted == nil {
		logger.Error("Failed to accept job")
		return
	}

	// Now that the job has been accepted, we can start it.
	a.jobRunner, err = JobRunner{
		Endpoint:           accepted.Endpoint,
		Agent:              a.Agent,
		AgentConfiguration: a.AgentConfiguration,
		Job:                accepted,
	}.Create()

	// Was there an error creating the job runner?
	if err != nil {
		logger.Error("Failed to initialize job: %s", err)
		return
	}

	// Start running the job
	if err = a.jobRunner.Run(); err != nil {
		logger.Error("Failed to run job: %s", err)
	}

	// No more job, no more runner.
	a.jobRunner = nil
}

// Disconnects the agent from the Buildkite Agent API, doesn't bother retrying
// because we want to disconnect as fast as possible.
func (a *AgentWorker) Disconnect() error {
	// Update the proc title
	a.UpdateProcTitle("disconnecting")

	_, err := a.APIClient.Agents.Disconnect()
	if err != nil {
		logger.Warn("There was an error sending the disconnect API call to Buildkite. If this agent still appears online, you may have to manually stop it (%s)", err)
	}

	return err
}

func (a *AgentWorker) UpdateProcTitle(action string) {
	proctitle.Replace(fmt.Sprintf("buildkite-agent v%s [%s]", Version(), action))
}
<|endoftext|>


<commit_before>package agent_test

import (
	"bytes"
	"testing"

	"github.com/google/go-cmp/cmp/cmpopts"

	"github.com/google/go-cmp/cmp"

	"github.com/influxdata/kapacitor/udf/agent"
)

func TestMessage_ReadWrite(t *testing.T) {
	req := &agent.Request{}
	req.Message = &agent.Request_Keepalive{
		Keepalive: &agent.KeepaliveRequest{
			Time: 42,
		},
	}

	var buf bytes.Buffer

	err := agent.WriteMessage(req, &buf)
	if err != nil {
		t.Fatal(err)
	}

	nreq := &agent.Request{}
	var b []byte
	err = agent.ReadMessage(&b, &buf, nreq)
	if err != nil {
		t.Fatal(err)
	}

	if !cmp.Equal(req, nreq, cmpopts.IgnoreUnexported(agent.Request{}, agent.KeepaliveRequest{})) {
		t.Errorf("unexpected request: \n%s", cmp.Diff(nreq, req))
	}
}

func TestMessage_ReadWriteMultiple(t *testing.T) {
	req := &agent.Request{}
	req.Message = &agent.Request_Keepalive{
		Keepalive: &agent.KeepaliveRequest{
			Time: 42,
		},
	}

	var buf bytes.Buffer

	var count int = 1e4
	for i := 0; i < count; i++ {
		err := agent.WriteMessage(req, &buf)
		if err != nil {
			t.Fatal(err)
		}
	}

	nreq := &agent.Request{}
	var b []byte

	for i := 0; i < count; i++ {
		err := agent.ReadMessage(&b, &buf, nreq)
		if err != nil {
			t.Fatal(err)
		}

		if !cmp.Equal(req, nreq, cmpopts.IgnoreUnexported(agent.Request{}, agent.KeepaliveRequest{})) {
			t.Fatalf("unexpected request: i:%d \n%s", i, cmp.Diff(nreq, req))
		}
	}
}
<commit_msg>chore: fix imports formatting<commit_after>package agent_test

import (
	"bytes"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"github.com/influxdata/kapacitor/udf/agent"
)

func TestMessage_ReadWrite(t *testing.T) {
	req := &agent.Request{}
	req.Message = &agent.Request_Keepalive{
		Keepalive: &agent.KeepaliveRequest{
			Time: 42,
		},
	}

	var buf bytes.Buffer

	err := agent.WriteMessage(req, &buf)
	if err != nil {
		t.Fatal(err)
	}

	nreq := &agent.Request{}
	var b []byte
	err = agent.ReadMessage(&b, &buf, nreq)
	if err != nil {
		t.Fatal(err)
	}

	if !cmp.Equal(req, nreq, cmpopts.IgnoreUnexported(agent.Request{}, agent.KeepaliveRequest{})) {
		t.Errorf("unexpected request: \n%s", cmp.Diff(nreq, req))
	}
}

func TestMessage_ReadWriteMultiple(t *testing.T) {
	req := &agent.Request{}
	req.Message = &agent.Request_Keepalive{
		Keepalive: &agent.KeepaliveRequest{
			Time: 42,
		},
	}

	var buf bytes.Buffer

	var count int = 1e4
	for i := 0; i < count; i++ {
		err := agent.WriteMessage(req, &buf)
		if err != nil {
			t.Fatal(err)
		}
	}

	nreq := &agent.Request{}
	var b []byte

	for i := 0; i < count; i++ {
		err := agent.ReadMessage(&b, &buf, nreq)
		if err != nil {
			t.Fatal(err)
		}

		if !cmp.Equal(req, nreq, cmpopts.IgnoreUnexported(agent.Request{}, agent.KeepaliveRequest{})) {
			t.Fatalf("unexpected request: i:%d \n%s", i, cmp.Diff(nreq, req))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
// vespa test command
// Author: jonmv

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/spf13/cobra"
	"github.com/vespa-engine/vespa/client/go/util"
	"github.com/vespa-engine/vespa/client/go/vespa"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

func init() {
	rootCmd.AddCommand(testCmd)
}

// TODO: add link to test doc at cloud.vespa.ai
var testCmd = &cobra.Command{
	Use:   "test [tests directory or test file]",
	Short: "Run a test suite, or a single test",
	Long: `Run a test suite, or a single test

Runs all JSON test files in the specified directory, or the single JSON
test file specified.

If no directory or file is specified, the working directory is used instead.`,
	Example: `$ vespa test src/test/application/tests/system-test
$ vespa test src/test/application/tests/system-test/feed-and-query.json`,
	Args:              cobra.MaximumNArgs(1),
	DisableAutoGenTag: true,
	Run: func(cmd *cobra.Command, args []string) {
		target := getTarget()
		testPath := "."
		if len(args) > 0 {
			testPath = args[0]
		}
		if count, failed := runTests(testPath, target); len(failed) != 0 {
			fmt.Fprintf(stdout, "\nFailed %d of %d tests:\n", len(failed), count)
			for _, test := range failed {
				fmt.Fprintln(stdout, test)
			}
			exitFunc(3)
		} else if count == 0 {
			fmt.Fprintf(stdout, "Failed to find any tests at '%v'\n", testPath)
			exitFunc(3)
		} else {
			fmt.Fprintf(stdout, "%d tests completed successfully\n", count)
		}
	},
}

func runTests(rootPath string, target vespa.Target) (int, []string) {
	count := 0
	failed := make([]string, 0)
	if stat, err := os.Stat(rootPath); err != nil {
		fatalErr(err, "Failed reading specified test path")
	} else if stat.IsDir() {
		tests, err := os.ReadDir(rootPath)
		if err != nil {
			fatalErr(err, "Failed reading specified test directory")
		}
		for _, test := range tests {
			if !test.IsDir() && filepath.Ext(test.Name()) == ".json" {
				testPath := path.Join(rootPath, test.Name())
				failure := runTest(testPath, target)
				if failure != "" {
					failed = append(failed, failure)
				}
				count++
			}
		}
	} else if strings.HasSuffix(stat.Name(), ".json") {
		failure := runTest(rootPath, target)
		if failure != "" {
			failed = append(failed, failure)
		}
		count++
	}
	return count, failed
}

// Runs the test at the given path, and returns the specified test name if the test fails
func runTest(testPath string, target vespa.Target) string {
	var test test
	testBytes, err := ioutil.ReadFile(testPath)
	if err != nil {
		fatalErr(err, fmt.Sprintf("Failed to read test file at '%s'", testPath))
	}
	if err = json.Unmarshal(testBytes, &test); err != nil {
		fatalErr(err, fmt.Sprintf("Failed to parse test file at '%s", testPath))
	}

	testName := test.Name
	if test.Name == "" {
		testName = testPath
	}
	fmt.Fprintf(stdout, "Running %s:", testName)

	defaultParameters, err := getParameters(test.Defaults.ParametersRaw, path.Dir(testPath))
	if err != nil {
		fatalErr(err, fmt.Sprintf("Invalid default parameters for '%s'", testName))
	}

	if len(test.Assertions) == 0 {
		fatalErr(fmt.Errorf("a test must have at least one assertion, but none were found in '%s'", testPath))
	}
	for i, assertion := range test.Assertions {
		assertionName := assertion.Name
		if assertionName == "" {
			assertionName = fmt.Sprintf("assertion %d", i)
		}
		failure, err := verify(assertion, path.Dir(testPath), test.Defaults.Cluster, defaultParameters, target)
		if err != nil {
			fatalErr(err, fmt.Sprintf("\nError verifying %s", assertionName))
		}
		if failure != "" {
			fmt.Fprintf(stdout, "\nFailed verifying %s:\n%s\n", assertionName, failure)
			return fmt.Sprintf("%v: %v", testName, assertionName)
		}
		if i == 0 {
			fmt.Fprintf(stdout, " ")
		}
		fmt.Fprint(stdout, ".")
	}
	fmt.Fprintln(stdout, " OK!")
	return ""
}

// Asserts specified response is obtained for request, or returns a failure message, or an error if this fails
func verify(assertion assertion, testsPath string, defaultCluster string, defaultParameters map[string]string, target vespa.Target) (string, error) {
	requestBody, err := getBody(assertion.Request.BodyRaw, testsPath)
	if err != nil {
		return "", err
	}

	parameters, err := getParameters(assertion.Request.ParametersRaw, testsPath)
	if err != nil {
		return "", err
	}
	for name, value := range defaultParameters {
		if _, present := parameters[name]; !present {
			parameters[name] = value
		}
	}

	cluster := assertion.Request.Cluster
	if cluster == "" {
		cluster = defaultCluster
	}

	service, err := target.Service("query", 0, 0, cluster)
	if err != nil {
		return "", err
	}

	method := assertion.Request.Method
	if method == "" {
		method = "GET"
	}

	pathAndQuery := assertion.Request.URI
	if pathAndQuery == "" {
		pathAndQuery = "/search/"
	}
	requestUrl, err := url.ParseRequestURI(service.BaseURL + pathAndQuery)
	if err != nil {
		return "", err
	}
	query := requestUrl.Query()
	for name, value := range parameters {
		query.Add(name, value)
	}
	requestUrl.RawQuery = query.Encode()

	header := http.Header{}
	header.Add("Content-Type", "application/json") // TODO: Not guaranteed to be true ...

	request := &http.Request{
		URL:    requestUrl,
		Method: method,
		Header: header,
		Body:   ioutil.NopCloser(bytes.NewReader(requestBody)),
	}
	defer request.Body.Close()

	response, err := service.Do(request, 600*time.Second) // Vespa should provide a response within the given request timeout
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	statusCode := assertion.Response.Code
	if statusCode == 0 {
		statusCode = 200
	}
	if statusCode != response.StatusCode {
		return fmt.Sprintf("Expected status code (%d) does not match actual (%d). Response body:\n%s", statusCode, response.StatusCode, util.ReaderToJSON(response.Body)), nil
	}

	responseBodySpecBytes, err := getBody(assertion.Response.BodyRaw, testsPath)
	if err != nil {
		return "", err
	}
	if responseBodySpecBytes == nil {
		return "", nil
	}
	var responseBodySpec interface{}
	err = json.Unmarshal(responseBodySpecBytes, &responseBodySpec)
	if err != nil {
		return "", err
	}

	responseBodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	var responseBody interface{}
	err = json.Unmarshal(responseBodyBytes, &responseBody)
	if err != nil {
		return "", fmt.Errorf("got non-JSON response; %w:\n%s", err, string(responseBodyBytes))
	}

	failure, err := compare(responseBodySpec, responseBody, "")
	if failure != "" {
		responsePretty, _ := json.MarshalIndent(responseBody, "", "  ")
		failure = failure + " Response body:\n" + string(responsePretty)
	}
	return failure, err
}

func compare(expected interface{}, actual interface{}, path string) (string, error) {
	typeMatch := false
	valueMatch := false
	switch u := expected.(type) {
	case nil:
		typeMatch = actual == nil
		valueMatch = actual == nil
	case bool:
		v, ok := actual.(bool)
		typeMatch = ok
		valueMatch = ok && u == v
	case float64:
		v, ok := actual.(float64)
		typeMatch = ok
		valueMatch = ok && math.Abs(u-v) < 1e-9
	case string:
		v, ok := actual.(string)
		typeMatch = ok
		valueMatch = ok && (u == v)
	case []interface{}:
		v, ok := actual.([]interface{})
		typeMatch = ok
		if ok {
			if len(u) == len(v) {
				for i, e := range u {
					result, err := compare(e, v[i], fmt.Sprintf("%s/%d", path, i))
					if result != "" || err != nil {
						return result, err
					}
				}
				valueMatch = true
			} else {
				return fmt.Sprintf("Expected number of elements at %s (%d) does not match actual (%d).", path, len(u), len(v)), nil
			}
		}
	case map[string]interface{}:
		v, ok := actual.(map[string]interface{})
		typeMatch = ok
		if ok {
			for n, e := range u {
				childPath := fmt.Sprintf("%s/%s", path, strings.ReplaceAll(strings.ReplaceAll(n, "~", "~0"), "/", "~1"))
				f, ok := v[n]
				if !ok {
					return fmt.Sprintf("Expected field at %s not present in actual data.", childPath), nil
				}
				result, err := compare(e, f, childPath)
				if result != "" || err != nil {
					return result, err
				}
			}
			valueMatch = true
		}
	default:
		return "", fmt.Errorf("unexpected expected JSON type for value '%v'", expected)
	}

	if !(typeMatch && valueMatch) {
		if path == "" {
			path = "root"
		}
		expectedJson, _ := json.MarshalIndent(expected, "", "  ")
		actualJson, _ := json.MarshalIndent(actual, "", "  ")
		return fmt.Sprintf("Expected JSON at %s (%s) does not match actual (%s).", path, expectedJson, actualJson), nil
	}
	return "", nil
}

func getParameters(parametersRaw []byte, testsPath string) (map[string]string, error) {
	if parametersRaw != nil {
		var parametersPath string
		if err := json.Unmarshal(parametersRaw, &parametersPath); err == nil {
			resolvedParametersPath := path.Join(testsPath, parametersPath)
			parametersRaw, err = ioutil.ReadFile(resolvedParametersPath)
			if err != nil {
				fatalErr(err, fmt.Sprintf("Failed to read request parameters file at '%s'", resolvedParametersPath))
			}
		}
		var parameters map[string]string
		if err := json.Unmarshal(parametersRaw, &parameters); err != nil {
			return nil, fmt.Errorf("request parameters must be JSON with only string values: %w", err)
		}
		return parameters, nil
	}
	return make(map[string]string), nil
}

func getBody(bodyRaw []byte, testsPath string) ([]byte, error) {
	var bodyPath string
	if err := json.Unmarshal(bodyRaw, &bodyPath); err == nil {
		resolvedBodyPath := path.Join(testsPath, bodyPath)
		bodyRaw, err = ioutil.ReadFile(resolvedBodyPath)
		if err != nil {
			fatalErr(err, fmt.Sprintf("Failed to read body file at '%s'", resolvedBodyPath))
		}
	}
	return bodyRaw, nil
}

type test struct {
	Name       string      `json:"name"`
	Defaults   defaults    `json:"defaults"`
	Assertions []assertion `json:"assertions"`
}

type defaults struct {
	Cluster       string          `json:"cluster"`
	ParametersRaw json.RawMessage `json:"parameters"`
}

type assertion struct {
	Name     string   `json:"name"`
	Request  request  `json:"request"`
	Response response `json:"response"`
}

type request struct {
	Cluster       string          `json:"cluster"`
	Method        string          `json:"method"`
	URI           string          `json:"uri"`
	ParametersRaw json.RawMessage `json:"parameters"`
	BodyRaw       json.RawMessage `json:"body"`
}

type response struct {
	Code    int             `json:"code"`
	BodyRaw json.RawMessage `json:"body"`
}
<commit_msg>use ioutil.ReadDir, which exists in go 1.15<commit_after>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
// vespa test command
// Author: jonmv

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/spf13/cobra"
	"github.com/vespa-engine/vespa/client/go/util"
	"github.com/vespa-engine/vespa/client/go/vespa"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

func init() {
	rootCmd.AddCommand(testCmd)
}

// TODO: add link to test doc at cloud.vespa.ai
var testCmd = &cobra.Command{
	Use:   "test [tests directory or test file]",
	Short: "Run a test suite, or a single test",
	Long: `Run a test suite, or a single test

Runs all JSON test files in the specified directory, or the single JSON
test file specified.

If no directory or file is specified, the working directory is used instead.`,
	Example: `$ vespa test src/test/application/tests/system-test
$ vespa test src/test/application/tests/system-test/feed-and-query.json`,
	Args:              cobra.MaximumNArgs(1),
	DisableAutoGenTag: true,
	Run: func(cmd *cobra.Command, args []string) {
		target := getTarget()
		testPath := "."
		if len(args) > 0 {
			testPath = args[0]
		}
		if count, failed := runTests(testPath, target); len(failed) != 0 {
			fmt.Fprintf(stdout, "\nFailed %d of %d tests:\n", len(failed), count)
			for _, test := range failed {
				fmt.Fprintln(stdout, test)
			}
			exitFunc(3)
		} else if count == 0 {
			fmt.Fprintf(stdout, "Failed to find any tests at '%v'\n", testPath)
			exitFunc(3)
		} else {
			fmt.Fprintf(stdout, "%d tests completed successfully\n", count)
		}
	},
}

func runTests(rootPath string, target vespa.Target) (int, []string) {
	count := 0
	failed := make([]string, 0)
	if stat, err := os.Stat(rootPath); err != nil {
		fatalErr(err, "Failed reading specified test path")
	} else if stat.IsDir() {
		tests, err := ioutil.ReadDir(rootPath) // TODO: Use os.ReadDir when >= 1.16 is required.
		if err != nil {
			fatalErr(err, "Failed reading specified test directory")
		}
		for _, test := range tests {
			if !test.IsDir() && filepath.Ext(test.Name()) == ".json" {
				testPath := path.Join(rootPath, test.Name())
				failure := runTest(testPath, target)
				if failure != "" {
					failed = append(failed, failure)
				}
				count++
			}
		}
	} else if strings.HasSuffix(stat.Name(), ".json") {
		failure := runTest(rootPath, target)
		if failure != "" {
			failed = append(failed, failure)
		}
		count++
	}
	return count, failed
}

// Runs the test at the given path, and returns the specified test name if the test fails
func runTest(testPath string, target vespa.Target) string {
	var test test
	testBytes, err := ioutil.ReadFile(testPath)
	if err != nil {
		fatalErr(err, fmt.Sprintf("Failed to read test file at '%s'", testPath))
	}
	if err = json.Unmarshal(testBytes, &test); err != nil {
		fatalErr(err, fmt.Sprintf("Failed to parse test file at '%s", testPath))
	}

	testName := test.Name
	if test.Name == "" {
		testName = testPath
	}
	fmt.Fprintf(stdout, "Running %s:", testName)

	defaultParameters, err := getParameters(test.Defaults.ParametersRaw, path.Dir(testPath))
	if err != nil {
		fatalErr(err, fmt.Sprintf("Invalid default parameters for '%s'", testName))
	}

	if len(test.Assertions) == 0 {
		fatalErr(fmt.Errorf("a test must have at least one assertion, but none were found in '%s'", testPath))
	}
	for i, assertion := range test.Assertions {
		assertionName := assertion.Name
		if assertionName == "" {
			assertionName = fmt.Sprintf("assertion %d", i)
		}
		failure, err := verify(assertion, path.Dir(testPath), test.Defaults.Cluster, defaultParameters, target)
		if err != nil {
			fatalErr(err, fmt.Sprintf("\nError verifying %s", assertionName))
		}
		if failure != "" {
			fmt.Fprintf(stdout, "\nFailed verifying %s:\n%s\n", assertionName, failure)
			return fmt.Sprintf("%v: %v", testName, assertionName)
		}
		if i == 0 {
			fmt.Fprintf(stdout, " ")
		}
		fmt.Fprint(stdout, ".")
	}
	fmt.Fprintln(stdout, " OK!")
	return ""
}

// Asserts specified response is obtained for request, or returns a failure message, or an error if this fails
func verify(assertion assertion, testsPath string, defaultCluster string, defaultParameters map[string]string, target vespa.Target) (string, error) {
	requestBody, err := getBody(assertion.Request.BodyRaw, testsPath)
	if err != nil {
		return "", err
	}

	parameters, err := getParameters(assertion.Request.ParametersRaw, testsPath)
	if err != nil {
		return "", err
	}
	for name, value := range defaultParameters {
		if _, present := parameters[name]; !present {
			parameters[name] = value
		}
	}

	cluster := assertion.Request.Cluster
	if cluster == "" {
		cluster = defaultCluster
	}

	service, err := target.Service("query", 0, 0, cluster)
	if err != nil {
		return "", err
	}

	method := assertion.Request.Method
	if method == "" {
		method = "GET"
	}

	pathAndQuery := assertion.Request.URI
	if pathAndQuery == "" {
		pathAndQuery = "/search/"
	}
	requestUrl, err := url.ParseRequestURI(service.BaseURL + pathAndQuery)
	if err != nil {
		return "", err
	}
	query := requestUrl.Query()
	for name, value := range parameters {
		query.Add(name, value)
	}
	requestUrl.RawQuery = query.Encode()

	header := http.Header{}
	header.Add("Content-Type", "application/json") // TODO: Not guaranteed to be true ...

	request := &http.Request{
		URL:    requestUrl,
		Method: method,
		Header: header,
		Body:   ioutil.NopCloser(bytes.NewReader(requestBody)),
	}
	defer request.Body.Close()

	response, err := service.Do(request, 600*time.Second) // Vespa should provide a response within the given request timeout
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	statusCode := assertion.Response.Code
	if statusCode == 0 {
		statusCode = 200
	}
	if statusCode != response.StatusCode {
		return fmt.Sprintf("Expected status code (%d) does not match actual (%d). Response body:\n%s", statusCode, response.StatusCode, util.ReaderToJSON(response.Body)), nil
	}

	responseBodySpecBytes, err := getBody(assertion.Response.BodyRaw, testsPath)
	if err != nil {
		return "", err
	}
	if responseBodySpecBytes == nil {
		return "", nil
	}
	var responseBodySpec interface{}
	err = json.Unmarshal(responseBodySpecBytes, &responseBodySpec)
	if err != nil {
		return "", err
	}

	responseBodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	var responseBody interface{}
	err = json.Unmarshal(responseBodyBytes, &responseBody)
	if err != nil {
		return "", fmt.Errorf("got non-JSON response; %w:\n%s", err, string(responseBodyBytes))
	}

	failure, err := compare(responseBodySpec, responseBody, "")
	if failure != "" {
		responsePretty, _ := json.MarshalIndent(responseBody, "", "  ")
		failure = failure + " Response body:\n" + string(responsePretty)
	}
	return failure, err
}

func compare(expected interface{}, actual interface{}, path string) (string, error) {
	typeMatch := false
	valueMatch := false
	switch u := expected.(type) {
	case nil:
		typeMatch = actual == nil
		valueMatch = actual == nil
	case bool:
		v, ok := actual.(bool)
		typeMatch = ok
		valueMatch = ok && u == v
	case float64:
		v, ok := actual.(float64)
		typeMatch = ok
		valueMatch = ok && math.Abs(u-v) < 1e-9
	case string:
		v, ok := actual.(string)
		typeMatch = ok
		valueMatch = ok && (u == v)
	case []interface{}:
		v, ok := actual.([]interface{})
		typeMatch = ok
		if ok {
			if len(u) == len(v) {
				for i, e := range u {
					result, err := compare(e, v[i], fmt.Sprintf("%s/%d", path, i))
					if result != "" || err != nil {
						return result, err
					}
				}
				valueMatch = true
			} else {
				return fmt.Sprintf("Expected number of elements at %s (%d) does not match actual (%d).", path, len(u), len(v)), nil
			}
		}
	case map[string]interface{}:
		v, ok := actual.(map[string]interface{})
		typeMatch = ok
		if ok {
			for n, e := range u {
				childPath := fmt.Sprintf("%s/%s", path, strings.ReplaceAll(strings.ReplaceAll(n, "~", "~0"), "/", "~1"))
				f, ok := v[n]
				if !ok {
					return fmt.Sprintf("Expected field at %s not present in actual data.", childPath), nil
				}
				result, err := compare(e, f, childPath)
				if result != "" || err != nil {
					return result, err
				}
			}
			valueMatch = true
		}
	default:
		return "", fmt.Errorf("unexpected expected JSON type for value '%v'", expected)
	}

	if !(typeMatch && valueMatch) {
		if path == "" {
			path = "root"
		}
		expectedJson, _ := json.MarshalIndent(expected, "", "  ")
		actualJson, _ := json.MarshalIndent(actual, "", "  ")
		return fmt.Sprintf("Expected JSON at %s (%s) does not match actual (%s).", path, expectedJson, actualJson), nil
	}
	return "", nil
}

func getParameters(parametersRaw []byte, testsPath string) (map[string]string, error) {
	if parametersRaw != nil {
		var parametersPath string
		if err := json.Unmarshal(parametersRaw, &parametersPath); err == nil {
			resolvedParametersPath := path.Join(testsPath, parametersPath)
			parametersRaw, err = ioutil.ReadFile(resolvedParametersPath)
			if err != nil {
				fatalErr(err, fmt.Sprintf("Failed to read request parameters file at '%s'", resolvedParametersPath))
			}
		}
		var parameters map[string]string
		if err := json.Unmarshal(parametersRaw, &parameters); err != nil {
			return nil, fmt.Errorf("request parameters must be JSON with only string values: %w", err)
		}
		return parameters, nil
	}
	return make(map[string]string), nil
}

func getBody(bodyRaw []byte, testsPath string) ([]byte, error) {
	var bodyPath string
	if err := json.Unmarshal(bodyRaw, &bodyPath); err == nil {
		resolvedBodyPath := path.Join(testsPath, bodyPath)
		bodyRaw, err = ioutil.ReadFile(resolvedBodyPath)
		if err != nil {
			fatalErr(err, fmt.Sprintf("Failed to read body file at '%s'", resolvedBodyPath))
		}
	}
	return bodyRaw, nil
}

type test struct {
	Name       string      `json:"name"`
	Defaults   defaults    `json:"defaults"`
	Assertions []assertion `json:"assertions"`
}

type defaults struct {
	Cluster       string          `json:"cluster"`
	ParametersRaw json.RawMessage `json:"parameters"`
}

type assertion struct {
	Name     string   `json:"name"`
	Request  request  `json:"request"`
	Response response `json:"response"`
}

type request struct {
	Cluster       string          `json:"cluster"`
	Method        string          `json:"method"`
	URI           string          `json:"uri"`
	ParametersRaw json.RawMessage `json:"parameters"`
	BodyRaw       json.RawMessage `json:"body"`
}

type response struct {
	Code    int             `json:"code"`
	BodyRaw json.RawMessage `json:"body"`
}
<|endoftext|>


<commit_before>package amqpfeed

import (
	"encoding/json"
	"github.com/StefanKjartansson/eventhub"
	"github.com/streadway/amqp"
	"log"
	"testing"
)

const (
	uri          = "amqp://guest:guest@localhost:5672/"
	exchange     = "ha.test-exchange"
	exchangeType = "direct"
	queue        = "test-queue"
	bindingKey   = "test-key"
	consumerTag  = "simple-consumer"
	lifetime     = 0
)

func publish(amqpURI, exchange, exchangeType, routingKey string, event eventhub.Event, reliable, durable, auto_delete bool) error {

	b, err := json.Marshal(event)
	if err != nil {
		return err
	}

	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		return err
	}
	defer connection.Close()

	channel, err := connection.Channel()
	if err != nil {
		return err
	}

	if err := channel.ExchangeDeclare(
		exchange,     // name
		exchangeType, // type
		durable,      // durable
		auto_delete,  // auto-deleted
		false,        // internal
		false,        // noWait
		nil,          // arguments
	); err != nil {
		return err
	}

	if reliable {
		if err := channel.Confirm(false); err != nil {
			return err
		}
		ack, nack := channel.NotifyConfirm(make(chan uint64, 1), make(chan uint64, 1))
		defer confirmOne(ack, nack)
	}

	if err = channel.Publish(
		exchange,   // publish to an exchange
		routingKey, // routing to 0 or more queues
		false,      // mandatory
		false,      // immediate
		amqp.Publishing{
			Headers:         amqp.Table{},
			ContentType:     "application/json",
			ContentEncoding: "",
			Body:            b,
			DeliveryMode:    amqp.Transient, // 1=non-persistent, 2=persistent
			Priority:        0,              // 0-9
		},
	); err != nil {
		return err
	}

	return nil
}

func confirmOne(ack, nack chan uint64) {
	select {
	case tag := <-ack:
		log.Printf("confirmed delivery with delivery tag: %d", tag)
	case tag := <-nack:
		log.Printf("failed delivery of delivery tag: %d", tag)
	}
}

func TestAMQPFeed(t *testing.T) {

	_, err := amqp.Dial(uri)
	if err != nil {
		t.Skip("Skipping amqp test, no rabbitmq available")
	}

	c, err := NewConsumer(uri, exchange, exchangeType, queue, bindingKey, consumerTag, true, true)
	if err != nil {
		t.Logf("%s", err)
	}

	t.Logf("%v", c)

	e := eventhub.Event{
		Key:         "foo.bar",
		Description: "My event",
		Importance:  3,
		Origin:      "mysystem",
		Entities:    []string{"ns/foo", "ns/moo"},
		Actors:      []string{"someone"},
	}

	err = publish(uri, exchange, exchangeType, bindingKey, e, true, true, true)

	if err != nil {
		t.Logf("%s", err)
	}

	i := <-c.Updates()

	t.Logf("%v", i)

	if c.Close() != nil {
		t.Logf("%s", err)
	}
}
<commit_msg>Use constructor function<commit_after>package amqpfeed

import (
	"encoding/json"
	"github.com/StefanKjartansson/eventhub"
	"github.com/streadway/amqp"
	"log"
	"testing"
)

const (
	uri          = "amqp://guest:guest@localhost:5672/"
	exchange     = "ha.test-exchange"
	exchangeType = "direct"
	queue        = "test-queue"
	bindingKey   = "test-key"
	consumerTag  = "simple-consumer"
	lifetime     = 0
)

func publish(amqpURI, exchange, exchangeType, routingKey string, event *eventhub.Event, reliable, durable, auto_delete bool) error {

	b, err := json.Marshal(event)
	if err != nil {
		return err
	}

	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		return err
	}
	defer connection.Close()

	channel, err := connection.Channel()
	if err != nil {
		return err
	}

	if err := channel.ExchangeDeclare(
		exchange,     // name
		exchangeType, // type
		durable,      // durable
		auto_delete,  // auto-deleted
		false,        // internal
		false,        // noWait
		nil,          // arguments
	); err != nil {
		return err
	}

	if reliable {
		if err := channel.Confirm(false); err != nil {
			return err
		}
		ack, nack := channel.NotifyConfirm(make(chan uint64, 1), make(chan uint64, 1))
		defer confirmOne(ack, nack)
	}

	if err = channel.Publish(
		exchange,   // publish to an exchange
		routingKey, // routing to 0 or more queues
		false,      // mandatory
		false,      // immediate
		amqp.Publishing{
			Headers:         amqp.Table{},
			ContentType:     "application/json",
			ContentEncoding: "",
			Body:            b,
			DeliveryMode:    amqp.Transient, // 1=non-persistent, 2=persistent
			Priority:        0,              // 0-9
		},
	); err != nil {
		return err
	}

	return nil
}

func confirmOne(ack, nack chan uint64) {
	select {
	case tag := <-ack:
		log.Printf("confirmed delivery with delivery tag: %d", tag)
	case tag := <-nack:
		log.Printf("failed delivery of delivery tag: %d", tag)
	}
}

func TestAMQPFeed(t *testing.T) {

	_, err := amqp.Dial(uri)
	if err != nil {
		t.Skip("Skipping amqp test, no rabbitmq available")
	}

	c, err := NewConsumer(uri, exchange, exchangeType, queue, bindingKey, consumerTag, true, true)
	if err != nil {
		t.Logf("%s", err)
	}

	t.Logf("%v", c)

	e := eventhub.NewEvent(
		"foo.bar",
		nil,
		nil,
		"My event",
		3,
		"mysystem",
		[]string{"ns/foo", "ns/bar"},
		[]string{"someone"},
		nil,
		nil)

	err = publish(uri, exchange, exchangeType, bindingKey, e, true, true, true)

	if err != nil {
		t.Logf("%s", err)
	}

	i := <-c.Updates()

	t.Logf("%v", i)

	if c.Close() != nil {
		t.Logf("%s", err)
	}
}
<|endoftext|>


<commit_before>package testutil

var IntrospectionQuery = `
  query IntrospectionQuery {
    __schema {
      queryType { name }
      mutationType { name }
      subscriptionType { name }
      types {
        ...FullType
      }
      directives {
        name
        description
		locations
        args {
          ...InputValue
        }
        # deprecated, but included for coverage till removed
		onOperation
        onFragment
        onField
      }
    }
  }

  fragment FullType on __Type {
    kind
    name
    description
    fields(includeDeprecated: true) {
      name
      description
      args {
        ...InputValue
      }
      type {
        ...TypeRef
      }
      isDeprecated
      deprecationReason
    }
    inputFields {
      ...InputValue
    }
    interfaces {
      ...TypeRef
    }
    enumValues(includeDeprecated: true) {
      name
      description
      isDeprecated
      deprecationReason
    }
    possibleTypes {
      ...TypeRef
    }
  }

  fragment InputValue on __InputValue {
    name
    description
    type { ...TypeRef }
    defaultValue
  }

  fragment TypeRef on __Type {
    kind
    name
    ofType {
      kind
      name
      ofType {
        kind
        name
        ofType {
          kind
          name
        }
      }
    }
  }
`
<commit_msg>Deepen introspection query<commit_after>package testutil

var IntrospectionQuery = `
  query IntrospectionQuery {
    __schema {
      queryType { name }
      mutationType { name }
      subscriptionType { name }
      types {
        ...FullType
      }
      directives {
        name
        description
		locations
        args {
          ...InputValue
        }
        # deprecated, but included for coverage till removed
		onOperation
        onFragment
        onField
      }
    }
  }

  fragment FullType on __Type {
    kind
    name
    description
    fields(includeDeprecated: true) {
      name
      description
      args {
        ...InputValue
      }
      type {
        ...TypeRef
      }
      isDeprecated
      deprecationReason
    }
    inputFields {
      ...InputValue
    }
    interfaces {
      ...TypeRef
    }
    enumValues(includeDeprecated: true) {
      name
      description
      isDeprecated
      deprecationReason
    }
    possibleTypes {
      ...TypeRef
    }
  }

  fragment InputValue on __InputValue {
    name
    description
    type { ...TypeRef }
    defaultValue
  }

  fragment TypeRef on __Type {
    kind
    name
    ofType {
      kind
      name
      ofType {
        kind
        name
        ofType {
          kind
          name
          ofType {
            kind
            name
            ofType {
              kind
              name
              ofType {
                kind
                name
                ofType {
                  kind
                  name
                }
              }
            }
          }
        }
      }
    }
  }
`
<|endoftext|>


<commit_before>package client

import (
	"irken/test"
	"testing"
)

func TestLexValid(t *testing.T) {
	message := ":prefix COMMAND param1 param2 :param 3 :-) yeah!?"
	prefix, command, params, err := lexMsg(message)
	if err != nil {
		t.Error(err)
	}
	test.Check(t, prefix, "prefix")
	test.Check(t, command, "COMMAND")
	test.Check(t, params[0], "param1")
	test.Check(t, params[1], "param2")
	test.Check(t, params[2], "param 3 :-) yeah!?")
}

func TestLexInValid(t *testing.T) {
	message := ":prefix"
	message2 := ":prefix "
	_, _, _, err := lexMsg(message)
	if err == nil {
		t.Errorf("Illegal message is not error reported")
	}
	_, _, _, err = lexMsg(message2)
	if err == nil {
		t.Errorf("Illegal message is not error reported")
	}

}

func TestLexNoParams(t *testing.T) {
	message := "COMMAND"
	prefix, command, params, err := lexMsg(message)
	if err != nil {
		t.Error(err)
	}
	test.Check(t, prefix, "")
	test.Check(t, command, "COMMAND")
	if len(params) != 0 {
		t.Errorf("Reported fake parameters")
	}
}

func TestJoin(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com JOIN #chan"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx has joined #chan"
	expCont := "#chan"
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestMode(t *testing.T) {
	// TODO
}

func TestQuit(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com QUIT :Later suckerz"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx has quit (Later suckerz)"
	expCont := ""
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestPart(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com PART #chan"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx has left #chan"
	expCont := "#chan"
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestPrivMsg(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com PRIVMSG #chan :Octotastic!"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx: Octotastic!"
	expCont := "#chan"
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestResolveNick(t *testing.T) {
	input := "_mrx!blabla@haxxor.com"
	nick, err := resolveNick(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	exp := "_mrx"
	test.Check(t, nick, exp)
}

func TestResolveInvalidNick(t *testing.T) {
	input := "_mrxblabla@haxxor.com"
	_, err := resolveNick(input)
	if err == nil {
		t.Errorf("Should not parse")
	}
}
<commit_msg>test mode<commit_after>package client

import (
	"irken/test"
	"testing"
)

func TestLexValid(t *testing.T) {
	message := ":prefix COMMAND param1 param2 :param 3 :-) yeah!?"
	prefix, command, params, err := lexMsg(message)
	if err != nil {
		t.Error(err)
	}
	test.Check(t, prefix, "prefix")
	test.Check(t, command, "COMMAND")
	test.Check(t, params[0], "param1")
	test.Check(t, params[1], "param2")
	test.Check(t, params[2], "param 3 :-) yeah!?")
}

func TestLexInValid(t *testing.T) {
	message := ":prefix"
	message2 := ":prefix "
	_, _, _, err := lexMsg(message)
	if err == nil {
		t.Errorf("Illegal message is not error reported")
	}
	_, _, _, err = lexMsg(message2)
	if err == nil {
		t.Errorf("Illegal message is not error reported")
	}

}

func TestLexNoParams(t *testing.T) {
	message := "COMMAND"
	prefix, command, params, err := lexMsg(message)
	if err != nil {
		t.Error(err)
	}
	test.Check(t, prefix, "")
	test.Check(t, command, "COMMAND")
	if len(params) != 0 {
		t.Errorf("Reported fake parameters")
	}
}

func TestJoin(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com JOIN #chan"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx has joined #chan"
	expCont := "#chan"
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestMode(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com MODE #chan +i -l"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx changed mode +i -l for #chan"
	expCont := "#chan"
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestQuit(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com QUIT :Later suckerz"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx has quit (Later suckerz)"
	expCont := ""
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestPart(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com PART #chan"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx has left #chan"
	expCont := "#chan"
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestPrivMsg(t *testing.T) {
	input := ":_mrx!blabla@haxxor.com PRIVMSG #chan :Octotastic! I like pie btw :)"
	msg, cont, err := ParseServerMsg(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	expMsg := "_mrx: Octotastic! I like pie btw :)"
	expCont := "#chan"
	test.Check(t, msg, expMsg)
	test.Check(t, cont, expCont)
}

func TestResolveNick(t *testing.T) {
	input := "_mrx!blabla@haxxor.com"
	nick, err := resolveNick(input)
	if err != nil {
		t.Errorf("Should parse!")
	}
	exp := "_mrx"
	test.Check(t, nick, exp)
}

func TestResolveInvalidNick(t *testing.T) {
	input := "_mrxblabla@haxxor.com"
	_, err := resolveNick(input)
	if err == nil {
		t.Errorf("Should not parse")
	}
}
<|endoftext|>


<commit_before>package client

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/structs"

	cstructs "github.com/hashicorp/nomad/client/driver/structs"
)

// TaskRunner is used to wrap a task within an allocation and provide the execution context.
type TaskRunner struct {
	config         *config.Config
	updater        TaskStateUpdater
	logger         *log.Logger
	ctx            *driver.ExecContext
	allocID        string
	restartTracker restartTracker

	task     *structs.Task
	state    *structs.TaskState
	updateCh chan *structs.Task
	handle   driver.DriverHandle

	destroy     bool
	destroyCh   chan struct{}
	destroyLock sync.Mutex
	waitCh      chan struct{}

	snapshotLock sync.Mutex
}

// taskRunnerState is used to snapshot the state of the task runner
type taskRunnerState struct {
	Task     *structs.Task
	HandleID string
}

// TaskStateUpdater is used to signal that tasks state has changed.
type TaskStateUpdater func(taskName string)

// NewTaskRunner is used to create a new task context
func NewTaskRunner(logger *log.Logger, config *config.Config,
	updater TaskStateUpdater, ctx *driver.ExecContext,
	allocID string, task *structs.Task, state *structs.TaskState,
	restartTracker restartTracker) *TaskRunner {

	tc := &TaskRunner{
		config:         config,
		updater:        updater,
		logger:         logger,
		restartTracker: restartTracker,
		ctx:            ctx,
		allocID:        allocID,
		task:           task,
		state:          state,
		updateCh:       make(chan *structs.Task, 8),
		destroyCh:      make(chan struct{}),
		waitCh:         make(chan struct{}),
	}
	return tc
}

// WaitCh returns a channel to wait for termination
func (r *TaskRunner) WaitCh() <-chan struct{} {
	return r.waitCh
}

// stateFilePath returns the path to our state file
func (r *TaskRunner) stateFilePath() string {
	// Get the MD5 of the task name
	hashVal := md5.Sum([]byte(r.task.Name))
	hashHex := hex.EncodeToString(hashVal[:])
	dirName := fmt.Sprintf("task-%s", hashHex)

	// Generate the path
	path := filepath.Join(r.config.StateDir, "alloc", r.allocID,
		dirName, "state.json")
	return path
}

// RestoreState is used to restore our state
func (r *TaskRunner) RestoreState() error {
	// Load the snapshot
	var snap taskRunnerState
	if err := restoreState(r.stateFilePath(), &snap); err != nil {
		return err
	}

	// Restore fields
	r.task = snap.Task

	// Restore the driver
	if snap.HandleID != "" {
		driver, err := r.createDriver()
		if err != nil {
			return err
		}

		handle, err := driver.Open(r.ctx, snap.HandleID)

		// In the case it fails, we relaunch the task in the Run() method.
		if err != nil {
			r.logger.Printf("[ERR] client: failed to open handle to task '%s' for alloc '%s': %v",
				r.task.Name, r.allocID, err)
			return nil
		}
		r.handle = handle
	}
	return nil
}

// SaveState is used to snapshot our state
func (r *TaskRunner) SaveState() error {
	r.snapshotLock.Lock()
	defer r.snapshotLock.Unlock()
	snap := taskRunnerState{
		Task: r.task,
	}
	if r.handle != nil {
		snap.HandleID = r.handle.ID()
	}
	return persistState(r.stateFilePath(), &snap)
}

// DestroyState is used to cleanup after ourselves
func (r *TaskRunner) DestroyState() error {
	return os.RemoveAll(r.stateFilePath())
}

func (r *TaskRunner) appendEvent(event *structs.TaskEvent) {
	capacity := 10
	if r.state.Events == nil {
		r.state.Events = make([]*structs.TaskEvent, 0, capacity)
	}

	// If we hit capacity, then shift it.
	if len(r.state.Events) == capacity {
		old := r.state.Events
		r.state.Events = make([]*structs.TaskEvent, 0, capacity)
		r.state.Events = append(r.state.Events, old[1:]...)
	}

	r.state.Events = append(r.state.Events, event)
}

// setState is used to update the state of the task runner
func (r *TaskRunner) setState(state string, event *structs.TaskEvent) {
	// Update the task.
	r.state.State = state
	r.appendEvent(event)

	// Persist our state to disk.
	if err := r.SaveState(); err != nil {
		r.logger.Printf("[ERR] client: failed to save state of Task Runner: %v", r.task.Name)
	}

	// Indicate the task has been updated.
	r.updater(r.task.Name)
}

// createDriver makes a driver for the task
func (r *TaskRunner) createDriver() (driver.Driver, error) {
	driverCtx := driver.NewDriverContext(r.task.Name, r.config, r.config.Node, r.logger)
	driver, err := driver.NewDriver(r.task.Driver, driverCtx)
	if err != nil {
		err = fmt.Errorf("failed to create driver '%s' for alloc %s: %v",
			r.task.Driver, r.allocID, err)
		r.logger.Printf("[ERR] client: %s", err)
	}
	return driver, err
}

// startTask is used to start the task if there is no handle
func (r *TaskRunner) startTask() error {
	// Create a driver
	driver, err := r.createDriver()
	if err != nil {
		e := structs.NewTaskEvent(structs.TaskDriverFailure).SetDriverError(err)
		r.setState(structs.TaskStateDead, e)
		return err
	}

	// Start the job
	handle, err := driver.Start(r.ctx, r.task)
	if err != nil {
		r.logger.Printf("[ERR] client: failed to start task '%s' for alloc '%s': %v",
			r.task.Name, r.allocID, err)
		e := structs.NewTaskEvent(structs.TaskDriverFailure).
			SetDriverError(fmt.Errorf("failed to start: %v", err))
		r.setState(structs.TaskStateDead, e)
		return err
	}
	r.handle = handle
	r.setState(structs.TaskStateRunning, structs.NewTaskEvent(structs.TaskStarted))
	return nil
}

// Run is a long running routine used to manage the task
func (r *TaskRunner) Run() {
	defer close(r.waitCh)
	r.logger.Printf("[DEBUG] client: starting task context for '%s' (alloc '%s')",
		r.task.Name, r.allocID)

	r.run(false)
	return
}

func (r *TaskRunner) run(forceStart bool) {
	// Start the task if not yet started or it is being forced.
	if r.handle == nil || forceStart {
		if err := r.startTask(); err != nil {
			return
		}
	}

	// Store the errors that caused use to stop waiting for updates.
	var waitRes *cstructs.WaitResult
	var destroyErr error
	destroyed := false

OUTER:
	// Wait for updates
	for {
		select {
		case waitRes = <-r.handle.WaitCh():
			break OUTER
		case update := <-r.updateCh:
			// Update
			r.task = update
			if err := r.handle.Update(update); err != nil {
				r.logger.Printf("[ERR] client: failed to update task '%s' for alloc '%s': %v", r.task.Name, r.allocID, err)
			}
		case <-r.destroyCh:
			// Send the kill signal, and use the WaitCh to block until complete
			if err := r.handle.Kill(); err != nil {
				r.logger.Printf("[ERR] client: failed to kill task '%s' for alloc '%s': %v", r.task.Name, r.allocID, err)
				destroyErr = err
			}
			destroyed = true
		}
	}

	// If the user destroyed the task, we do not attempt to do any restarts.
	if destroyed {
		r.setState(structs.TaskStateDead, structs.NewTaskEvent(structs.TaskKilled).SetKillError(destroyErr))
		return
	}

	// Log whether the task was successful or not.
	if !waitRes.Successful() {
		r.logger.Printf("[ERR] client: failed to complete task '%s' for alloc '%s': %v", r.task.Name, r.allocID, waitRes)
	} else {
		r.logger.Printf("[INFO] client: completed task '%s' for alloc '%s'", r.task.Name, r.allocID)
	}

	// Check if we should restart. If not mark task as dead and exit.
	waitEvent := r.waitErrorToEvent(waitRes)
	shouldRestart, when := r.restartTracker.nextRestart()
	if !shouldRestart {
		r.logger.Printf("[INFO] client: Not restarting task: %v for alloc: %v ", r.task.Name, r.allocID)
		r.setState(structs.TaskStateDead, waitEvent)
		return
	}

	r.logger.Printf("[INFO] client: Restarting Task: %v", r.task.Name)
	r.logger.Printf("[DEBUG] client: Sleeping for %v before restarting Task %v", when, r.task.Name)
	r.setState(structs.TaskStatePending, waitEvent)

	// Sleep but watch for destroy events.
	select {
	case <-time.After(when):
	case <-r.destroyCh:
	}

	// Destroyed while we were waiting to restart, so abort.
	r.destroyLock.Lock()
	destroyed = r.destroy
	r.destroyLock.Unlock()
	if destroyed {
		r.logger.Printf("[DEBUG] client: Not restarting task: %v because it's destroyed by user", r.task.Name)
		r.setState(structs.TaskStateDead, structs.NewTaskEvent(structs.TaskKilled))
		return
	}

	// Recurse on ourselves and force the start since we are restarting the task.
	r.run(true)
	// TODO: Alex
	return
}

// Helper function for converting a WaitResult into a TaskTerminated event.
func (r *TaskRunner) waitErrorToEvent(res *cstructs.WaitResult) *structs.TaskEvent {
	e := structs.NewTaskEvent(structs.TaskTerminated).SetExitCode(res.ExitCode).SetSignal(res.Signal)
	if res.Err != nil {
		e.SetExitMessage(res.Err.Error())
	}
	return e
}

// Update is used to update the task of the context
func (r *TaskRunner) Update(update *structs.Task) {
	select {
	case r.updateCh <- update:
	default:
		r.logger.Printf("[ERR] client: dropping task update '%s' (alloc '%s')",
			update.Name, r.allocID)
	}
}

// Destroy is used to indicate that the task context should be destroyed
func (r *TaskRunner) Destroy() {
	r.destroyLock.Lock()
	defer r.destroyLock.Unlock()

	if r.destroy {
		return
	}
	r.destroy = true
	close(r.destroyCh)
}
<commit_msg>Use loop not recursion<commit_after>package client

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/structs"

	cstructs "github.com/hashicorp/nomad/client/driver/structs"
)

// TaskRunner is used to wrap a task within an allocation and provide the execution context.
type TaskRunner struct {
	config         *config.Config
	updater        TaskStateUpdater
	logger         *log.Logger
	ctx            *driver.ExecContext
	allocID        string
	restartTracker restartTracker

	task     *structs.Task
	state    *structs.TaskState
	updateCh chan *structs.Task
	handle   driver.DriverHandle

	destroy     bool
	destroyCh   chan struct{}
	destroyLock sync.Mutex
	waitCh      chan struct{}

	snapshotLock sync.Mutex
}

// taskRunnerState is used to snapshot the state of the task runner
type taskRunnerState struct {
	Task     *structs.Task
	HandleID string
}

// TaskStateUpdater is used to signal that tasks state has changed.
type TaskStateUpdater func(taskName string)

// NewTaskRunner is used to create a new task context
func NewTaskRunner(logger *log.Logger, config *config.Config,
	updater TaskStateUpdater, ctx *driver.ExecContext,
	allocID string, task *structs.Task, state *structs.TaskState,
	restartTracker restartTracker) *TaskRunner {

	tc := &TaskRunner{
		config:         config,
		updater:        updater,
		logger:         logger,
		restartTracker: restartTracker,
		ctx:            ctx,
		allocID:        allocID,
		task:           task,
		state:          state,
		updateCh:       make(chan *structs.Task, 8),
		destroyCh:      make(chan struct{}),
		waitCh:         make(chan struct{}),
	}
	return tc
}

// WaitCh returns a channel to wait for termination
func (r *TaskRunner) WaitCh() <-chan struct{} {
	return r.waitCh
}

// stateFilePath returns the path to our state file
func (r *TaskRunner) stateFilePath() string {
	// Get the MD5 of the task name
	hashVal := md5.Sum([]byte(r.task.Name))
	hashHex := hex.EncodeToString(hashVal[:])
	dirName := fmt.Sprintf("task-%s", hashHex)

	// Generate the path
	path := filepath.Join(r.config.StateDir, "alloc", r.allocID,
		dirName, "state.json")
	return path
}

// RestoreState is used to restore our state
func (r *TaskRunner) RestoreState() error {
	// Load the snapshot
	var snap taskRunnerState
	if err := restoreState(r.stateFilePath(), &snap); err != nil {
		return err
	}

	// Restore fields
	r.task = snap.Task

	// Restore the driver
	if snap.HandleID != "" {
		driver, err := r.createDriver()
		if err != nil {
			return err
		}

		handle, err := driver.Open(r.ctx, snap.HandleID)

		// In the case it fails, we relaunch the task in the Run() method.
		if err != nil {
			r.logger.Printf("[ERR] client: failed to open handle to task '%s' for alloc '%s': %v",
				r.task.Name, r.allocID, err)
			return nil
		}
		r.handle = handle
	}
	return nil
}

// SaveState is used to snapshot our state
func (r *TaskRunner) SaveState() error {
	r.snapshotLock.Lock()
	defer r.snapshotLock.Unlock()
	snap := taskRunnerState{
		Task: r.task,
	}
	if r.handle != nil {
		snap.HandleID = r.handle.ID()
	}
	return persistState(r.stateFilePath(), &snap)
}

// DestroyState is used to cleanup after ourselves
func (r *TaskRunner) DestroyState() error {
	return os.RemoveAll(r.stateFilePath())
}

func (r *TaskRunner) appendEvent(event *structs.TaskEvent) {
	capacity := 10
	if r.state.Events == nil {
		r.state.Events = make([]*structs.TaskEvent, 0, capacity)
	}

	// If we hit capacity, then shift it.
	if len(r.state.Events) == capacity {
		old := r.state.Events
		r.state.Events = make([]*structs.TaskEvent, 0, capacity)
		r.state.Events = append(r.state.Events, old[1:]...)
	}

	r.state.Events = append(r.state.Events, event)
}

// setState is used to update the state of the task runner
func (r *TaskRunner) setState(state string, event *structs.TaskEvent) {
	// Update the task.
	r.state.State = state
	r.appendEvent(event)

	// Persist our state to disk.
	if err := r.SaveState(); err != nil {
		r.logger.Printf("[ERR] client: failed to save state of Task Runner: %v", r.task.Name)
	}

	// Indicate the task has been updated.
	r.updater(r.task.Name)
}

// createDriver makes a driver for the task
func (r *TaskRunner) createDriver() (driver.Driver, error) {
	driverCtx := driver.NewDriverContext(r.task.Name, r.config, r.config.Node, r.logger)
	driver, err := driver.NewDriver(r.task.Driver, driverCtx)
	if err != nil {
		err = fmt.Errorf("failed to create driver '%s' for alloc %s: %v",
			r.task.Driver, r.allocID, err)
		r.logger.Printf("[ERR] client: %s", err)
	}
	return driver, err
}

// startTask is used to start the task if there is no handle
func (r *TaskRunner) startTask() error {
	// Create a driver
	driver, err := r.createDriver()
	if err != nil {
		e := structs.NewTaskEvent(structs.TaskDriverFailure).SetDriverError(err)
		r.setState(structs.TaskStateDead, e)
		return err
	}

	// Start the job
	handle, err := driver.Start(r.ctx, r.task)
	if err != nil {
		r.logger.Printf("[ERR] client: failed to start task '%s' for alloc '%s': %v",
			r.task.Name, r.allocID, err)
		e := structs.NewTaskEvent(structs.TaskDriverFailure).
			SetDriverError(fmt.Errorf("failed to start: %v", err))
		r.setState(structs.TaskStateDead, e)
		return err
	}
	r.handle = handle
	r.setState(structs.TaskStateRunning, structs.NewTaskEvent(structs.TaskStarted))
	return nil
}

// Run is a long running routine used to manage the task
func (r *TaskRunner) Run() {
	defer close(r.waitCh)
	r.logger.Printf("[DEBUG] client: starting task context for '%s' (alloc '%s')",
		r.task.Name, r.allocID)

	r.run()
	return
}

func (r *TaskRunner) run() {
	var forceStart bool
	for {
		// Start the task if not yet started or it is being forced.
		if r.handle == nil || forceStart {
			forceStart = false
			if err := r.startTask(); err != nil {
				return
			}
		}

		// Store the errors that caused use to stop waiting for updates.
		var waitRes *cstructs.WaitResult
		var destroyErr error
		destroyed := false

	OUTER:
		// Wait for updates
		for {
			select {
			case waitRes = <-r.handle.WaitCh():
				break OUTER
			case update := <-r.updateCh:
				// Update
				r.task = update
				if err := r.handle.Update(update); err != nil {
					r.logger.Printf("[ERR] client: failed to update task '%s' for alloc '%s': %v", r.task.Name, r.allocID, err)
				}
			case <-r.destroyCh:
				// Send the kill signal, and use the WaitCh to block until complete
				if err := r.handle.Kill(); err != nil {
					r.logger.Printf("[ERR] client: failed to kill task '%s' for alloc '%s': %v", r.task.Name, r.allocID, err)
					destroyErr = err
				}
				destroyed = true
			}
		}

		// If the user destroyed the task, we do not attempt to do any restarts.
		if destroyed {
			r.setState(structs.TaskStateDead, structs.NewTaskEvent(structs.TaskKilled).SetKillError(destroyErr))
			return
		}

		// Log whether the task was successful or not.
		if !waitRes.Successful() {
			r.logger.Printf("[ERR] client: failed to complete task '%s' for alloc '%s': %v", r.task.Name, r.allocID, waitRes)
		} else {
			r.logger.Printf("[INFO] client: completed task '%s' for alloc '%s'", r.task.Name, r.allocID)
		}

		// Check if we should restart. If not mark task as dead and exit.
		waitEvent := r.waitErrorToEvent(waitRes)
		shouldRestart, when := r.restartTracker.nextRestart()
		if !shouldRestart {
			r.logger.Printf("[INFO] client: Not restarting task: %v for alloc: %v ", r.task.Name, r.allocID)
			r.setState(structs.TaskStateDead, waitEvent)
			return
		}

		r.logger.Printf("[INFO] client: Restarting Task: %v", r.task.Name)
		r.logger.Printf("[DEBUG] client: Sleeping for %v before restarting Task %v", when, r.task.Name)
		r.setState(structs.TaskStatePending, waitEvent)

		// Sleep but watch for destroy events.
		select {
		case <-time.After(when):
		case <-r.destroyCh:
		}

		// Destroyed while we were waiting to restart, so abort.
		r.destroyLock.Lock()
		destroyed = r.destroy
		r.destroyLock.Unlock()
		if destroyed {
			r.logger.Printf("[DEBUG] client: Not restarting task: %v because it's destroyed by user", r.task.Name)
			r.setState(structs.TaskStateDead, structs.NewTaskEvent(structs.TaskKilled))
			return
		}

		// Set force start because we are restarting the task.
		forceStart = true
	}
	return
}

// Helper function for converting a WaitResult into a TaskTerminated event.
func (r *TaskRunner) waitErrorToEvent(res *cstructs.WaitResult) *structs.TaskEvent {
	e := structs.NewTaskEvent(structs.TaskTerminated).SetExitCode(res.ExitCode).SetSignal(res.Signal)
	if res.Err != nil {
		e.SetExitMessage(res.Err.Error())
	}
	return e
}

// Update is used to update the task of the context
func (r *TaskRunner) Update(update *structs.Task) {
	select {
	case r.updateCh <- update:
	default:
		r.logger.Printf("[ERR] client: dropping task update '%s' (alloc '%s')",
			update.Name, r.allocID)
	}
}

// Destroy is used to indicate that the task context should be destroyed
func (r *TaskRunner) Destroy() {
	r.destroyLock.Lock()
	defer r.destroyLock.Unlock()

	if r.destroy {
		return
	}
	r.destroy = true
	close(r.destroyCh)
}
<|endoftext|>