<commit_before>package logger

import (
	"bytes"
	"errors"
	"net/http"
	"encoding/json"
)

const PROVIDER_TELEGRAM = "telegram"

type TelegramProvider struct {
	url     string
	chatIds []string
}

func NewTelegramProvider(url string, chatIds []string) (*TelegramProvider, error) {

	if len(url) == 0 {
		return nil, errors.New("Empty telegram url.")
	} else if chatIds == nil {
		return nil, errors.New("Empty telegram chat ids.")
	}

	provider := &TelegramProvider{
		url:     url,
		chatIds: chatIds,
	}

	return provider, nil
}

func (p TelegramProvider) GetID() string {
	return PROVIDER_TELEGRAM
}

func (p TelegramProvider) Log(msg []byte) {
	p.send("Log message", msg)
}

func (p TelegramProvider) Error(msg []byte) {
	p.send("Error message", msg)
}

func (p TelegramProvider) Fatal(msg []byte) {
	p.send("Fatal message", msg)
}

func (p TelegramProvider) Debug(msg []byte) {
	p.send("Debug message", msg)
}

func (p TelegramProvider) send(subject string, body []byte) {
	go tg_send(p.url, p.chatIds, subject, body)
}

func tg_send(url string, chatIds []string, subject string, body []byte) {
	msg := map[string]interface{}{
		"chat_id": "",
		"text": subject + "\n" + string(body),
	}
	client := &http.Client{}

	for _, chatId := range chatIds {
		msg["chat_id"] = chatId
		jsonStr, err := json.Marshal(msg)
		if err != nil {
			return;
		}

		req, _ := http.NewRequest("POST", url, bytes.NewBuffer(jsonStr))
		req.Header.Set("Content-Type", "application/json")

		resp, _ := client.Do(req)
		defer resp.Body.Close()
	}
}
<commit_msg>check if body is empty<commit_after>package logger

import (
	"bytes"
	"encoding/json"
	"errors"
	"net/http"
)

const PROVIDER_TELEGRAM = "telegram"

type TelegramProvider struct {
	url     string
	chatIds []string
}

func NewTelegramProvider(url string, chatIds []string) (*TelegramProvider, error) {

	if len(url) == 0 {
		return nil, errors.New("Empty telegram url.")
	} else if chatIds == nil {
		return nil, errors.New("Empty telegram chat ids.")
	}

	provider := &TelegramProvider{
		url:     url,
		chatIds: chatIds,
	}

	return provider, nil
}

func (p TelegramProvider) GetID() string {
	return PROVIDER_TELEGRAM
}

func (p TelegramProvider) Log(msg []byte) {
	p.send("Log message", msg)
}

func (p TelegramProvider) Error(msg []byte) {
	p.send("Error message", msg)
}

func (p TelegramProvider) Fatal(msg []byte) {
	p.send("Fatal message", msg)
}

func (p TelegramProvider) Debug(msg []byte) {
	p.send("Debug message", msg)
}

func (p TelegramProvider) send(subject string, body []byte) {
	go tg_send(p.url, p.chatIds, subject, body)
}

func tg_send(url string, chatIds []string, subject string, body []byte) {
	msg := map[string]interface{}{
		"chat_id": "",
		"text":    subject + "\n" + string(body),
	}
	client := &http.Client{}

	for _, chatId := range chatIds {
		msg["chat_id"] = chatId
		jsonStr, err := json.Marshal(msg)
		if err != nil {
			return
		}

		req, _ := http.NewRequest("POST", url, bytes.NewBuffer(jsonStr))
		req.Header.Set("Content-Type", "application/json")

		resp, err := client.Do(req)
		if err != nil {
			return
		}

		defer resp.Body.Close()
	}
}
<|endoftext|>


<commit_before>package template

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/HeavyHorst/memkv"
	"github.com/HeavyHorst/remco/backends"
	"github.com/HeavyHorst/remco/log"
	"github.com/HeavyHorst/remco/template/fileutil"
	"github.com/Sirupsen/logrus"
	"github.com/flosch/pongo2"
	flag "github.com/spf13/pflag"
)

type StoreConfig struct {
	backends.StoreClient
	Name     string
	Onetime  bool
	Watch    bool
	Prefix   string
	Interval int
	Keys     []string
	store    *memkv.Store
}

// Resource is the representation of a parsed template resource.
type Resource struct {
	storeClients []StoreConfig
	funcMap      map[string]interface{}
	store        memkv.Store
	sources      []*SrcDst
}

var ErrEmptySrc = errors.New("empty src template")

// NewResource creates a Resource.
func NewResource(storeClients []StoreConfig, sources []*SrcDst) (*Resource, error) {
	if len(storeClients) == 0 {
		return nil, errors.New("A valid StoreClient is required.")
	}

	for _, v := range sources {
		if v.Src == "" {
			return nil, ErrEmptySrc
		}
	}

	tr := &Resource{
		storeClients: storeClients,
		store:        memkv.New(),
		funcMap:      newFuncMap(),
		sources:      sources,
	}

	// initialize the inidividual backend memkv Stores
	for i := range tr.storeClients {
		store := memkv.New()
		tr.storeClients[i].store = &store
	}

	addFuncs(tr.funcMap, tr.store.FuncMap)

	return tr, nil
}

// NewResourceFromFlags creates a Resource from the provided commandline flags.
func NewResourceFromFlags(s backends.Store, flags *flag.FlagSet, watch bool) (*Resource, error) {
	src, _ := flags.GetString("src")
	dst, _ := flags.GetString("dst")
	keys, _ := flags.GetStringSlice("keys")
	fileMode, _ := flags.GetString("fileMode")
	prefix, _ := flags.GetString("prefix")
	reloadCmd, _ := flags.GetString("reload_cmd")
	checkCmd, _ := flags.GetString("check_cmd")
	onetime, _ := flags.GetBool("onetime")
	interval, _ := flags.GetInt("interval")
	UID := os.Geteuid()
	GID := os.Getegid()

	sd := &SrcDst{
		Src:       src,
		Dst:       dst,
		Mode:      fileMode,
		ReloadCmd: reloadCmd,
		CheckCmd:  checkCmd,
		GID:       GID,
		UID:       UID,
	}

	b := StoreConfig{
		StoreClient: s.Client,
		Name:        s.Name,
		Onetime:     onetime,
		Prefix:      prefix,
		Watch:       watch,
		Interval:    interval,
		Keys:        keys,
	}

	return NewResource([]StoreConfig{b}, []*SrcDst{sd})
}

// setVars sets the Vars for template resource.
func (t *Resource) setVars(storeClient StoreConfig) error {
	var err error

	log.WithFields(logrus.Fields{
		"store":      storeClient.Name,
		"key_prefix": storeClient.Prefix,
	}).Debug("Retrieving keys")

	result, err := storeClient.GetValues(appendPrefix(storeClient.Prefix, storeClient.Keys))
	if err != nil {
		return err
	}

	storeClient.store.Purge()

	for key, value := range result {
		storeClient.store.Set(path.Join("/", strings.TrimPrefix(key, storeClient.Prefix)), value)
	}

	//merge all stores
	t.store.Purge()
	for _, v := range t.storeClients {
		for _, kv := range v.store.GetAllKVs() {
			if t.store.Exists(kv.Key) {
				log.Warning("Key collision - " + kv.Key)
			}
			t.store.Set(kv.Key, kv.Value)
		}
	}

	return nil
}

// createStageFile stages the src configuration file by processing the src
// template and setting the desired owner, group, and mode. It also sets the
// StageFile for the template resource.
// It returns an error if any.
func (t *Resource) createStageFileAndSync() error {
	for _, s := range t.sources {
		if !fileutil.IsFileExist(s.Src) {
			return errors.New("Missing template: " + s.Src)
		}

		log.WithFields(logrus.Fields{
			"template": s.Src,
		}).Debug("Compiling source template")

		tmpl, err := pongo2.FromFile(s.Src)
		if err != nil {
			return fmt.Errorf("Unable to process template %s, %s", s.Src, err)
		}

		// create TempFile in Dest directory to avoid cross-filesystem issues
		temp, err := ioutil.TempFile(filepath.Dir(s.Dst), "."+filepath.Base(s.Dst))
		if err != nil {
			return err
		}

		if err = tmpl.ExecuteWriter(t.funcMap, temp); err != nil {
			temp.Close()
			os.Remove(temp.Name())
			return err
		}

		defer temp.Close()

		// Set the owner, group, and mode on the stage file now to make it easier to
		// compare against the destination configuration file later.
		os.Chmod(temp.Name(), s.fileMode)
		os.Chown(temp.Name(), s.UID, s.GID)
		s.stageFile = temp

		if err = t.sync(s); err != nil {
			return err
		}
	}
	return nil
}

// sync compares the staged and dest config files and attempts to sync them
// if they differ. sync will run a config check command if set before
// overwriting the target config file. Finally, sync will run a reload command
// if set to have the application or service pick up the changes.
// It returns an error if any.
func (t *Resource) sync(s *SrcDst) error {
	staged := s.stageFile.Name()
	defer os.Remove(staged)

	log.WithFields(logrus.Fields{
		"staged": path.Base(staged),
		"dest":   s.Dst,
	}).Debug("Comparing staged and dest config files")

	ok, err := fileutil.SameFile(staged, s.Dst)
	if err != nil {
		log.Error(err.Error())
	}

	if !ok {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config out of sync")

		if s.CheckCmd != "" {
			if err := s.check(staged); err != nil {
				return errors.New("Config check failed: " + err.Error())
			}
		}
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Overwriting target config")
		err := os.Rename(staged, s.Dst)
		if err != nil {
			if strings.Contains(err.Error(), "device or resource busy") {
				log.Debug("Rename failed - target is likely a mount. Trying to write instead")
				// try to open the file and write to it
				var contents []byte
				var rerr error
				contents, rerr = ioutil.ReadFile(staged)
				if rerr != nil {
					return rerr
				}
				err := ioutil.WriteFile(s.Dst, contents, s.fileMode)
				// make sure owner and group match the temp file, in case the file was created with WriteFile
				os.Chown(s.Dst, s.UID, s.GID)
				if err != nil {
					return err
				}
			} else {
				return err
			}
		}
		if s.ReloadCmd != "" {
			if err := s.reload(); err != nil {
				return err
			}
		}

		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config has been updated")

	} else {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Target config in sync")

	}
	return nil
}

// setFileMode sets the FileMode.
func (t *Resource) setFileMode() error {
	for _, s := range t.sources {
		if err := s.setFileMode(); err != nil {
			return nil
		}
	}
	return nil
}

// Process is a convenience function that wraps calls to the three main tasks
// required to keep local configuration files in sync. First we gather vars
// from the store, then we stage a candidate configuration file, and finally sync
// things up.
// It returns an error if any.
func (t *Resource) process(storeClient StoreConfig) error {
	if err := t.setFileMode(); err != nil {
		return err
	}
	if err := t.setVars(storeClient); err != nil {
		return err
	}
	if err := t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

func (t *Resource) processAll() error {
	//load all KV-pairs the first time
	var err error
	if err = t.setFileMode(); err != nil {
		return err
	}
	for _, storeClient := range t.storeClients {
		if err := t.setVars(storeClient); err != nil {
			return err
		}
	}
	if err = t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

func (s StoreConfig) watch(stopChan chan bool, processChan chan StoreConfig) {
	var lastIndex uint64
	keysPrefix := appendPrefix(s.Prefix, s.Keys)

	for {
		select {
		case <-stopChan:
			return
		default:
			index, err := s.WatchPrefix(s.Prefix, keysPrefix, lastIndex, stopChan)
			if err != nil {
				log.Error(err)
				// Prevent backend errors from consuming all resources.
				time.Sleep(time.Second * 2)
				continue
			}
			processChan <- s
			lastIndex = index
		}
	}
}

func (s StoreConfig) interval(stopChan chan bool, processChan chan StoreConfig) {
	if s.Onetime {
		return
	}
	for {
		select {
		case <-stopChan:
			return
		case <-time.After(time.Duration(s.Interval) * time.Second):
			processChan <- s
		}
	}
}

func (t *Resource) Monitor(stopChan chan bool) {
	wg := &sync.WaitGroup{}
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	done := make(chan bool)
	stopIntervalWatch := make(chan bool)
	processChan := make(chan StoreConfig)

	defer close(processChan)

	if err := t.processAll(); err != nil {
		log.Error(err)
	}

	for _, sc := range t.storeClients {
		wg.Add(1)
		if sc.Watch {
			go func(s StoreConfig) {
				defer wg.Done()
				//t.watch(s, stopChan, processChan)
				s.watch(stopIntervalWatch, processChan)
			}(sc)
		} else {
			go func(s StoreConfig) {
				defer wg.Done()
				s.interval(stopIntervalWatch, processChan)
			}(sc)
		}
	}

	go func() {
		// If there is no goroutine left - quit
		wg.Wait()
		close(done)
	}()

	for {
		select {
		case storeClient := <-processChan:
			if err := t.process(storeClient); err != nil {
				log.Error(err)
			}
		case <-stopChan:
			close(stopIntervalWatch)
			go func() {
				for range processChan {
				}
			}()
			wg.Wait()
			return
		case <-done:
			return
		}
	}
}
<commit_msg>handle intervals <= 0<commit_after>package template

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/HeavyHorst/memkv"
	"github.com/HeavyHorst/remco/backends"
	"github.com/HeavyHorst/remco/log"
	"github.com/HeavyHorst/remco/template/fileutil"
	"github.com/Sirupsen/logrus"
	"github.com/flosch/pongo2"
	flag "github.com/spf13/pflag"
)

type StoreConfig struct {
	backends.StoreClient
	Name     string
	Onetime  bool
	Watch    bool
	Prefix   string
	Interval int
	Keys     []string
	store    *memkv.Store
}

// Resource is the representation of a parsed template resource.
type Resource struct {
	storeClients []StoreConfig
	funcMap      map[string]interface{}
	store        memkv.Store
	sources      []*SrcDst
}

var ErrEmptySrc = errors.New("empty src template")

// NewResource creates a Resource.
func NewResource(storeClients []StoreConfig, sources []*SrcDst) (*Resource, error) {
	if len(storeClients) == 0 {
		return nil, errors.New("A valid StoreClient is required.")
	}

	for _, v := range sources {
		if v.Src == "" {
			return nil, ErrEmptySrc
		}
	}

	for i := range storeClients {
		if storeClients[i].Interval <= 0 {
			log.Warning("Interval needs to be > 0: setting interval to 60")
			storeClients[i].Interval = 60
		}
	}

	tr := &Resource{
		storeClients: storeClients,
		store:        memkv.New(),
		funcMap:      newFuncMap(),
		sources:      sources,
	}

	// initialize the inidividual backend memkv Stores
	for i := range tr.storeClients {
		store := memkv.New()
		tr.storeClients[i].store = &store
	}

	addFuncs(tr.funcMap, tr.store.FuncMap)

	return tr, nil
}

// NewResourceFromFlags creates a Resource from the provided commandline flags.
func NewResourceFromFlags(s backends.Store, flags *flag.FlagSet, watch bool) (*Resource, error) {
	src, _ := flags.GetString("src")
	dst, _ := flags.GetString("dst")
	keys, _ := flags.GetStringSlice("keys")
	fileMode, _ := flags.GetString("fileMode")
	prefix, _ := flags.GetString("prefix")
	reloadCmd, _ := flags.GetString("reload_cmd")
	checkCmd, _ := flags.GetString("check_cmd")
	onetime, _ := flags.GetBool("onetime")
	interval, _ := flags.GetInt("interval")

	UID := os.Geteuid()
	GID := os.Getegid()

	sd := &SrcDst{
		Src:       src,
		Dst:       dst,
		Mode:      fileMode,
		ReloadCmd: reloadCmd,
		CheckCmd:  checkCmd,
		GID:       GID,
		UID:       UID,
	}

	b := StoreConfig{
		StoreClient: s.Client,
		Name:        s.Name,
		Onetime:     onetime,
		Prefix:      prefix,
		Watch:       watch,
		Interval:    interval,
		Keys:        keys,
	}

	return NewResource([]StoreConfig{b}, []*SrcDst{sd})
}

// setVars sets the Vars for template resource.
func (t *Resource) setVars(storeClient StoreConfig) error {
	var err error

	log.WithFields(logrus.Fields{
		"store":      storeClient.Name,
		"key_prefix": storeClient.Prefix,
	}).Debug("Retrieving keys")

	result, err := storeClient.GetValues(appendPrefix(storeClient.Prefix, storeClient.Keys))
	if err != nil {
		return err
	}

	storeClient.store.Purge()

	for key, value := range result {
		storeClient.store.Set(path.Join("/", strings.TrimPrefix(key, storeClient.Prefix)), value)
	}

	//merge all stores
	t.store.Purge()
	for _, v := range t.storeClients {
		for _, kv := range v.store.GetAllKVs() {
			if t.store.Exists(kv.Key) {
				log.Warning("Key collision - " + kv.Key)
			}
			t.store.Set(kv.Key, kv.Value)
		}
	}

	return nil
}

// createStageFile stages the src configuration file by processing the src
// template and setting the desired owner, group, and mode. It also sets the
// StageFile for the template resource.
// It returns an error if any.
func (t *Resource) createStageFileAndSync() error {
	for _, s := range t.sources {
		if !fileutil.IsFileExist(s.Src) {
			return errors.New("Missing template: " + s.Src)
		}

		log.WithFields(logrus.Fields{
			"template": s.Src,
		}).Debug("Compiling source template")

		tmpl, err := pongo2.FromFile(s.Src)
		if err != nil {
			return fmt.Errorf("Unable to process template %s, %s", s.Src, err)
		}

		// create TempFile in Dest directory to avoid cross-filesystem issues
		temp, err := ioutil.TempFile(filepath.Dir(s.Dst), "."+filepath.Base(s.Dst))
		if err != nil {
			return err
		}

		if err = tmpl.ExecuteWriter(t.funcMap, temp); err != nil {
			temp.Close()
			os.Remove(temp.Name())
			return err
		}

		defer temp.Close()

		// Set the owner, group, and mode on the stage file now to make it easier to
		// compare against the destination configuration file later.
		os.Chmod(temp.Name(), s.fileMode)
		os.Chown(temp.Name(), s.UID, s.GID)
		s.stageFile = temp

		if err = t.sync(s); err != nil {
			return err
		}
	}
	return nil
}

// sync compares the staged and dest config files and attempts to sync them
// if they differ. sync will run a config check command if set before
// overwriting the target config file. Finally, sync will run a reload command
// if set to have the application or service pick up the changes.
// It returns an error if any.
func (t *Resource) sync(s *SrcDst) error {
	staged := s.stageFile.Name()
	defer os.Remove(staged)

	log.WithFields(logrus.Fields{
		"staged": path.Base(staged),
		"dest":   s.Dst,
	}).Debug("Comparing staged and dest config files")

	ok, err := fileutil.SameFile(staged, s.Dst)
	if err != nil {
		log.Error(err.Error())
	}

	if !ok {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config out of sync")

		if s.CheckCmd != "" {
			if err := s.check(staged); err != nil {
				return errors.New("Config check failed: " + err.Error())
			}
		}
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Overwriting target config")
		err := os.Rename(staged, s.Dst)
		if err != nil {
			if strings.Contains(err.Error(), "device or resource busy") {
				log.Debug("Rename failed - target is likely a mount. Trying to write instead")
				// try to open the file and write to it
				var contents []byte
				var rerr error
				contents, rerr = ioutil.ReadFile(staged)
				if rerr != nil {
					return rerr
				}
				err := ioutil.WriteFile(s.Dst, contents, s.fileMode)
				// make sure owner and group match the temp file, in case the file was created with WriteFile
				os.Chown(s.Dst, s.UID, s.GID)
				if err != nil {
					return err
				}
			} else {
				return err
			}
		}
		if s.ReloadCmd != "" {
			if err := s.reload(); err != nil {
				return err
			}
		}

		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config has been updated")

	} else {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Target config in sync")

	}
	return nil
}

// setFileMode sets the FileMode.
func (t *Resource) setFileMode() error {
	for _, s := range t.sources {
		if err := s.setFileMode(); err != nil {
			return nil
		}
	}
	return nil
}

// Process is a convenience function that wraps calls to the three main tasks
// required to keep local configuration files in sync. First we gather vars
// from the store, then we stage a candidate configuration file, and finally sync
// things up.
// It returns an error if any.
func (t *Resource) process(storeClient StoreConfig) error {
	if err := t.setFileMode(); err != nil {
		return err
	}
	if err := t.setVars(storeClient); err != nil {
		return err
	}
	if err := t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

func (t *Resource) processAll() error {
	//load all KV-pairs the first time
	var err error
	if err = t.setFileMode(); err != nil {
		return err
	}
	for _, storeClient := range t.storeClients {
		if err := t.setVars(storeClient); err != nil {
			return err
		}
	}
	if err = t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

func (s StoreConfig) watch(stopChan chan bool, processChan chan StoreConfig) {
	var lastIndex uint64
	keysPrefix := appendPrefix(s.Prefix, s.Keys)

	for {
		select {
		case <-stopChan:
			return
		default:
			index, err := s.WatchPrefix(s.Prefix, keysPrefix, lastIndex, stopChan)
			if err != nil {
				log.Error(err)
				// Prevent backend errors from consuming all resources.
				time.Sleep(time.Second * 2)
				continue
			}
			processChan <- s
			lastIndex = index
		}
	}
}

func (s StoreConfig) interval(stopChan chan bool, processChan chan StoreConfig) {
	if s.Onetime {
		return
	}
	for {
		select {
		case <-stopChan:
			return
		case <-time.After(time.Duration(s.Interval) * time.Second):
			processChan <- s
		}
	}
}

func (t *Resource) Monitor(stopChan chan bool) {
	wg := &sync.WaitGroup{}
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	done := make(chan bool)
	stopIntervalWatch := make(chan bool)
	processChan := make(chan StoreConfig)

	defer close(processChan)

	if err := t.processAll(); err != nil {
		log.Error(err)
	}

	for _, sc := range t.storeClients {
		wg.Add(1)
		if sc.Watch {
			go func(s StoreConfig) {
				defer wg.Done()
				//t.watch(s, stopChan, processChan)
				s.watch(stopIntervalWatch, processChan)
			}(sc)
		} else {
			go func(s StoreConfig) {
				defer wg.Done()
				s.interval(stopIntervalWatch, processChan)
			}(sc)
		}
	}

	go func() {
		// If there is no goroutine left - quit
		wg.Wait()
		close(done)
	}()

	for {
		select {
		case storeClient := <-processChan:
			if err := t.process(storeClient); err != nil {
				log.Error(err)
			}
		case <-stopChan:
			close(stopIntervalWatch)
			go func() {
				for range processChan {
				}
			}()
			wg.Wait()
			return
		case <-done:
			return
		}
	}
}
<|endoftext|>


<commit_before>/*
 * This file is part of remco.
 * Based on code from confd. https://github.com/kelseyhightower/confd
 * © 2013 Kelsey Hightower
 * © 2016 The Remco Authors
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

package template

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/HeavyHorst/easyKV"
	"github.com/HeavyHorst/memkv"
	"github.com/HeavyHorst/pongo2"
	berr "github.com/HeavyHorst/remco/backends/error"
	"github.com/HeavyHorst/remco/log"
	"github.com/HeavyHorst/remco/template/fileutil"
	"github.com/Sirupsen/logrus"
)

// A BackendConfig - Every backend implements this interface. If Connect is called a new connection to the underlaying kv-store will be established.
// Connect should also set the name and the StoreClient of the Backend. The other values of Backend will be loaded from the configuration file.
type BackendConfig interface {
	Connect() (Backend, error)
}

// Backend is the representation of a template backend like etcd or consul
type Backend struct {
	easyKV.ReadWatcher
	Name     string
	Onetime  bool
	Watch    bool
	Prefix   string
	Interval int
	Keys     []string
	store    *memkv.Store
}

type Backends []Backend

// Close is calling the close method on all backends
func (b Backends) Close() {
	for _, v := range b {
		log.WithFields(logrus.Fields{
			"backend": v.Name,
		}).Debug("Closing client connection")
		v.ReadWatcher.Close()
	}
}

// Resource is the representation of a parsed template resource.
type Resource struct {
	backends []Backend
	funcMap  map[string]interface{}
	store    memkv.Store
	sources  []*ProcessConfig
}

// ErrEmptySrc is returned if an emty src template is passed to NewResource
var ErrEmptySrc = errors.New("empty src template")

// NewResource creates a Resource.
func NewResource(backends Backends, sources []*ProcessConfig) (*Resource, error) {
	if len(backends) == 0 {
		return nil, errors.New("A valid StoreClient is required.")
	}

	for _, v := range sources {
		if v.Src == "" {
			return nil, ErrEmptySrc
		}
	}

	tr := &Resource{
		backends: backends,
		store:    memkv.New(),
		funcMap:  newFuncMap(),
		sources:  sources,
	}

	// initialize the inidividual backend memkv Stores
	for i := range tr.backends {
		store := memkv.New()
		tr.backends[i].store = &store

		if tr.backends[i].Interval <= 0 && tr.backends[i].Onetime == false && tr.backends[i].Watch == false {
			log.Warning("Interval needs to be > 0: setting interval to 60")
			tr.backends[i].Interval = 60
		}
	}

	addFuncs(tr.funcMap, tr.store.FuncMap)

	return tr, nil
}

// setVars sets the Vars for a template resource.
func (t *Resource) setVars(storeClient Backend) error {
	var err error

	log.WithFields(logrus.Fields{
		"backend":    storeClient.Name,
		"key_prefix": storeClient.Prefix,
	}).Debug("Retrieving keys")

	result, err := storeClient.GetValues(appendPrefix(storeClient.Prefix, storeClient.Keys))
	if err != nil {
		return err
	}

	storeClient.store.Purge()

	for key, value := range result {
		storeClient.store.Set(path.Join("/", strings.TrimPrefix(key, storeClient.Prefix)), value)
	}

	//merge all stores
	t.store.Purge()
	for _, v := range t.backends {
		for _, kv := range v.store.GetAllKVs() {
			if t.store.Exists(kv.Key) {
				log.Warning("Key collision - " + kv.Key)
			}
			t.store.Set(kv.Key, kv.Value)
		}
	}

	return nil
}

// createStageFile stages the src configuration file by processing the src
// template and setting the desired owner, group, and mode. It also sets the
// StageFile for the template resource.
// It returns an error if any.
func (t *Resource) createStageFileAndSync() error {
	for _, s := range t.sources {
		if !fileutil.IsFileExist(s.Src) {
			return errors.New("Missing template: " + s.Src)
		}

		log.WithFields(logrus.Fields{
			"template": s.Src,
		}).Debug("Compiling source template")

		set := pongo2.NewSet("local", &pongo2.LocalFilesystemLoader{})
		tmpl, err := set.FromFile(s.Src)
		if err != nil {
			return fmt.Errorf("Unable to process template %s, %s", s.Src, err)
		}

		// create TempFile in Dest directory to avoid cross-filesystem issues
		temp, err := ioutil.TempFile(filepath.Dir(s.Dst), "."+filepath.Base(s.Dst))
		if err != nil {
			return err
		}

		if err = tmpl.ExecuteWriter(t.funcMap, temp); err != nil {
			temp.Close()
			os.Remove(temp.Name())
			return err
		}

		temp.Close()

		fileMode, err := s.getFileMode()
		if err != nil {
			return err
		}
		// Set the owner, group, and mode on the stage file now to make it easier to
		// compare against the destination configuration file later.
		os.Chmod(temp.Name(), fileMode)
		os.Chown(temp.Name(), s.UID, s.GID)
		s.stageFile = temp

		if err = t.sync(s); err != nil {
			return err
		}
	}
	return nil
}

// sync compares the staged and dest config files and attempts to sync them
// if they differ. sync will run a config check command if set before
// overwriting the target config file. Finally, sync will run a reload command
// if set to have the application or service pick up the changes.
// It returns an error if any.
func (t *Resource) sync(s *ProcessConfig) error {
	staged := s.stageFile.Name()
	defer os.Remove(staged)

	log.WithFields(logrus.Fields{
		"staged": path.Base(staged),
		"dest":   s.Dst,
	}).Debug("Comparing staged and dest config files")

	ok, err := fileutil.SameFile(staged, s.Dst)
	if err != nil {
		log.Error(err.Error())
	}

	if !ok {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config out of sync")

		if err := s.check(staged); err != nil {
			return errors.New("Config check failed: " + err.Error())
		}

		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Overwriting target config")

		fileMode, err := s.getFileMode()
		if err != nil {
			return err
		}
		if err := fileutil.ReplaceFile(staged, s.Dst, fileMode); err != nil {
			return err
		}

		// make sure owner and group match the temp file, in case the file was created with WriteFile
		os.Chown(s.Dst, s.UID, s.GID)

		if err := s.reload(); err != nil {
			return err
		}

		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config has been updated")

	} else {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Target config in sync")

	}
	return nil
}

// Process is a convenience function that wraps calls to the three main tasks
// required to keep local configuration files in sync. First we gather vars
// from the store, then we stage a candidate configuration file, and finally sync
// things up.
// It returns an error if any.
func (t *Resource) process(storeClients []Backend) error {
	for _, storeClient := range storeClients {
		if err := t.setVars(storeClient); err != nil {
			return berr.BackendError{
				Message: err.Error(),
				Backend: storeClient.Name,
			}
		}
	}
	if err := t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

func (s Backend) watch(ctx context.Context, processChan chan Backend) {
	var lastIndex uint64
	keysPrefix := appendPrefix(s.Prefix, s.Keys)

	for {
		select {
		case <-ctx.Done():
			return
		default:
			index, err := s.WatchPrefix(s.Prefix, ctx, easyKV.WithKeys(keysPrefix), easyKV.WithWaitIndex(lastIndex))
			if err != nil {
				if err != easyKV.ErrWatchCanceled {
					log.WithFields(logrus.Fields{
						"backend": s.Name,
					}).Error(err)
					time.Sleep(2 * time.Second)
				}
				continue
			}
			processChan <- s
			lastIndex = index
		}
	}
}

func (s Backend) interval(ctx context.Context, processChan chan Backend) {
	if s.Onetime {
		return
	}
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(time.Duration(s.Interval) * time.Second):
			processChan <- s
		}
	}
}

// Monitor will start to monitor all given Backends for changes.
// It will process all given tamplates on changes.
func (t *Resource) Monitor(ctx context.Context) {
	wg := &sync.WaitGroup{}
	done := make(chan struct{})

	processChan := make(chan Backend)
	defer close(processChan)

	if err := t.process(t.backends); err != nil {
		switch err.(type) {
		case berr.BackendError:
			err := err.(berr.BackendError)
			log.WithFields(logrus.Fields{
				"backend": err.Backend,
			}).Error(err)
		default:
			log.Error(err)
		}
	}

	for _, sc := range t.backends {
		wg.Add(1)
		if sc.Watch {
			go func(s Backend) {
				defer wg.Done()
				s.watch(ctx, processChan)
			}(sc)
		} else {
			go func(s Backend) {
				defer wg.Done()
				s.interval(ctx, processChan)
			}(sc)
		}
	}

	go func() {
		// If there is no goroutine left - quit
		wg.Wait()
		close(done)
	}()

	for {
		select {
		case storeClient := <-processChan:
			if err := t.process([]Backend{storeClient}); err != nil {
				if _, ok := err.(berr.BackendError); ok {
					log.WithFields(logrus.Fields{
						"backend": storeClient.Name,
					}).Error(err)
				} else {
					log.Error(err)
				}
			}
		case <-ctx.Done():
			go func() {
				for range processChan {
				}
			}()
			wg.Wait()
			return
		case <-done:
			return
		}
	}
}
<commit_msg>watch and interval not mutual exclusive anymore<commit_after>/*
 * This file is part of remco.
 * Based on code from confd. https://github.com/kelseyhightower/confd
 * © 2013 Kelsey Hightower
 * © 2016 The Remco Authors
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

package template

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/HeavyHorst/easyKV"
	"github.com/HeavyHorst/memkv"
	"github.com/HeavyHorst/pongo2"
	berr "github.com/HeavyHorst/remco/backends/error"
	"github.com/HeavyHorst/remco/log"
	"github.com/HeavyHorst/remco/template/fileutil"
	"github.com/Sirupsen/logrus"
)

// A BackendConfig - Every backend implements this interface. If Connect is called a new connection to the underlaying kv-store will be established.
// Connect should also set the name and the StoreClient of the Backend. The other values of Backend will be loaded from the configuration file.
type BackendConfig interface {
	Connect() (Backend, error)
}

// Backend is the representation of a template backend like etcd or consul
type Backend struct {
	easyKV.ReadWatcher
	Name     string
	Onetime  bool
	Watch    bool
	Prefix   string
	Interval int
	Keys     []string
	store    *memkv.Store
}

type Backends []Backend

// Close is calling the close method on all backends
func (b Backends) Close() {
	for _, v := range b {
		log.WithFields(logrus.Fields{
			"backend": v.Name,
		}).Debug("Closing client connection")
		v.ReadWatcher.Close()
	}
}

// Resource is the representation of a parsed template resource.
type Resource struct {
	backends []Backend
	funcMap  map[string]interface{}
	store    memkv.Store
	sources  []*ProcessConfig
}

// ErrEmptySrc is returned if an emty src template is passed to NewResource
var ErrEmptySrc = errors.New("empty src template")

// NewResource creates a Resource.
func NewResource(backends Backends, sources []*ProcessConfig) (*Resource, error) {
	if len(backends) == 0 {
		return nil, errors.New("A valid StoreClient is required.")
	}

	for _, v := range sources {
		if v.Src == "" {
			return nil, ErrEmptySrc
		}
	}

	tr := &Resource{
		backends: backends,
		store:    memkv.New(),
		funcMap:  newFuncMap(),
		sources:  sources,
	}

	// initialize the inidividual backend memkv Stores
	for i := range tr.backends {
		store := memkv.New()
		tr.backends[i].store = &store

		if tr.backends[i].Interval <= 0 && tr.backends[i].Onetime == false && tr.backends[i].Watch == false {
			log.Warning("Interval needs to be > 0: setting interval to 60")
			tr.backends[i].Interval = 60
		}
	}

	addFuncs(tr.funcMap, tr.store.FuncMap)

	return tr, nil
}

// setVars sets the Vars for a template resource.
func (t *Resource) setVars(storeClient Backend) error {
	var err error

	log.WithFields(logrus.Fields{
		"backend":    storeClient.Name,
		"key_prefix": storeClient.Prefix,
	}).Debug("Retrieving keys")

	result, err := storeClient.GetValues(appendPrefix(storeClient.Prefix, storeClient.Keys))
	if err != nil {
		return err
	}

	storeClient.store.Purge()

	for key, value := range result {
		storeClient.store.Set(path.Join("/", strings.TrimPrefix(key, storeClient.Prefix)), value)
	}

	//merge all stores
	t.store.Purge()
	for _, v := range t.backends {
		for _, kv := range v.store.GetAllKVs() {
			if t.store.Exists(kv.Key) {
				log.Warning("Key collision - " + kv.Key)
			}
			t.store.Set(kv.Key, kv.Value)
		}
	}

	return nil
}

// createStageFile stages the src configuration file by processing the src
// template and setting the desired owner, group, and mode. It also sets the
// StageFile for the template resource.
// It returns an error if any.
func (t *Resource) createStageFileAndSync() error {
	for _, s := range t.sources {
		if !fileutil.IsFileExist(s.Src) {
			return errors.New("Missing template: " + s.Src)
		}

		log.WithFields(logrus.Fields{
			"template": s.Src,
		}).Debug("Compiling source template")

		set := pongo2.NewSet("local", &pongo2.LocalFilesystemLoader{})
		tmpl, err := set.FromFile(s.Src)
		if err != nil {
			return fmt.Errorf("Unable to process template %s, %s", s.Src, err)
		}

		// create TempFile in Dest directory to avoid cross-filesystem issues
		temp, err := ioutil.TempFile(filepath.Dir(s.Dst), "."+filepath.Base(s.Dst))
		if err != nil {
			return err
		}

		if err = tmpl.ExecuteWriter(t.funcMap, temp); err != nil {
			temp.Close()
			os.Remove(temp.Name())
			return err
		}

		temp.Close()

		fileMode, err := s.getFileMode()
		if err != nil {
			return err
		}
		// Set the owner, group, and mode on the stage file now to make it easier to
		// compare against the destination configuration file later.
		os.Chmod(temp.Name(), fileMode)
		os.Chown(temp.Name(), s.UID, s.GID)
		s.stageFile = temp

		if err = t.sync(s); err != nil {
			return err
		}
	}
	return nil
}

// sync compares the staged and dest config files and attempts to sync them
// if they differ. sync will run a config check command if set before
// overwriting the target config file. Finally, sync will run a reload command
// if set to have the application or service pick up the changes.
// It returns an error if any.
func (t *Resource) sync(s *ProcessConfig) error {
	staged := s.stageFile.Name()
	defer os.Remove(staged)

	log.WithFields(logrus.Fields{
		"staged": path.Base(staged),
		"dest":   s.Dst,
	}).Debug("Comparing staged and dest config files")

	ok, err := fileutil.SameFile(staged, s.Dst)
	if err != nil {
		log.Error(err.Error())
	}

	if !ok {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config out of sync")

		if err := s.check(staged); err != nil {
			return errors.New("Config check failed: " + err.Error())
		}

		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Overwriting target config")

		fileMode, err := s.getFileMode()
		if err != nil {
			return err
		}
		if err := fileutil.ReplaceFile(staged, s.Dst, fileMode); err != nil {
			return err
		}

		// make sure owner and group match the temp file, in case the file was created with WriteFile
		os.Chown(s.Dst, s.UID, s.GID)

		if err := s.reload(); err != nil {
			return err
		}

		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Info("Target config has been updated")

	} else {
		log.WithFields(logrus.Fields{
			"config": s.Dst,
		}).Debug("Target config in sync")

	}
	return nil
}

// Process is a convenience function that wraps calls to the three main tasks
// required to keep local configuration files in sync. First we gather vars
// from the store, then we stage a candidate configuration file, and finally sync
// things up.
// It returns an error if any.
func (t *Resource) process(storeClients []Backend) error {
	for _, storeClient := range storeClients {
		if err := t.setVars(storeClient); err != nil {
			return berr.BackendError{
				Message: err.Error(),
				Backend: storeClient.Name,
			}
		}
	}
	if err := t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

func (s Backend) watch(ctx context.Context, processChan chan Backend) {
	if s.Onetime {
		return
	}

	var lastIndex uint64
	keysPrefix := appendPrefix(s.Prefix, s.Keys)

	for {
		select {
		case <-ctx.Done():
			return
		default:
			index, err := s.WatchPrefix(s.Prefix, ctx, easyKV.WithKeys(keysPrefix), easyKV.WithWaitIndex(lastIndex))
			if err != nil {
				if err != easyKV.ErrWatchCanceled {
					log.WithFields(logrus.Fields{
						"backend": s.Name,
					}).Error(err)
					time.Sleep(2 * time.Second)
				}
				continue
			}
			processChan <- s
			lastIndex = index
		}
	}
}

func (s Backend) interval(ctx context.Context, processChan chan Backend) {
	if s.Onetime {
		return
	}
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(time.Duration(s.Interval) * time.Second):
			processChan <- s
		}
	}
}

// Monitor will start to monitor all given Backends for changes.
// It will process all given tamplates on changes.
func (t *Resource) Monitor(ctx context.Context) {
	wg := &sync.WaitGroup{}

	processChan := make(chan Backend)
	defer close(processChan)

	if err := t.process(t.backends); err != nil {
		switch err.(type) {
		case berr.BackendError:
			err := err.(berr.BackendError)
			log.WithFields(logrus.Fields{
				"backend": err.Backend,
			}).Error(err)
		default:
			log.Error(err)
		}
	}

	for _, sc := range t.backends {
		if sc.Watch {
			wg.Add(1)
			go func(s Backend) {
				defer wg.Done()
				s.watch(ctx, processChan)
			}(sc)
		}

		if sc.Interval > 0 {
			wg.Add(1)
			go func(s Backend) {
				defer wg.Done()
				s.interval(ctx, processChan)
			}(sc)
		}
	}

	done := make(chan struct{})
	go func() {
		// If there is no goroutine left - quit
		wg.Wait()
		close(done)
	}()

	for {
		select {
		case storeClient := <-processChan:
			if err := t.process([]Backend{storeClient}); err != nil {
				if _, ok := err.(berr.BackendError); ok {
					log.WithFields(logrus.Fields{
						"backend": storeClient.Name,
					}).Error(err)
				} else {
					log.Error(err)
				}
			}
		case <-ctx.Done():
			go func() {
				for range processChan {
				}
			}()
			wg.Wait()
			return
		case <-done:
			return
		}
	}
}
<|endoftext|>


<commit_before>// +build e2e

/*
Copyright 2018 Knative Authors LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package test

import (
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/knative/build-pipeline/pkg/apis/pipeline/v1alpha1"
	clientset "github.com/knative/build-pipeline/pkg/client/clientset/versioned/typed/pipeline/v1alpha1"
	tb "github.com/knative/build-pipeline/test/builder"
	knativetest "github.com/knative/pkg/test"
	"github.com/knative/pkg/test/logging"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TestDAGPipelineRun creates a graph of arbitrary Tasks, then looks at the corresponding
// TaskRun start times to ensure they were run in the order intended, which is:
//                               |
//                        pipeline-task-1
//                       /               \
//   pipeline-task-2-parallel-1    pipeline-task-2-parallel-2
//                       \                /
//                        pipeline-task-3
//                               |
//                        pipeline-task-4
func TestDAGPipelineRun(t *testing.T) {
	logger := logging.GetContextLogger(t.Name())
	c, namespace := setup(t, logger)

	knativetest.CleanupOnInterrupt(func() { tearDown(t, logger, c, namespace) }, logger)
	defer tearDown(t, logger, c, namespace)

	// Create the Task that echoes text
	echoTask := tb.Task("echo-task", namespace, tb.TaskSpec(
		tb.TaskInputs(
			tb.InputsResource("repo", v1alpha1.PipelineResourceTypeGit),
			tb.InputsParam("text", tb.ParamDescription("The text that should be echoed")),
		),
		tb.TaskOutputs(tb.OutputsResource("repo", v1alpha1.PipelineResourceTypeGit)),
		tb.Step("echo-text", "busybox", tb.Command("echo"), tb.Args("${inputs.params.text}")),
	))
	if _, err := c.TaskClient.Create(echoTask); err != nil {
		t.Fatalf("Failed to create echo Task: %s", err)
	}

	// Create the repo PipelineResource (doesn't really matter which repo we use)
	repoResource := tb.PipelineResource("repo", namespace, tb.PipelineResourceSpec(
		v1alpha1.PipelineResourceTypeGit,
		tb.PipelineResourceSpecParam("Url", "https://github.com/githubtraining/example-basic"),
	))
	if _, err := c.PipelineResourceClient.Create(repoResource); err != nil {
		t.Fatalf("Failed to create simple repo PipelineResource: %s", err)
	}

	// Intentionally delcaring Tasks in a mixed up order to ensure the order
	// of execution isn't at all dependent on the order they are declared in
	pipeline := tb.Pipeline("dag-pipeline", namespace, tb.PipelineSpec(
		tb.PipelineDeclaredResource("repo", "repo"),
		tb.PipelineTask("pipeline-task-3", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo", tb.From("pipeline-task-2-parallel-1", "pipeline-task-2-parallel-2")),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "wow"),
		),
		tb.PipelineTask("pipeline-task-2-parallel-2", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo", tb.From("pipeline-task-1")), tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "such parallel"),
		),
		tb.PipelineTask("pipeline-task-4", "echo-task",
			tb.RunAfter("pipeline-task-3"),
			tb.PipelineTaskInputResource("repo", "repo"),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "very cloud native"),
		),
		tb.PipelineTask("pipeline-task-2-parallel-1", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo", tb.From("pipeline-task-1")),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "much graph"),
		),
		tb.PipelineTask("pipeline-task-1", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo"),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "how to ci/cd?"),
		),
	))
	if _, err := c.PipelineClient.Create(pipeline); err != nil {
		t.Fatalf("Failed to create dag-pipeline: %s", err)
	}
	pipelineRun := tb.PipelineRun("dag-pipeline-run", namespace, tb.PipelineRunSpec("dag-pipeline",
		tb.PipelineRunResourceBinding("repo", tb.PipelineResourceBindingRef("repo")),
	))
	if _, err := c.PipelineRunClient.Create(pipelineRun); err != nil {
		t.Fatalf("Failed to create dag-pipeline-run PipelineRun: %s", err)
	}
	logger.Infof("Waiting for DAG pipeline to complete")
	if err := WaitForPipelineRunState(c, "dag-pipeline-run", pipelineRunTimeout, PipelineRunSucceed("dag-pipeline-run"), "PipelineRunSuccess"); err != nil {
		t.Fatalf("Error waiting for PipelineRun to finish: %s", err)
	}

	logger.Infof("Verifying order of execution")
	times := getTaskStartTimes(t, c.TaskRunClient)
	vefifyExpectedOrder(t, times)
}

type runTime struct {
	name string
	t    time.Time
}

type runTimes []runTime

func (f runTimes) Len() int {
	return len(f)
}

func (f runTimes) Less(i, j int) bool {
	return f[i].t.Before(f[j].t)
}

func (f runTimes) Swap(i, j int) {
	f[i], f[j] = f[j], f[i]
}

func getTaskStartTimes(t *testing.T, c clientset.TaskRunInterface) runTimes {
	taskRuns, err := c.List(metav1.ListOptions{})
	if err != nil {
		t.Fatalf("Couldn't get TaskRuns (so that we could check when they executed): %v", err)
	}
	times := runTimes{}
	for _, t := range taskRuns.Items {
		times = append(times, runTime{
			name: t.Name,
			t:    t.Status.StartTime.Time,
		})
	}
	return times
}

func vefifyExpectedOrder(t *testing.T, times runTimes) {
	if len(times) != 5 {
		t.Fatalf("Expected 5 Taskruns to have executed but only got start times for %d Taskruns", len(times))
	}

	sort.Sort(times)

	if !strings.HasPrefix(times[0].name, "dag-pipeline-run-pipeline-task-1") {
		t.Errorf("Expected first task to execute first, but %q was first", times[0].name)
	}
	if !strings.HasPrefix(times[1].name, "dag-pipeline-run-pipeline-task-2") {
		t.Errorf("Expected parallel tasks to run second & third, but %q was second", times[1].name)
	}
	if !strings.HasPrefix(times[2].name, "dag-pipeline-run-pipeline-task-2") {
		t.Errorf("Expected parallel tasks to run second & third, but %q was third", times[2].name)
	}
	if !strings.HasPrefix(times[3].name, "dag-pipeline-run-pipeline-task-3") {
		t.Errorf("Expected third task to execute third, but %q was third", times[3].name)
	}
	if !strings.HasPrefix(times[4].name, "dag-pipeline-run-pipeline-task-4") {
		t.Errorf("Expected fourth task to execute fourth, but %q was fourth", times[4].name)
	}

	// Check that the two tasks that can run in parallel did
	parallelDiff := times[2].t.Sub(times[1].t)
	if parallelDiff > (time.Second * 5) {
		t.Errorf("Expected parallel tasks to execute more or less at the same time, but they were %v apart", parallelDiff)
	}
}
<commit_msg>Fix spelling errors<commit_after>// +build e2e

/*
Copyright 2018 Knative Authors LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package test

import (
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/knative/build-pipeline/pkg/apis/pipeline/v1alpha1"
	clientset "github.com/knative/build-pipeline/pkg/client/clientset/versioned/typed/pipeline/v1alpha1"
	tb "github.com/knative/build-pipeline/test/builder"
	knativetest "github.com/knative/pkg/test"
	"github.com/knative/pkg/test/logging"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TestDAGPipelineRun creates a graph of arbitrary Tasks, then looks at the corresponding
// TaskRun start times to ensure they were run in the order intended, which is:
//                               |
//                        pipeline-task-1
//                       /               \
//   pipeline-task-2-parallel-1    pipeline-task-2-parallel-2
//                       \                /
//                        pipeline-task-3
//                               |
//                        pipeline-task-4
func TestDAGPipelineRun(t *testing.T) {
	logger := logging.GetContextLogger(t.Name())
	c, namespace := setup(t, logger)

	knativetest.CleanupOnInterrupt(func() { tearDown(t, logger, c, namespace) }, logger)
	defer tearDown(t, logger, c, namespace)

	// Create the Task that echoes text
	echoTask := tb.Task("echo-task", namespace, tb.TaskSpec(
		tb.TaskInputs(
			tb.InputsResource("repo", v1alpha1.PipelineResourceTypeGit),
			tb.InputsParam("text", tb.ParamDescription("The text that should be echoed")),
		),
		tb.TaskOutputs(tb.OutputsResource("repo", v1alpha1.PipelineResourceTypeGit)),
		tb.Step("echo-text", "busybox", tb.Command("echo"), tb.Args("${inputs.params.text}")),
	))
	if _, err := c.TaskClient.Create(echoTask); err != nil {
		t.Fatalf("Failed to create echo Task: %s", err)
	}

	// Create the repo PipelineResource (doesn't really matter which repo we use)
	repoResource := tb.PipelineResource("repo", namespace, tb.PipelineResourceSpec(
		v1alpha1.PipelineResourceTypeGit,
		tb.PipelineResourceSpecParam("Url", "https://github.com/githubtraining/example-basic"),
	))
	if _, err := c.PipelineResourceClient.Create(repoResource); err != nil {
		t.Fatalf("Failed to create simple repo PipelineResource: %s", err)
	}

	// Intentionally declaring Tasks in a mixed up order to ensure the order
	// of execution isn't at all dependent on the order they are declared in
	pipeline := tb.Pipeline("dag-pipeline", namespace, tb.PipelineSpec(
		tb.PipelineDeclaredResource("repo", "repo"),
		tb.PipelineTask("pipeline-task-3", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo", tb.From("pipeline-task-2-parallel-1", "pipeline-task-2-parallel-2")),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "wow"),
		),
		tb.PipelineTask("pipeline-task-2-parallel-2", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo", tb.From("pipeline-task-1")), tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "such parallel"),
		),
		tb.PipelineTask("pipeline-task-4", "echo-task",
			tb.RunAfter("pipeline-task-3"),
			tb.PipelineTaskInputResource("repo", "repo"),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "very cloud native"),
		),
		tb.PipelineTask("pipeline-task-2-parallel-1", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo", tb.From("pipeline-task-1")),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "much graph"),
		),
		tb.PipelineTask("pipeline-task-1", "echo-task",
			tb.PipelineTaskInputResource("repo", "repo"),
			tb.PipelineTaskOutputResource("repo", "repo"),
			tb.PipelineTaskParam("text", "how to ci/cd?"),
		),
	))
	if _, err := c.PipelineClient.Create(pipeline); err != nil {
		t.Fatalf("Failed to create dag-pipeline: %s", err)
	}
	pipelineRun := tb.PipelineRun("dag-pipeline-run", namespace, tb.PipelineRunSpec("dag-pipeline",
		tb.PipelineRunResourceBinding("repo", tb.PipelineResourceBindingRef("repo")),
	))
	if _, err := c.PipelineRunClient.Create(pipelineRun); err != nil {
		t.Fatalf("Failed to create dag-pipeline-run PipelineRun: %s", err)
	}
	logger.Infof("Waiting for DAG pipeline to complete")
	if err := WaitForPipelineRunState(c, "dag-pipeline-run", pipelineRunTimeout, PipelineRunSucceed("dag-pipeline-run"), "PipelineRunSuccess"); err != nil {
		t.Fatalf("Error waiting for PipelineRun to finish: %s", err)
	}

	logger.Infof("Verifying order of execution")
	times := getTaskStartTimes(t, c.TaskRunClient)
	vefifyExpectedOrder(t, times)
}

type runTime struct {
	name string
	t    time.Time
}

type runTimes []runTime

func (f runTimes) Len() int {
	return len(f)
}

func (f runTimes) Less(i, j int) bool {
	return f[i].t.Before(f[j].t)
}

func (f runTimes) Swap(i, j int) {
	f[i], f[j] = f[j], f[i]
}

func getTaskStartTimes(t *testing.T, c clientset.TaskRunInterface) runTimes {
	taskRuns, err := c.List(metav1.ListOptions{})
	if err != nil {
		t.Fatalf("Couldn't get TaskRuns (so that we could check when they executed): %v", err)
	}
	times := runTimes{}
	for _, t := range taskRuns.Items {
		times = append(times, runTime{
			name: t.Name,
			t:    t.Status.StartTime.Time,
		})
	}
	return times
}

func vefifyExpectedOrder(t *testing.T, times runTimes) {
	if len(times) != 5 {
		t.Fatalf("Expected 5 Taskruns to have executed but only got start times for %d Taskruns", len(times))
	}

	sort.Sort(times)

	if !strings.HasPrefix(times[0].name, "dag-pipeline-run-pipeline-task-1") {
		t.Errorf("Expected first task to execute first, but %q was first", times[0].name)
	}
	if !strings.HasPrefix(times[1].name, "dag-pipeline-run-pipeline-task-2") {
		t.Errorf("Expected parallel tasks to run second & third, but %q was second", times[1].name)
	}
	if !strings.HasPrefix(times[2].name, "dag-pipeline-run-pipeline-task-2") {
		t.Errorf("Expected parallel tasks to run second & third, but %q was third", times[2].name)
	}
	if !strings.HasPrefix(times[3].name, "dag-pipeline-run-pipeline-task-3") {
		t.Errorf("Expected third task to execute third, but %q was third", times[3].name)
	}
	if !strings.HasPrefix(times[4].name, "dag-pipeline-run-pipeline-task-4") {
		t.Errorf("Expected fourth task to execute fourth, but %q was fourth", times[4].name)
	}

	// Check that the two tasks that can run in parallel did
	parallelDiff := times[2].t.Sub(times[1].t)
	if parallelDiff > (time.Second * 5) {
		t.Errorf("Expected parallel tasks to execute more or less at the same time, but they were %v apart", parallelDiff)
	}
}
<|endoftext|>


<commit_before>// Package pubsub implements a pub-sub model with a single publisher (Server)
// and multiple subscribers (clients).
//
// Though you can have multiple publishers by sharing a pointer to a server or
// by giving the same channel to each publisher and publishing messages from
// that channel (fan-in).
//
// Clients subscribe for messages, which could be of any type, using a query.
// When some message is published, we match it with all queries. If there is a
// match, this message will be pushed to all clients, subscribed to that query.
// See query subpackage for our implementation.
package pubsub

import (
	"context"

	cmn "github.com/tendermint/tmlibs/common"
)

type operation int

const (
	sub operation = iota
	pub
	unsub
	shutdown
)

type cmd struct {
	op       operation
	query    Query
	ch       chan<- interface{}
	clientID string
	msg      interface{}
	tags     map[string]interface{}
}

// Query defines an interface for a query to be used for subscribing.
type Query interface {
	Matches(tags map[string]interface{}) bool
}

// Server allows clients to subscribe/unsubscribe for messages, publishing
// messages with or without tags, and manages internal state.
type Server struct {
	cmn.BaseService

	cmds    chan cmd
	cmdsCap int
}

// Option sets a parameter for the server.
type Option func(*Server)

// NewServer returns a new server. See the commentary on the Option functions
// for a detailed description of how to configure buffering. If no options are
// provided, the resulting server's queue is unbuffered.
func NewServer(options ...Option) *Server {
	s := &Server{}
	s.BaseService = *cmn.NewBaseService(nil, "PubSub", s)

	for _, option := range options {
		option(s)
	}

	// if BufferCapacity option was not set, the channel is unbuffered
	s.cmds = make(chan cmd, s.cmdsCap)

	return s
}

// BufferCapacity allows you to specify capacity for the internal server's
// queue. Since the server, given Y subscribers, could only process X messages,
// this option could be used to survive spikes (e.g. high amount of
// transactions during peak hours).
func BufferCapacity(cap int) Option {
	return func(s *Server) {
		if cap > 0 {
			s.cmdsCap = cap
		}
	}
}

// BufferCapacity returns capacity of the internal server's queue.
func (s Server) BufferCapacity() int {
	return s.cmdsCap
}

// Subscribe creates a subscription for the given client. It accepts a channel
// on which messages matching the given query can be received. If the
// subscription already exists, the old channel will be closed. An error will
// be returned to the caller if the context is canceled.
func (s *Server) Subscribe(ctx context.Context, clientID string, query Query, out chan<- interface{}) error {
	select {
	case s.cmds <- cmd{op: sub, clientID: clientID, query: query, ch: out}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Unsubscribe removes the subscription on the given query. An error will be
// returned to the caller if the context is canceled.
func (s *Server) Unsubscribe(ctx context.Context, clientID string, query Query) error {
	select {
	case s.cmds <- cmd{op: unsub, clientID: clientID, query: query}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Unsubscribe removes all client subscriptions. An error will be returned to
// the caller if the context is canceled.
func (s *Server) UnsubscribeAll(ctx context.Context, clientID string) error {
	select {
	case s.cmds <- cmd{op: unsub, clientID: clientID}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Publish publishes the given message. An error will be returned to the caller
// if the context is canceled.
func (s *Server) Publish(ctx context.Context, msg interface{}) error {
	return s.PublishWithTags(ctx, msg, make(map[string]interface{}))
}

// PublishWithTags publishes the given message with the set of tags. The set is
// matched with clients queries. If there is a match, the message is sent to
// the client.
func (s *Server) PublishWithTags(ctx context.Context, msg interface{}, tags map[string]interface{}) error {
	select {
	case s.cmds <- cmd{op: pub, msg: msg, tags: tags}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// OnStop implements Service.OnStop by shutting down the server.
func (s *Server) OnStop() {
	s.cmds <- cmd{op: shutdown}
}

// NOTE: not goroutine safe
type state struct {
	// query -> client -> ch
	queries map[Query]map[string]chan<- interface{}
	// client -> query -> struct{}
	clients map[string]map[Query]struct{}
}

// OnStart implements Service.OnStart by starting the server.
func (s *Server) OnStart() error {
	go s.loop(state{
		queries: make(map[Query]map[string]chan<- interface{}),
		clients: make(map[string]map[Query]struct{}),
	})
	return nil
}

func (s *Server) loop(state state) {
loop:
	for cmd := range s.cmds {
		switch cmd.op {
		case unsub:
			if cmd.query != nil {
				state.remove(cmd.clientID, cmd.query)
			} else {
				state.removeAll(cmd.clientID)
			}
		case shutdown:
			for clientID, _ := range state.clients {
				state.removeAll(clientID)
			}
			break loop
		case sub:
			state.add(cmd.clientID, cmd.query, cmd.ch)
		case pub:
			state.send(cmd.msg, cmd.tags)
		}
	}
}

func (state *state) add(clientID string, q Query, ch chan<- interface{}) {
	// add query if needed
	if clientToChannelMap, ok := state.queries[q]; !ok {
		state.queries[q] = make(map[string]chan<- interface{})
	} else {
		// check if already subscribed
		if oldCh, ok := clientToChannelMap[clientID]; ok {
			close(oldCh)
		}
	}

	// create subscription
	state.queries[q][clientID] = ch

	// add client if needed
	if _, ok := state.clients[clientID]; !ok {
		state.clients[clientID] = make(map[Query]struct{})
	}
	state.clients[clientID][q] = struct{}{}
}

func (state *state) remove(clientID string, q Query) {
	clientToChannelMap, ok := state.queries[q]
	if !ok {
		return
	}

	ch, ok := clientToChannelMap[clientID]
	if ok {
		close(ch)

		delete(state.clients[clientID], q)

		// if it not subscribed to anything else, remove the client
		if len(state.clients[clientID]) == 0 {
			delete(state.clients, clientID)
		}

		delete(state.queries[q], clientID)
	}
}

func (state *state) removeAll(clientID string) {
	queryMap, ok := state.clients[clientID]
	if !ok {
		return
	}

	for q, _ := range queryMap {
		ch := state.queries[q][clientID]
		close(ch)

		delete(state.queries[q], clientID)
	}

	delete(state.clients, clientID)
}

func (state *state) send(msg interface{}, tags map[string]interface{}) {
	for q, clientToChannelMap := range state.queries {
		if q.Matches(tags) {
			for _, ch := range clientToChannelMap {
				ch <- msg
			}
		}
	}
}
<commit_msg>fixes from gometalinter review<commit_after>// Package pubsub implements a pub-sub model with a single publisher (Server)
// and multiple subscribers (clients).
//
// Though you can have multiple publishers by sharing a pointer to a server or
// by giving the same channel to each publisher and publishing messages from
// that channel (fan-in).
//
// Clients subscribe for messages, which could be of any type, using a query.
// When some message is published, we match it with all queries. If there is a
// match, this message will be pushed to all clients, subscribed to that query.
// See query subpackage for our implementation.
package pubsub

import (
	"context"

	cmn "github.com/tendermint/tmlibs/common"
)

type operation int

const (
	sub operation = iota
	pub
	unsub
	shutdown
)

type cmd struct {
	op       operation
	query    Query
	ch       chan<- interface{}
	clientID string
	msg      interface{}
	tags     map[string]interface{}
}

// Query defines an interface for a query to be used for subscribing.
type Query interface {
	Matches(tags map[string]interface{}) bool
}

// Server allows clients to subscribe/unsubscribe for messages, publishing
// messages with or without tags, and manages internal state.
type Server struct {
	cmn.BaseService

	cmds    chan cmd
	cmdsCap int
}

// Option sets a parameter for the server.
type Option func(*Server)

// NewServer returns a new server. See the commentary on the Option functions
// for a detailed description of how to configure buffering. If no options are
// provided, the resulting server's queue is unbuffered.
func NewServer(options ...Option) *Server {
	s := &Server{}
	s.BaseService = *cmn.NewBaseService(nil, "PubSub", s)

	for _, option := range options {
		option(s)
	}

	// if BufferCapacity option was not set, the channel is unbuffered
	s.cmds = make(chan cmd, s.cmdsCap)

	return s
}

// BufferCapacity allows you to specify capacity for the internal server's
// queue. Since the server, given Y subscribers, could only process X messages,
// this option could be used to survive spikes (e.g. high amount of
// transactions during peak hours).
func BufferCapacity(cap int) Option {
	return func(s *Server) {
		if cap > 0 {
			s.cmdsCap = cap
		}
	}
}

// BufferCapacity returns capacity of the internal server's queue.
func (s Server) BufferCapacity() int {
	return s.cmdsCap
}

// Subscribe creates a subscription for the given client. It accepts a channel
// on which messages matching the given query can be received. If the
// subscription already exists, the old channel will be closed. An error will
// be returned to the caller if the context is canceled.
func (s *Server) Subscribe(ctx context.Context, clientID string, query Query, out chan<- interface{}) error {
	select {
	case s.cmds <- cmd{op: sub, clientID: clientID, query: query, ch: out}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Unsubscribe removes the subscription on the given query. An error will be
// returned to the caller if the context is canceled.
func (s *Server) Unsubscribe(ctx context.Context, clientID string, query Query) error {
	select {
	case s.cmds <- cmd{op: unsub, clientID: clientID, query: query}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// UnsubscribeAll removes all client subscriptions. An error will be returned
// to the caller if the context is canceled.
func (s *Server) UnsubscribeAll(ctx context.Context, clientID string) error {
	select {
	case s.cmds <- cmd{op: unsub, clientID: clientID}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Publish publishes the given message. An error will be returned to the caller
// if the context is canceled.
func (s *Server) Publish(ctx context.Context, msg interface{}) error {
	return s.PublishWithTags(ctx, msg, make(map[string]interface{}))
}

// PublishWithTags publishes the given message with the set of tags. The set is
// matched with clients queries. If there is a match, the message is sent to
// the client.
func (s *Server) PublishWithTags(ctx context.Context, msg interface{}, tags map[string]interface{}) error {
	select {
	case s.cmds <- cmd{op: pub, msg: msg, tags: tags}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// OnStop implements Service.OnStop by shutting down the server.
func (s *Server) OnStop() {
	s.cmds <- cmd{op: shutdown}
}

// NOTE: not goroutine safe
type state struct {
	// query -> client -> ch
	queries map[Query]map[string]chan<- interface{}
	// client -> query -> struct{}
	clients map[string]map[Query]struct{}
}

// OnStart implements Service.OnStart by starting the server.
func (s *Server) OnStart() error {
	go s.loop(state{
		queries: make(map[Query]map[string]chan<- interface{}),
		clients: make(map[string]map[Query]struct{}),
	})
	return nil
}

func (s *Server) loop(state state) {
loop:
	for cmd := range s.cmds {
		switch cmd.op {
		case unsub:
			if cmd.query != nil {
				state.remove(cmd.clientID, cmd.query)
			} else {
				state.removeAll(cmd.clientID)
			}
		case shutdown:
			for clientID := range state.clients {
				state.removeAll(clientID)
			}
			break loop
		case sub:
			state.add(cmd.clientID, cmd.query, cmd.ch)
		case pub:
			state.send(cmd.msg, cmd.tags)
		}
	}
}

func (state *state) add(clientID string, q Query, ch chan<- interface{}) {
	// add query if needed
	if clientToChannelMap, ok := state.queries[q]; !ok {
		state.queries[q] = make(map[string]chan<- interface{})
	} else {
		// check if already subscribed
		if oldCh, ok := clientToChannelMap[clientID]; ok {
			close(oldCh)
		}
	}

	// create subscription
	state.queries[q][clientID] = ch

	// add client if needed
	if _, ok := state.clients[clientID]; !ok {
		state.clients[clientID] = make(map[Query]struct{})
	}
	state.clients[clientID][q] = struct{}{}
}

func (state *state) remove(clientID string, q Query) {
	clientToChannelMap, ok := state.queries[q]
	if !ok {
		return
	}

	ch, ok := clientToChannelMap[clientID]
	if ok {
		close(ch)

		delete(state.clients[clientID], q)

		// if it not subscribed to anything else, remove the client
		if len(state.clients[clientID]) == 0 {
			delete(state.clients, clientID)
		}

		delete(state.queries[q], clientID)
	}
}

func (state *state) removeAll(clientID string) {
	queryMap, ok := state.clients[clientID]
	if !ok {
		return
	}

	for q := range queryMap {
		ch := state.queries[q][clientID]
		close(ch)

		delete(state.queries[q], clientID)
	}

	delete(state.clients, clientID)
}

func (state *state) send(msg interface{}, tags map[string]interface{}) {
	for q, clientToChannelMap := range state.queries {
		if q.Matches(tags) {
			for _, ch := range clientToChannelMap {
				ch <- msg
			}
		}
	}
}
<|endoftext|>


<commit_before>// e2e provides an E2E test for TfJobs.
//
// The test creates TfJobs and runs various checks to ensure various operations work as intended.
// The test is intended to run as a helm test that ensures the TfJob operator is working correctly.
// Thus, the program returns non-zero exit status on error.
//
// TODO(jlewi): Do we need to make the test output conform to the TAP(https://testanything.org/)
// protocol so we can fit into the K8s dashboard
//
package main

import (
  "mlkube.io/pkg/util/k8sutil"
  log "github.com/golang/glog"
  "mlkube.io/pkg/spec"
  metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
  "mlkube.io/pkg/util"
  "github.com/gogo/protobuf/proto"
  "k8s.io/client-go/pkg/api/v1"
  "flag"
  "time"
  "strings"
  "fmt"
  k8s_errors "k8s.io/apimachinery/pkg/api/errors"
  "os"
)

const (
  Namespace = "default"
)

var (
  image = flag.String("image", "gcr.io/tf-on-k8s-dogfood/tf_sample:latest", "The Docker image to use with the TfJob.")
)

func run() error {
  kubeCli := k8sutil.MustNewKubeClient()
  tfJobClient, err := k8sutil.NewTfJobClient()
  if err != nil {
    return err
  }

  name := "e2e-test-job-" + util.RandString(4)

  original := &spec.TfJob{
    Metadata: metav1.ObjectMeta{
      Name: name,
      Labels: map[string] string {
        "test.mlkube.io": "",
      },
    },
    Spec: spec.TfJobSpec{
      ReplicaSpecs: []*spec.TfReplicaSpec{
        {
          Replicas:      proto.Int32(1),
          TfPort:        proto.Int32(2222),
          TfReplicaType: spec.MASTER,
          Template: &v1.PodTemplateSpec{
            Spec: v1.PodSpec{
              Containers: []v1.Container{
                {
                  Name:  "tensorflow",
                  Image: *image,
                },
              },
              RestartPolicy: v1.RestartPolicyOnFailure,
            },
          },
        },
        {
          Replicas:      proto.Int32(1),
          TfPort:        proto.Int32(2222),
          TfReplicaType: spec.PS,
          Template: &v1.PodTemplateSpec{
            Spec: v1.PodSpec{
              Containers: []v1.Container{
                {
                  Name:  "tensorflow",
                  Image: *image,
                },
              },
              RestartPolicy: v1.RestartPolicyOnFailure,
            },
          },
        },
        {
          Replicas:      proto.Int32(1),
          TfPort:        proto.Int32(2222),
          TfReplicaType: spec.WORKER,
          Template: &v1.PodTemplateSpec{
            Spec: v1.PodSpec{
              Containers: []v1.Container{
                {
                  Name:  "tensorflow",
                  Image: *image,
                },
              },
              RestartPolicy: v1.RestartPolicyOnFailure,
            },
          },
        },
      },
    },
  }

  _, err = tfJobClient.Create(Namespace, original)

  if err != nil {
    return err
  }

  // Wait for the job to complete for up to 2 minutes.
  var tfJob *spec.TfJob
  for endTime := time.Now().Add(2 * time.Minute); time.Now().Before(endTime); {
    tfJob, err = tfJobClient.Get(Namespace, name)
    if err != nil {
      log.Warningf("There was a problem getting TfJob: %v; error %v", name, err)
    }

    if tfJob.Status.State == spec.StateSucceeded  || tfJob.Status.State == spec.StateFailed {
      break
    }
    log.Infof("Waiting for job %v to finish", name)
    time.Sleep(5 * time.Second)
  }

  if tfJob == nil {
    return fmt.Errorf("Failed to get TfJob %v", name)
  }

  if tfJob.Status.State != spec.StateSucceeded {
    // TODO(jlewi): Should we clean up the job.
    return fmt.Errorf("TfJob %v did not succeed;\n %v", name, util.Pformat(tfJob))
  }

  if tfJob.Spec.RuntimeId == "" {
    return fmt.Errorf("TfJob %v doesn't have a RuntimeId", name)
  }

  // Loop over each replica and make sure the expected resources were created.
  for _, r := range original.Spec.ReplicaSpecs {
    baseName :=  strings.ToLower(string(r.TfReplicaType))

    for i := 0;  i < int(*r.Replicas); i += 1 {
      jobName := fmt.Sprintf("%v-%v-%v", baseName, tfJob.Spec.RuntimeId, i)

      _, err := kubeCli.BatchV1().Jobs(Namespace).Get(jobName, metav1.GetOptions{})

      if err != nil {
        return fmt.Errorf("Tfob %v did not create Job %v for ReplicaType %v Index %v", name, jobName, r.TfReplicaType, i)
      }
    }
  }

  // Delete the job and make sure all subresources are properly garbage collected.
  if err := tfJobClient.Delete(Namespace, name); err != nil {
    log.Fatal("Failed to delete TfJob %v; error %v", name, err)
  }

  // Define sets to keep track of Job controllers corresponding to Replicas
  // that still exist.
  jobs := make(map[string] bool)

  // Loop over each replica and make sure the expected resources are being deleted.
  for _, r := range original.Spec.ReplicaSpecs {
    baseName :=  strings.ToLower(string(r.TfReplicaType))

    for i := 0;  i < int(*r.Replicas); i += 1 {
      jobName := fmt.Sprintf("%v-%v-%v", baseName, tfJob.Spec.RuntimeId, i)

      jobs[jobName] = true
    }
  }

  // Wait for all jobs to be deleted.
  for endTime := time.Now().Add(5 * time.Minute); time.Now().Before(endTime) && len(jobs) >0; {
    for k := range jobs {
      _, err := kubeCli.BatchV1().Jobs(Namespace).Get(k, metav1.GetOptions{})
      if k8s_errors.IsNotFound(err) {
        // Deleting map entry during loop is safe.
        // See: https://stackoverflow.com/questions/23229975/is-it-safe-to-remove-selected-keys-from-golang-map-within-a-range-loop
        delete(jobs, k)
      } else {
        log.Infof("Job %v still exists", k)
      }
    }

    if len(jobs) > 0 {
      time.Sleep(5 * time.Second)
    }
  }

  if len(jobs) > 0 {
    return fmt.Errorf("Not all Job controllers were successfully deleted.")
  }
  return nil
}

func main() {
  flag.Parse()

  err := run()

  // Generate TAP (https://testanything.org/) output
  fmt.Println("1..1")
  if err == nil {
    fmt.Println("ok 1 - Successfully ran TfJob")
  } else {
    fmt.Println("not ok 1 - Running TfJob failed %v", err)
    // Exit with non zero exit code for Helm tests.
    os.Exit(1)
  }
}
<commit_msg>Fix a Print statement.<commit_after>// e2e provides an E2E test for TfJobs.
//
// The test creates TfJobs and runs various checks to ensure various operations work as intended.
// The test is intended to run as a helm test that ensures the TfJob operator is working correctly.
// Thus, the program returns non-zero exit status on error.
//
// TODO(jlewi): Do we need to make the test output conform to the TAP(https://testanything.org/)
// protocol so we can fit into the K8s dashboard
//
package main

import (
  "mlkube.io/pkg/util/k8sutil"
  log "github.com/golang/glog"
  "mlkube.io/pkg/spec"
  metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
  "mlkube.io/pkg/util"
  "github.com/gogo/protobuf/proto"
  "k8s.io/client-go/pkg/api/v1"
  "flag"
  "time"
  "strings"
  "fmt"
  k8s_errors "k8s.io/apimachinery/pkg/api/errors"
  "os"
)

const (
  Namespace = "default"
)

var (
  image = flag.String("image", "gcr.io/tf-on-k8s-dogfood/tf_sample:latest", "The Docker image to use with the TfJob.")
)

func run() error {
  kubeCli := k8sutil.MustNewKubeClient()
  tfJobClient, err := k8sutil.NewTfJobClient()
  if err != nil {
    return err
  }

  name := "e2e-test-job-" + util.RandString(4)

  original := &spec.TfJob{
    Metadata: metav1.ObjectMeta{
      Name: name,
      Labels: map[string] string {
        "test.mlkube.io": "",
      },
    },
    Spec: spec.TfJobSpec{
      ReplicaSpecs: []*spec.TfReplicaSpec{
        {
          Replicas:      proto.Int32(1),
          TfPort:        proto.Int32(2222),
          TfReplicaType: spec.MASTER,
          Template: &v1.PodTemplateSpec{
            Spec: v1.PodSpec{
              Containers: []v1.Container{
                {
                  Name:  "tensorflow",
                  Image: *image,
                },
              },
              RestartPolicy: v1.RestartPolicyOnFailure,
            },
          },
        },
        {
          Replicas:      proto.Int32(1),
          TfPort:        proto.Int32(2222),
          TfReplicaType: spec.PS,
          Template: &v1.PodTemplateSpec{
            Spec: v1.PodSpec{
              Containers: []v1.Container{
                {
                  Name:  "tensorflow",
                  Image: *image,
                },
              },
              RestartPolicy: v1.RestartPolicyOnFailure,
            },
          },
        },
        {
          Replicas:      proto.Int32(1),
          TfPort:        proto.Int32(2222),
          TfReplicaType: spec.WORKER,
          Template: &v1.PodTemplateSpec{
            Spec: v1.PodSpec{
              Containers: []v1.Container{
                {
                  Name:  "tensorflow",
                  Image: *image,
                },
              },
              RestartPolicy: v1.RestartPolicyOnFailure,
            },
          },
        },
      },
    },
  }

  _, err = tfJobClient.Create(Namespace, original)

  if err != nil {
    return err
  }

  // Wait for the job to complete for up to 2 minutes.
  var tfJob *spec.TfJob
  for endTime := time.Now().Add(2 * time.Minute); time.Now().Before(endTime); {
    tfJob, err = tfJobClient.Get(Namespace, name)
    if err != nil {
      log.Warningf("There was a problem getting TfJob: %v; error %v", name, err)
    }

    if tfJob.Status.State == spec.StateSucceeded  || tfJob.Status.State == spec.StateFailed {
      break
    }
    log.Infof("Waiting for job %v to finish", name)
    time.Sleep(5 * time.Second)
  }

  if tfJob == nil {
    return fmt.Errorf("Failed to get TfJob %v", name)
  }

  if tfJob.Status.State != spec.StateSucceeded {
    // TODO(jlewi): Should we clean up the job.
    return fmt.Errorf("TfJob %v did not succeed;\n %v", name, util.Pformat(tfJob))
  }

  if tfJob.Spec.RuntimeId == "" {
    return fmt.Errorf("TfJob %v doesn't have a RuntimeId", name)
  }

  // Loop over each replica and make sure the expected resources were created.
  for _, r := range original.Spec.ReplicaSpecs {
    baseName :=  strings.ToLower(string(r.TfReplicaType))

    for i := 0;  i < int(*r.Replicas); i += 1 {
      jobName := fmt.Sprintf("%v-%v-%v", baseName, tfJob.Spec.RuntimeId, i)

      _, err := kubeCli.BatchV1().Jobs(Namespace).Get(jobName, metav1.GetOptions{})

      if err != nil {
        return fmt.Errorf("Tfob %v did not create Job %v for ReplicaType %v Index %v", name, jobName, r.TfReplicaType, i)
      }
    }
  }

  // Delete the job and make sure all subresources are properly garbage collected.
  if err := tfJobClient.Delete(Namespace, name); err != nil {
    log.Fatal("Failed to delete TfJob %v; error %v", name, err)
  }

  // Define sets to keep track of Job controllers corresponding to Replicas
  // that still exist.
  jobs := make(map[string] bool)

  // Loop over each replica and make sure the expected resources are being deleted.
  for _, r := range original.Spec.ReplicaSpecs {
    baseName :=  strings.ToLower(string(r.TfReplicaType))

    for i := 0;  i < int(*r.Replicas); i += 1 {
      jobName := fmt.Sprintf("%v-%v-%v", baseName, tfJob.Spec.RuntimeId, i)

      jobs[jobName] = true
    }
  }

  // Wait for all jobs to be deleted.
  for endTime := time.Now().Add(5 * time.Minute); time.Now().Before(endTime) && len(jobs) >0; {
    for k := range jobs {
      _, err := kubeCli.BatchV1().Jobs(Namespace).Get(k, metav1.GetOptions{})
      if k8s_errors.IsNotFound(err) {
        // Deleting map entry during loop is safe.
        // See: https://stackoverflow.com/questions/23229975/is-it-safe-to-remove-selected-keys-from-golang-map-within-a-range-loop
        delete(jobs, k)
      } else {
        log.Infof("Job %v still exists", k)
      }
    }

    if len(jobs) > 0 {
      time.Sleep(5 * time.Second)
    }
  }

  if len(jobs) > 0 {
    return fmt.Errorf("Not all Job controllers were successfully deleted.")
  }
  return nil
}

func main() {
  flag.Parse()

  err := run()

  // Generate TAP (https://testanything.org/) output
  fmt.Println("1..1")
  if err == nil {
    fmt.Println("ok 1 - Successfully ran TfJob")
  } else {
    fmt.Printf("not ok 1 - Running TfJob failed %v \n", err)
    // Exit with non zero exit code for Helm tests.
    os.Exit(1)
  }
}
<|endoftext|>


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

package k8sTest

import (
	"fmt"
	"strings"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
)

var _ = Describe("K8sCLI", func() {
	var kubectl *helpers.Kubectl
	var ciliumFilename string

	BeforeAll(func() {
		kubectl = helpers.CreateKubectl(helpers.K8s1VMName(), logger)

		ciliumFilename = helpers.TimestampFilename("cilium.yaml")
		DeployCiliumAndDNS(kubectl, ciliumFilename)
		ExpectCiliumReady(kubectl)
	})

	AfterAll(func() {
		UninstallCiliumFromManifest(kubectl, ciliumFilename)
	})

	JustAfterEach(func() {
		kubectl.ValidateNoErrorsInLogs(CurrentGinkgoTestDescription().Duration)
	})

	SkipContextIf(helpers.DoesNotRunOnGKE, "Identity CLI testing", func() {
		const (
			manifestYAML = "test-cli.yaml"
			fooID        = "foo"
			fooSHA       = "a83c739e630049e46b9ac6883dc2682b31bf8472b09c8bb81d87092a51d14ddf"
			fooNode      = "k8s1"
			// These labels are automatically added to all pods in the default namespace.
			defaultLabels = "k8s:io.cilium.k8s.policy.cluster=default " +
				"k8s:io.cilium.k8s.policy.serviceaccount=default k8s:io.kubernetes.pod.namespace=default"
		)

		var (
			cliManifest string
			ciliumPod   string
			err         error
			identity    int64
		)

		BeforeAll(func() {
			cliManifest = helpers.ManifestGet(kubectl.BasePath(), manifestYAML)
			res := kubectl.ApplyDefault(cliManifest)
			res.ExpectSuccess("Unable to apply %s", cliManifest)
			err = kubectl.WaitforPods(helpers.DefaultNamespace, "-l id", helpers.HelperTimeout)
			Expect(err).Should(BeNil(), "The pods were not ready after timeout")

			ciliumPod, err = kubectl.GetCiliumPodOnNodeWithLabel(fooNode)
			Expect(err).Should(BeNil())

			err := kubectl.WaitForCEPIdentity(helpers.DefaultNamespace, fooID)
			Expect(err).Should(BeNil())

			ep, err := kubectl.GetCiliumEndpoint(helpers.DefaultNamespace, fooID)
			Expect(err).Should(BeNil(), fmt.Sprintf("Unable to get CEP for pod %s", fooID))
			identity = ep.Identity.ID
		})

		AfterAll(func() {
			_ = kubectl.Delete(cliManifest)
			ExpectAllPodsTerminated(kubectl)
		})

		It("Test labelsSHA256", func() {
			cmd := fmt.Sprintf("cilium identity get %d -o json", identity)
			res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, cmd)
			res.ExpectSuccess()
			out, err := res.Filter("{[0].labelsSHA256}")
			Expect(err).Should(BeNil(), "Error getting SHA from identity")
			Expect(out.String()).Should(Equal(fooSHA))
		})

		It("Test identity list", func() {
			By("Testing 'cilium identity list' for an endpoint's identity")
			cmd := fmt.Sprintf("cilium identity list k8s:id=%s %s", fooID, defaultLabels)
			res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, cmd)
			res.ExpectSuccess(fmt.Sprintf("Unable to get identity list output for label k8s:id=%s %s", fooID, defaultLabels))

			resSingleOut := res.SingleOut()
			containsIdentity := strings.Contains(resSingleOut, fmt.Sprintf("%d", identity))
			Expect(containsIdentity).To(BeTrue(), "Identity %d of endpoint %s not in 'cilium identity list' output", identity, resSingleOut)

			By("Testing 'cilium identity list' for reserved identities")
			res = kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium identity list")
			res.ExpectSuccess("Unable to get identity list output")
			resSingleOut = res.SingleOut()

			reservedIdentities := []string{"health", "host", "world", "init"}
			for _, id := range reservedIdentities {
				By("Checking that reserved identity '%s' is in 'cilium identity list' output", id)
				containsReservedIdentity := strings.Contains(resSingleOut, id)
				Expect(containsReservedIdentity).To(BeTrue(), "Reserved identity '%s' not in 'cilium identity list' output", id)
			}
		})
	})

	SkipContextIf(helpers.DoesNotRunOnGKE, "stdout/stderr testing", func() {
		var (
			ciliumPod string
			err       error
		)

		BeforeAll(func() {
			ciliumPod, err = kubectl.GetCiliumPodOnNodeWithLabel("k8s1")
			Expect(err).Should(BeNil())
		})

		It("Root command help should print to stdout", func() {
			res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium help")
			Expect(res.Stdout()).Should(ContainSubstring("Use \"cilium [command] --help\" for more information about a command."))
		})

		It("Subcommand help should print to stdout", func() {
			res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium help bpf")
			Expect(res.Stdout()).Should(ContainSubstring("Use \"cilium bpf [command] --help\" for more information about a command."))
		})

		It("Failed subcommand should print help to stdout", func() {
			res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium endpoint confi 173")
			Expect(res.Stdout()).Should(ContainSubstring("Use \"cilium endpoint [command] --help\" for more information about a command."))
		})
	})
})
<commit_msg>test: Avoid installing Cilium for K8sCLI if quarantined<commit_after>// Copyright 2018-2020 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package k8sTest

import (
	"fmt"
	"strings"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
)

var _ = Describe("K8sCLI", func() {
	SkipContextIf(helpers.DoesNotRunOnGKE, "CLI", func() {
		var kubectl *helpers.Kubectl
		var ciliumFilename string

		BeforeAll(func() {
			kubectl = helpers.CreateKubectl(helpers.K8s1VMName(), logger)

			ciliumFilename = helpers.TimestampFilename("cilium.yaml")
			DeployCiliumAndDNS(kubectl, ciliumFilename)
			ExpectCiliumReady(kubectl)
		})

		AfterAll(func() {
			UninstallCiliumFromManifest(kubectl, ciliumFilename)
		})

		JustAfterEach(func() {
			kubectl.ValidateNoErrorsInLogs(CurrentGinkgoTestDescription().Duration)
		})

		Context("Identity CLI testing", func() {
			const (
				manifestYAML = "test-cli.yaml"
				fooID        = "foo"
				fooSHA       = "a83c739e630049e46b9ac6883dc2682b31bf8472b09c8bb81d87092a51d14ddf"
				fooNode      = "k8s1"
				// These labels are automatically added to all pods in the default namespace.
				defaultLabels = "k8s:io.cilium.k8s.policy.cluster=default " +
					"k8s:io.cilium.k8s.policy.serviceaccount=default k8s:io.kubernetes.pod.namespace=default"
			)

			var (
				cliManifest string
				ciliumPod   string
				err         error
				identity    int64
			)

			BeforeAll(func() {
				cliManifest = helpers.ManifestGet(kubectl.BasePath(), manifestYAML)
				res := kubectl.ApplyDefault(cliManifest)
				res.ExpectSuccess("Unable to apply %s", cliManifest)
				err = kubectl.WaitforPods(helpers.DefaultNamespace, "-l id", helpers.HelperTimeout)
				Expect(err).Should(BeNil(), "The pods were not ready after timeout")

				ciliumPod, err = kubectl.GetCiliumPodOnNodeWithLabel(fooNode)
				Expect(err).Should(BeNil())

				err := kubectl.WaitForCEPIdentity(helpers.DefaultNamespace, fooID)
				Expect(err).Should(BeNil())

				ep, err := kubectl.GetCiliumEndpoint(helpers.DefaultNamespace, fooID)
				Expect(err).Should(BeNil(), fmt.Sprintf("Unable to get CEP for pod %s", fooID))
				identity = ep.Identity.ID
			})

			AfterAll(func() {
				_ = kubectl.Delete(cliManifest)
				ExpectAllPodsTerminated(kubectl)
			})

			It("Test labelsSHA256", func() {
				cmd := fmt.Sprintf("cilium identity get %d -o json", identity)
				res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, cmd)
				res.ExpectSuccess()
				out, err := res.Filter("{[0].labelsSHA256}")
				Expect(err).Should(BeNil(), "Error getting SHA from identity")
				Expect(out.String()).Should(Equal(fooSHA))
			})

			It("Test identity list", func() {
				By("Testing 'cilium identity list' for an endpoint's identity")
				cmd := fmt.Sprintf("cilium identity list k8s:id=%s %s", fooID, defaultLabels)
				res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, cmd)
				res.ExpectSuccess(fmt.Sprintf("Unable to get identity list output for label k8s:id=%s %s", fooID, defaultLabels))

				resSingleOut := res.SingleOut()
				containsIdentity := strings.Contains(resSingleOut, fmt.Sprintf("%d", identity))
				Expect(containsIdentity).To(BeTrue(), "Identity %d of endpoint %s not in 'cilium identity list' output", identity, resSingleOut)

				By("Testing 'cilium identity list' for reserved identities")
				res = kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium identity list")
				res.ExpectSuccess("Unable to get identity list output")
				resSingleOut = res.SingleOut()

				reservedIdentities := []string{"health", "host", "world", "init"}
				for _, id := range reservedIdentities {
					By("Checking that reserved identity '%s' is in 'cilium identity list' output", id)
					containsReservedIdentity := strings.Contains(resSingleOut, id)
					Expect(containsReservedIdentity).To(BeTrue(), "Reserved identity '%s' not in 'cilium identity list' output", id)
				}
			})
		})

		Context("stdout/stderr testing", func() {
			var (
				ciliumPod string
				err       error
			)

			BeforeAll(func() {
				ciliumPod, err = kubectl.GetCiliumPodOnNodeWithLabel("k8s1")
				Expect(err).Should(BeNil())
			})

			It("Root command help should print to stdout", func() {
				res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium help")
				Expect(res.Stdout()).Should(ContainSubstring("Use \"cilium [command] --help\" for more information about a command."))
			})

			It("Subcommand help should print to stdout", func() {
				res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium help bpf")
				Expect(res.Stdout()).Should(ContainSubstring("Use \"cilium bpf [command] --help\" for more information about a command."))
			})

			It("Failed subcommand should print help to stdout", func() {
				res := kubectl.ExecPodCmd(helpers.CiliumNamespace, ciliumPod, "cilium endpoint confi 173")
				Expect(res.Stdout()).Should(ContainSubstring("Use \"cilium endpoint [command] --help\" for more information about a command."))
			})
		})
	})
})
<|endoftext|>


<commit_before>// +build darwin linux
// run

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

// Test that maps don't go quadratic for NaNs and other values.

package main

import (
	"fmt"
	"math"
	"time"
)

// checkLinear asserts that the running time of f(n) is in O(n).
// tries is the initial number of iterations.
func checkLinear(typ string, tries int, f func(n int)) {
	// Depending on the machine and OS, this test might be too fast
	// to measure with accurate enough granularity. On failure,
	// make it run longer, hoping that the timing granularity
	// is eventually sufficient.

	timeF := func(n int) time.Duration {
		t1 := time.Now()
		f(n)
		return time.Since(t1)
	}

	t0 := time.Now()

	n := tries
	fails := 0
	for {
		t1 := timeF(n)
		t2 := timeF(2 * n)

		// should be 2x (linear); allow up to 3x
		if t2 < 3*t1 {
			if false {
				fmt.Println(typ, "\t", time.Since(t0))
			}
			return
		}
		fails++
		if fails == 6 {
			panic(fmt.Sprintf("%s: too slow: %d inserts: %v; %d inserts: %v\n",
				typ, n, t1, 2*n, t2))
		}
		if fails < 4 {
			n *= 2
		}
	}
}

type I interface {
	f()
}

type C int

func (C) f() {}

func main() {
	// NaNs. ~31ms on a 1.6GHz Zeon.
	checkLinear("NaN", 30000, func(n int) {
		m := map[float64]int{}
		nan := math.NaN()
		for i := 0; i < n; i++ {
			m[nan] = 1
		}
		if len(m) != n {
			panic("wrong size map after nan insertion")
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("eface", 10000, func(n int) {
		m := map[interface{}]int{}
		for i := 0; i < n; i++ {
			m[i] = 1
		}
	})

	// ~7ms on a 1.6GHz Zeon.
	// Regression test for CL 119360043.
	checkLinear("iface", 10000, func(n int) {
		m := map[I]int{}
		for i := 0; i < n; i++ {
			m[C(i)] = 1
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("int", 10000, func(n int) {
		m := map[int]int{}
		for i := 0; i < n; i++ {
			m[i] = 1
		}
	})

	// ~18ms on a 1.6GHz Zeon.
	checkLinear("string", 10000, func(n int) {
		m := map[string]int{}
		for i := 0; i < n; i++ {
			m[fmt.Sprint(i)] = 1
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("float32", 10000, func(n int) {
		m := map[float32]int{}
		for i := 0; i < n; i++ {
			m[float32(i)] = 1
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("float64", 10000, func(n int) {
		m := map[float64]int{}
		for i := 0; i < n; i++ {
			m[float64(i)] = 1
		}
	})

	// ~22ms on a 1.6GHz Zeon.
	checkLinear("complex64", 10000, func(n int) {
		m := map[complex64]int{}
		for i := 0; i < n; i++ {
			m[complex(float32(i), float32(i))] = 1
		}
	})

	// ~32ms on a 1.6GHz Zeon.
	checkLinear("complex128", 10000, func(n int) {
		m := map[complex128]int{}
		for i := 0; i < n; i++ {
			m[complex(float64(i), float64(i))] = 1
		}
	})

	// ~70ms on a 1.6GHz Zeon.
	// The iterate/delete idiom currently takes expected
	// O(n lg n) time.  Fortunately, the checkLinear test
	// leaves enough wiggle room to include n lg n time
	// (it actually tests for O(n^log_2(3)).
	checkLinear("iterdelete", 10000, func(n int) {
		m := map[int]int{}
		for i := 0; i < n; i++ {
			m[i] = i
		}
		for i := 0; i < n; i++ {
			for k := range m {
				delete(m, k)
				break
			}
		}
	})
}
<commit_msg>test: make maplinear iterdelete test less flaky<commit_after>// +build darwin linux
// run

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

// Test that maps don't go quadratic for NaNs and other values.

package main

import (
	"fmt"
	"math"
	"time"
)

// checkLinear asserts that the running time of f(n) is in O(n).
// tries is the initial number of iterations.
func checkLinear(typ string, tries int, f func(n int)) {
	// Depending on the machine and OS, this test might be too fast
	// to measure with accurate enough granularity. On failure,
	// make it run longer, hoping that the timing granularity
	// is eventually sufficient.

	timeF := func(n int) time.Duration {
		t1 := time.Now()
		f(n)
		return time.Since(t1)
	}

	t0 := time.Now()

	n := tries
	fails := 0
	for {
		t1 := timeF(n)
		t2 := timeF(2 * n)

		// should be 2x (linear); allow up to 3x
		if t2 < 3*t1 {
			if false {
				fmt.Println(typ, "\t", time.Since(t0))
			}
			return
		}
		fails++
		if fails == 6 {
			panic(fmt.Sprintf("%s: too slow: %d inserts: %v; %d inserts: %v\n",
				typ, n, t1, 2*n, t2))
		}
		if fails < 4 {
			n *= 2
		}
	}
}

type I interface {
	f()
}

type C int

func (C) f() {}

func main() {
	// NaNs. ~31ms on a 1.6GHz Zeon.
	checkLinear("NaN", 30000, func(n int) {
		m := map[float64]int{}
		nan := math.NaN()
		for i := 0; i < n; i++ {
			m[nan] = 1
		}
		if len(m) != n {
			panic("wrong size map after nan insertion")
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("eface", 10000, func(n int) {
		m := map[interface{}]int{}
		for i := 0; i < n; i++ {
			m[i] = 1
		}
	})

	// ~7ms on a 1.6GHz Zeon.
	// Regression test for CL 119360043.
	checkLinear("iface", 10000, func(n int) {
		m := map[I]int{}
		for i := 0; i < n; i++ {
			m[C(i)] = 1
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("int", 10000, func(n int) {
		m := map[int]int{}
		for i := 0; i < n; i++ {
			m[i] = 1
		}
	})

	// ~18ms on a 1.6GHz Zeon.
	checkLinear("string", 10000, func(n int) {
		m := map[string]int{}
		for i := 0; i < n; i++ {
			m[fmt.Sprint(i)] = 1
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("float32", 10000, func(n int) {
		m := map[float32]int{}
		for i := 0; i < n; i++ {
			m[float32(i)] = 1
		}
	})

	// ~6ms on a 1.6GHz Zeon.
	checkLinear("float64", 10000, func(n int) {
		m := map[float64]int{}
		for i := 0; i < n; i++ {
			m[float64(i)] = 1
		}
	})

	// ~22ms on a 1.6GHz Zeon.
	checkLinear("complex64", 10000, func(n int) {
		m := map[complex64]int{}
		for i := 0; i < n; i++ {
			m[complex(float32(i), float32(i))] = 1
		}
	})

	// ~32ms on a 1.6GHz Zeon.
	checkLinear("complex128", 10000, func(n int) {
		m := map[complex128]int{}
		for i := 0; i < n; i++ {
			m[complex(float64(i), float64(i))] = 1
		}
	})

	// ~70ms on a 1.6GHz Zeon.
	// The iterate/delete idiom currently takes expected
	// O(n lg n) time.  Fortunately, the checkLinear test
	// leaves enough wiggle room to include n lg n time
	// (it actually tests for O(n^log_2(3)).
	// To prevent false positives, average away variation
	// by doing multiple rounds within a single run.
	checkLinear("iterdelete", 2500, func(n int) {
		for round := 0; round < 4; round++ {
			m := map[int]int{}
			for i := 0; i < n; i++ {
				m[i] = i
			}
			for i := 0; i < n; i++ {
				for k := range m {
					delete(m, k)
					break
				}
			}
		}
	})
}
<|endoftext|>