<commit_before>package rcf

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"io"
	"os"
	"reflect"
	"sync"

	"github.com/ugorji/go/codec"
)

type File struct {
	sync.Mutex
	file         *os.File
	path         string
	colSets      [][]string
	colSetsFn    func(int) interface{}
	validateOnce sync.Once
}

var codecHandle = new(codec.CborHandle)

func (f *File) Sync() error {
	f.Lock()
	defer f.Unlock()
	return f.file.Sync()
}

func (f *File) Close() error {
	return f.file.Close()
}

func New(path string, colSetsFn func(int) interface{}) (*File, error) {
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return nil, makeErr(err, "open file")
	}
	n := 0
	colSets := [][]string{}
	for {
		v := colSetsFn(n)
		if v == nil {
			break
		}
		t := reflect.TypeOf(v).Elem()
		set := []string{}
		for i, max := 0, t.NumField(); i < max; i++ {
			set = append(set, t.Field(i).Name)
		}
		colSets = append(colSets, set)
		n++
	}
	ret := &File{
		file:      file,
		path:      path,
		colSets:   colSets,
		colSetsFn: colSetsFn,
	}
	return ret, nil
}

func (f *File) validate() (err error) {
	f.validateOnce.Do(func() {
	read:
		// read number of sets
		var numSets uint8
		err = binary.Read(f.file, binary.LittleEndian, &numSets)
		if err == io.EOF { // no more
			return
		}
		if err != nil {
			err = makeErr(err, "read number of column sets")
			return
		}
		// read meta length
		var sum, l uint32
		err = binary.Read(f.file, binary.LittleEndian, &l)
		if err != nil {
			err = makeErr(err, "read meta length")
			return
		}
		sum += l
		// read sets length
		for i, max := 0, int(numSets); i < max; i++ {
			err = binary.Read(f.file, binary.LittleEndian, &l)
			if err != nil {
				err = makeErr(err, "read column set length")
				return
			}
			sum += l
		}
		_, err = f.file.Seek(int64(sum), os.SEEK_CUR)
		if err != nil {
			err = makeErr(err, "validate seek")
			return
		}
		goto read
		return
	})
	return
}

func encode(o interface{}) (bs []byte, err error) {
	buf := new(bytes.Buffer)
	w := gzip.NewWriter(buf)
	err = codec.NewEncoder(w, codecHandle).Encode(o)
	if err != nil {
		return nil, err
	}
	err = w.Close()
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decode(bs []byte, target interface{}) (err error) {
	r, err := gzip.NewReader(bytes.NewReader(bs))
	if err != nil {
		return err
	}
	return codec.NewDecoder(r, codecHandle).Decode(target)
}

func (f *File) Append(rows, meta interface{}) error {
	f.validate()
	// encode meta
	metaBin, err := encode(meta)
	if err != nil {
		return makeErr(err, "encode meta")
	}
	// column slices
	rowsValue := reflect.ValueOf(rows)
	if rowsValue.Type().Kind() != reflect.Slice {
		return makeErr(nil, "rows is not slice")
	}
	columns := make(map[string]reflect.Value)
	for i, l := 0, rowsValue.Len(); i < l; i++ {
		row := rowsValue.Index(i)
		// make colums slices
		if i == 0 {
			for _, set := range f.colSets {
				for _, col := range set {
					if _, ok := columns[col]; ok {
						continue
					}
					columns[col] = reflect.MakeSlice(reflect.SliceOf(row.FieldByName(col).Type()), 0, 0)
				}
			}
		}
		// append colum values
		for name, col := range columns {
			columns[name] = reflect.Append(col, row.FieldByName(name))
		}
	}
	// column sets
	var bins [][]byte
	for n, set := range f.colSets {
		s := reflect.ValueOf(f.colSetsFn(n))
		for _, col := range set {
			s.Elem().FieldByName(col).Set(columns[col])
		}
		bin, err := encode(s)
		if err != nil {
			return makeErr(err, "encode column set")
		}
		bins = append(bins, bin)
	}
	// write header
	if len(bins) > 255 {
		return makeErr(nil, "more than 255 column sets")
	}
	f.Lock()
	defer f.Unlock()
	err = binary.Write(f.file, binary.LittleEndian, uint8(len(bins)))
	if err != nil {
		return makeErr(err, "write length length")
	}
	err = binary.Write(f.file, binary.LittleEndian, uint32(len(metaBin)))
	if err != nil {
		return makeErr(err, "write meta length")
	}
	for _, bin := range bins {
		err = binary.Write(f.file, binary.LittleEndian, uint32(len(bin)))
		if err != nil {
			return makeErr(err, "write column set length")
		}
	}
	// write encoded
	_, err = f.file.Write(metaBin)
	if err != nil {
		return makeErr(err, "write meta")
	}
	for _, bin := range bins {
		_, err = f.file.Write(bin)
		if err != nil {
			return makeErr(err, "write column set")
		}
	}
	return nil
}

func (f *File) IterMetas(fn interface{}) error {
	f.Sync()
	file, err := os.Open(f.path)
	if err != nil {
		return makeErr(err, "open file")
	}
	defer file.Close()
	fnValue := reflect.ValueOf(fn)
	fnType := fnValue.Type()
	metaType := fnType.In(0)
	decodeValue := reflect.ValueOf(decode)
	meta := reflect.New(metaType)
read:
	// read number of sets
	var numSets uint8
	err = binary.Read(file, binary.LittleEndian, &numSets)
	if err == io.EOF { // no more
		return nil
	}
	if err != nil {
		return makeErr(err, "read number of column sets")
	}
	// read meta length
	var metaLength uint32
	err = binary.Read(file, binary.LittleEndian, &metaLength)
	if err != nil {
		return makeErr(err, "read meta length")
	}
	// read sets length
	var sum, l uint32
	for i, max := 0, int(numSets); i < max; i++ {
		err = binary.Read(file, binary.LittleEndian, &l)
		if err != nil {
			return makeErr(err, "read column set length")
		}
		sum += l
	}
	// read meta
	bs := make([]byte, metaLength)
	_, err = io.ReadFull(file, bs)
	if err != nil {
		return makeErr(err, "read meta")
	}
	// decode meta
	ret := decodeValue.Call([]reflect.Value{
		reflect.ValueOf(bs),
		meta,
	})
	if e := ret[0].Interface(); e != nil {
		return makeErr(e.(error), "decode meta")
	}
	if !fnValue.Call([]reflect.Value{meta.Elem()})[0].Bool() {
		return nil
	}
	// skip sets
	_, err = file.Seek(int64(sum), os.SEEK_CUR)
	if err != nil {
		return makeErr(err, "skip column sets")
	}
	goto read
	return nil
}

func (f *File) IterRows(cols []string, cb interface{}) error {
	f.Sync()
	file, err := os.Open(f.path)
	if err != nil {
		return makeErr(err, "open file")
	}
	defer file.Close()
	cbValue := reflect.ValueOf(cb)
	// determine which set to decode
	toDecode := make([]bool, len(f.colSets))
	for n, set := range f.colSets {
	loop:
		for _, col := range set {
			for _, c := range cols {
				if c == col {
					toDecode[n] = true
					break loop
				}
			}
		}
	}
	// read number of sets
read:
	var numSets uint8
	err = binary.Read(file, binary.LittleEndian, &numSets)
	if err == io.EOF { // no more
		return nil
	}
	if err != nil {
		return makeErr(err, "read number of column sets")
	}
	// read meta length
	var metaLength uint32
	err = binary.Read(file, binary.LittleEndian, &metaLength)
	if err != nil {
		return makeErr(err, "read meta length")
	}
	// read sets length
	var lens []uint32
	var l uint32
	for i, max := 0, int(numSets); i < max; i++ {
		err = binary.Read(file, binary.LittleEndian, &l)
		if err != nil {
			return makeErr(err, "read column set length")
		}
		lens = append(lens, l)
	}
	// skip meta
	_, err = file.Seek(int64(metaLength), os.SEEK_CUR)
	if err != nil {
		return makeErr(err, "skip meta")
	}
	// decode sets
	columns := make(map[string]reflect.Value)
	for n, l := range lens {
		if toDecode[n] { // decode
			bs := make([]byte, l)
			_, err = io.ReadFull(file, bs)
			if err != nil {
				return makeErr(err, "read column set")
			}
			s := f.colSetsFn(n)
			err = decode(bs, &s)
			if err != nil {
				return makeErr(err, "decode column set")
			}
			sValue := reflect.ValueOf(s).Elem()
			sType := sValue.Type()
			for i, max := 0, sValue.NumField(); i < max; i++ {
				columns[sType.Field(i).Name] = sValue.Field(i)
			}
		} else { // skip
			_, err = file.Seek(int64(l), os.SEEK_CUR)
			if err != nil {
				return makeErr(err, "skip column set")
			}
		}
	}
	// call
	for i, max := 0, columns[cols[0]].Len(); i < max; i++ {
		var args []reflect.Value
		for _, col := range cols {
			args = append(args, columns[col].Index(i))
		}
		rets := cbValue.Call(args)
		if !rets[0].Bool() {
			return nil
		}
	}
	goto read
	return nil
}
<commit_msg>utilize multi-core in IterRows<commit_after>package rcf

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"io"
	"os"
	"reflect"
	"runtime"
	"sync"

	"github.com/ugorji/go/codec"
)

type File struct {
	sync.Mutex
	file         *os.File
	path         string
	colSets      [][]string
	colSetsFn    func(int) interface{}
	validateOnce sync.Once
}

var codecHandle = new(codec.CborHandle)

func (f *File) Sync() error {
	f.Lock()
	defer f.Unlock()
	return f.file.Sync()
}

func (f *File) Close() error {
	return f.file.Close()
}

func New(path string, colSetsFn func(int) interface{}) (*File, error) {
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return nil, makeErr(err, "open file")
	}
	n := 0
	colSets := [][]string{}
	for {
		v := colSetsFn(n)
		if v == nil {
			break
		}
		t := reflect.TypeOf(v).Elem()
		set := []string{}
		for i, max := 0, t.NumField(); i < max; i++ {
			set = append(set, t.Field(i).Name)
		}
		colSets = append(colSets, set)
		n++
	}
	ret := &File{
		file:      file,
		path:      path,
		colSets:   colSets,
		colSetsFn: colSetsFn,
	}
	return ret, nil
}

func (f *File) validate() (err error) {
	f.validateOnce.Do(func() {
	read:
		// read number of sets
		var numSets uint8
		err = binary.Read(f.file, binary.LittleEndian, &numSets)
		if err == io.EOF { // no more
			return
		}
		if err != nil {
			err = makeErr(err, "read number of column sets")
			return
		}
		// read meta length
		var sum, l uint32
		err = binary.Read(f.file, binary.LittleEndian, &l)
		if err != nil {
			err = makeErr(err, "read meta length")
			return
		}
		sum += l
		// read sets length
		for i, max := 0, int(numSets); i < max; i++ {
			err = binary.Read(f.file, binary.LittleEndian, &l)
			if err != nil {
				err = makeErr(err, "read column set length")
				return
			}
			sum += l
		}
		_, err = f.file.Seek(int64(sum), os.SEEK_CUR)
		if err != nil {
			err = makeErr(err, "validate seek")
			return
		}
		goto read
		return
	})
	return
}

func encode(o interface{}) (bs []byte, err error) {
	buf := new(bytes.Buffer)
	w := gzip.NewWriter(buf)
	err = codec.NewEncoder(w, codecHandle).Encode(o)
	if err != nil {
		return nil, err
	}
	err = w.Close()
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decode(bs []byte, target interface{}) (err error) {
	r, err := gzip.NewReader(bytes.NewReader(bs))
	if err != nil {
		return err
	}
	return codec.NewDecoder(r, codecHandle).Decode(target)
}

func (f *File) Append(rows, meta interface{}) error {
	f.validate()
	// encode meta
	metaBin, err := encode(meta)
	if err != nil {
		return makeErr(err, "encode meta")
	}
	// column slices
	rowsValue := reflect.ValueOf(rows)
	if rowsValue.Type().Kind() != reflect.Slice {
		return makeErr(nil, "rows is not slice")
	}
	columns := make(map[string]reflect.Value)
	for i, l := 0, rowsValue.Len(); i < l; i++ {
		row := rowsValue.Index(i)
		// make colums slices
		if i == 0 {
			for _, set := range f.colSets {
				for _, col := range set {
					if _, ok := columns[col]; ok {
						continue
					}
					columns[col] = reflect.MakeSlice(reflect.SliceOf(row.FieldByName(col).Type()), 0, 0)
				}
			}
		}
		// append colum values
		for name, col := range columns {
			columns[name] = reflect.Append(col, row.FieldByName(name))
		}
	}
	// column sets
	var bins [][]byte
	for n, set := range f.colSets {
		s := reflect.ValueOf(f.colSetsFn(n))
		for _, col := range set {
			s.Elem().FieldByName(col).Set(columns[col])
		}
		bin, err := encode(s)
		if err != nil {
			return makeErr(err, "encode column set")
		}
		bins = append(bins, bin)
	}
	// write header
	if len(bins) > 255 {
		return makeErr(nil, "more than 255 column sets")
	}
	f.Lock()
	defer f.Unlock()
	err = binary.Write(f.file, binary.LittleEndian, uint8(len(bins)))
	if err != nil {
		return makeErr(err, "write length length")
	}
	err = binary.Write(f.file, binary.LittleEndian, uint32(len(metaBin)))
	if err != nil {
		return makeErr(err, "write meta length")
	}
	for _, bin := range bins {
		err = binary.Write(f.file, binary.LittleEndian, uint32(len(bin)))
		if err != nil {
			return makeErr(err, "write column set length")
		}
	}
	// write encoded
	_, err = f.file.Write(metaBin)
	if err != nil {
		return makeErr(err, "write meta")
	}
	for _, bin := range bins {
		_, err = f.file.Write(bin)
		if err != nil {
			return makeErr(err, "write column set")
		}
	}
	return nil
}

func (f *File) IterMetas(fn interface{}) error {
	f.Sync()
	file, err := os.Open(f.path)
	if err != nil {
		return makeErr(err, "open file")
	}
	defer file.Close()
	fnValue := reflect.ValueOf(fn)
	fnType := fnValue.Type()
	metaType := fnType.In(0)
	decodeValue := reflect.ValueOf(decode)
	meta := reflect.New(metaType)
read:
	// read number of sets
	var numSets uint8
	err = binary.Read(file, binary.LittleEndian, &numSets)
	if err == io.EOF { // no more
		return nil
	}
	if err != nil {
		return makeErr(err, "read number of column sets")
	}
	// read meta length
	var metaLength uint32
	err = binary.Read(file, binary.LittleEndian, &metaLength)
	if err != nil {
		return makeErr(err, "read meta length")
	}
	// read sets length
	var sum, l uint32
	for i, max := 0, int(numSets); i < max; i++ {
		err = binary.Read(file, binary.LittleEndian, &l)
		if err != nil {
			return makeErr(err, "read column set length")
		}
		sum += l
	}
	// read meta
	bs := make([]byte, metaLength)
	_, err = io.ReadFull(file, bs)
	if err != nil {
		return makeErr(err, "read meta")
	}
	// decode meta
	ret := decodeValue.Call([]reflect.Value{
		reflect.ValueOf(bs),
		meta,
	})
	if e := ret[0].Interface(); e != nil {
		return makeErr(e.(error), "decode meta")
	}
	if !fnValue.Call([]reflect.Value{meta.Elem()})[0].Bool() {
		return nil
	}
	// skip sets
	_, err = file.Seek(int64(sum), os.SEEK_CUR)
	if err != nil {
		return makeErr(err, "skip column sets")
	}
	goto read
	return nil
}

func (f *File) IterRows(cols []string, cb interface{}) error {
	f.Sync()
	file, err := os.Open(f.path)
	if err != nil {
		return makeErr(err, "open file")
	}
	defer file.Close()
	cbValue := reflect.ValueOf(cb)
	// determine which set to decode
	toDecode := make([]bool, len(f.colSets))
	for n, set := range f.colSets {
	loop:
		for _, col := range set {
			for _, c := range cols {
				if c == col {
					toDecode[n] = true
					break loop
				}
			}
		}
	}

	done := make(chan struct{})
	setErr := func(e error) {
		err = e
	}

	// read bytes
	bins := make(chan [][]byte)
	go func() {
	loop:
		for {
			// read number of sets
			var numSets uint8
			err := binary.Read(file, binary.LittleEndian, &numSets)
			if err == io.EOF { // no more
				break loop
			}
			if err != nil {
				setErr(makeErr(err, "read number of column sets"))
				break loop
			}
			// read meta length
			var metaLength uint32
			err = binary.Read(file, binary.LittleEndian, &metaLength)
			if err != nil {
				setErr(makeErr(err, "read meta length"))
				break loop
			}
			// read sets length
			var lens []uint32
			var l uint32
			for i, max := 0, int(numSets); i < max; i++ {
				err = binary.Read(file, binary.LittleEndian, &l)
				if err != nil {
					setErr(makeErr(err, "read column set length"))
					break loop
				}
				lens = append(lens, l)
			}
			// skip meta
			_, err = file.Seek(int64(metaLength), os.SEEK_CUR)
			if err != nil {
				setErr(makeErr(err, "skip meta"))
				break loop
			}
			// read bytes
			var bss [][]byte
			for n, l := range lens {
				if toDecode[n] { // decode
					bs := make([]byte, l)
					_, err = io.ReadFull(file, bs)
					if err != nil {
						setErr(makeErr(err, "read column set"))
						break loop
					}
					bss = append(bss, bs)
				} else { // skip
					_, err = file.Seek(int64(l), os.SEEK_CUR)
					if err != nil {
						setErr(makeErr(err, "skip column set"))
						break loop
					}
					bss = append(bss, nil)
				}
			}
			select {
			case bins <- bss:
			case <-done:
				break loop
			}
		}
		close(bins)
	}()

	argsChan := make(chan []reflect.Value)
	ncpu := runtime.NumCPU()
	wg := new(sync.WaitGroup)
	wg.Add(ncpu)
	for i := 0; i < ncpu; i++ {
		go func() {
			defer wg.Done()
		loop:
			for bss := range bins {
				columns := make(map[string]reflect.Value)
				for n, bs := range bss {
					if bs == nil {
						continue
					}
					s := f.colSetsFn(n)
					err := decode(bs, &s)
					if err != nil {
						setErr(makeErr(err, "decode column set"))
						break loop
					}
					sValue := reflect.ValueOf(s).Elem()
					sType := sValue.Type()
					for i, max := 0, sValue.NumField(); i < max; i++ {
						columns[sType.Field(i).Name] = sValue.Field(i)
					}
				}
				for i, max := 0, columns[cols[0]].Len(); i < max; i++ {
					var args []reflect.Value
					for _, col := range cols {
						args = append(args, columns[col].Index(i))
					}
					select {
					case argsChan <- args:
					case <-done:
						break loop
					}
				}
			}
		}()
	}
	go func() {
		wg.Wait()
		close(argsChan)
	}()

	// call
	for args := range argsChan {
		rets := cbValue.Call(args)
		if !rets[0].Bool() {
			return err
		}
	}

	return err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/c-14/grue/config"
	"github.com/mmcdole/gofeed"
	"math"
	"time"
)

type FeedParser struct {
	parser   *gofeed.Parser
	messages chan *Email
	sem      chan int
	finished chan int
}

type RSSFeed struct {
	config      config.AccountConfig
	LastFetched int64               `json:",omitempty"`
	LastQueried int64               `json:",omitempty"`
	NextQuery   int64               `json:",omitempty"`
	Tries       int                 `json:",omitempty"`
	GUIDList    map[string]struct{} `json:",omitempty"`
}

type DateType int

const (
	NoDate DateType = iota
	DateNewer
	DateOlder
)

func hasNewerDate(item *gofeed.Item, lastFetched int64) (time.Time, DateType) {
	if item.PublishedParsed != nil {
		if item.PublishedParsed.Unix() > lastFetched {
			return *item.PublishedParsed, DateNewer
		} else {
			return *item.PublishedParsed, DateOlder
		}
	} else if item.UpdatedParsed != nil {
		if item.UpdatedParsed.Unix() > lastFetched {
			return *item.UpdatedParsed, DateNewer
		} else {
			return *item.UpdatedParsed, DateOlder
		}
	} else if date, exists := item.Extensions["dc"]["date"]; exists {
		dateParsed, err := time.Parse(time.RFC3339, date[0].Value)
		if err != nil {
			fmt.Printf("Can't parse (%v) as dc:date for (%v)\n", date, item.Link)
			return time.Now(), NoDate
		}
		if dateParsed.Unix() > lastFetched {
			return dateParsed, DateNewer
		} else {
			return dateParsed, DateOlder
		}
	}
	return time.Now(), NoDate
}

func fetchFeed(fp FeedParser, feedName string, account *RSSFeed, config *config.GrueConfig) {
	// if account.UserAgent != nil {
	// 	feed.SetUserAgent(*account.UserAgent)
	// }
	now := time.Now()
	if account.NextQuery > now.Unix() {
		return
	}
	feed, err := fp.parser.ParseURL(account.config.URI)
	account.LastQueried = now.Unix()
	if err != nil {
		if account.Tries > 0 {
			account.NextQuery = now.Add(time.Duration(math.Exp2(float64(account.Tries+4))) * time.Minute).Unix()
		}
		account.Tries++
		fmt.Printf("Caught error when parsing %s: %s\n", account.config.URI, err)
		return
	}
	account.NextQuery = 0
	account.Tries = 0
	guids := account.GUIDList
	if float64(len(guids)) > 1.2*float64(len(feed.Items)) {
		account.GUIDList = make(map[string]struct{})
	}
	for _, item := range feed.Items {
		date, newer := hasNewerDate(item, account.LastFetched)
		if newer == DateNewer {
			e := createEmail(feedName, feed.Title, item, date, account.config, config)
			err = e.Send(fp.messages)
			if err != nil {
				break
			}
		} else if newer == NoDate {
			_, exists := guids[item.GUID]
			if !exists {
				e := createEmail(feedName, feed.Title, item, date, account.config, config)
				err = e.Send(fp.messages)
			}
			if err == nil {
				account.GUIDList[item.GUID] = struct{}{}
			} else {
				break
			}
		}
	}
	if err == nil {
		account.LastFetched = account.LastQueried
	}

	<-fp.sem
	fp.finished <- 1
}

func fetchFeeds(ret chan error, conf *config.GrueConfig, init bool) {
	defer close(ret)
	hist, err := ReadHistory()
	if err != nil {
		ret <- err
		return
	}
	ch := make(chan *Email)
	defer close(ch)
	if !init {
		s, err := setupDialer(conf)
		if err != nil {
			ret <- err
			return
		}
		go startDialing(s, ch, ret)
	} else {
		go func() {
			for m := range ch {
				m.ret <- nil
			}
		}()
	}

	fp := FeedParser{parser: gofeed.NewParser(), messages: ch, sem: make(chan int, 10), finished: make(chan int)}
	go func() {
		for name, accountConfig := range conf.Accounts {
			fp.sem <- 1
			account, exist := hist.Feeds[name]
			if !exist {
				account = new(RSSFeed)
				account.GUIDList = make(map[string]struct{})
				hist.Feeds[name] = account
			} else if len(account.GUIDList) == 0 {
				account.GUIDList = make(map[string]struct{})
			}
			account.config = accountConfig
			go fetchFeed(fp, name, account, conf)
		}
	}()
	for range conf.Accounts {
		<-fp.finished
	}
	ret <- hist.Write()
}
<commit_msg>Unlock semaphore and send finished on error cases<commit_after>package main

import (
	"fmt"
	"github.com/c-14/grue/config"
	"github.com/mmcdole/gofeed"
	"math"
	"time"
)

type FeedParser struct {
	parser   *gofeed.Parser
	messages chan *Email
	sem      chan int
	finished chan int
}

type RSSFeed struct {
	config      config.AccountConfig
	LastFetched int64               `json:",omitempty"`
	LastQueried int64               `json:",omitempty"`
	NextQuery   int64               `json:",omitempty"`
	Tries       int                 `json:",omitempty"`
	GUIDList    map[string]struct{} `json:",omitempty"`
}

type DateType int

const (
	NoDate DateType = iota
	DateNewer
	DateOlder
)

func hasNewerDate(item *gofeed.Item, lastFetched int64) (time.Time, DateType) {
	if item.PublishedParsed != nil {
		if item.PublishedParsed.Unix() > lastFetched {
			return *item.PublishedParsed, DateNewer
		} else {
			return *item.PublishedParsed, DateOlder
		}
	} else if item.UpdatedParsed != nil {
		if item.UpdatedParsed.Unix() > lastFetched {
			return *item.UpdatedParsed, DateNewer
		} else {
			return *item.UpdatedParsed, DateOlder
		}
	} else if date, exists := item.Extensions["dc"]["date"]; exists {
		dateParsed, err := time.Parse(time.RFC3339, date[0].Value)
		if err != nil {
			fmt.Printf("Can't parse (%v) as dc:date for (%v)\n", date, item.Link)
			return time.Now(), NoDate
		}
		if dateParsed.Unix() > lastFetched {
			return dateParsed, DateNewer
		} else {
			return dateParsed, DateOlder
		}
	}
	return time.Now(), NoDate
}

func fetchFeed(fp FeedParser, feedName string, account *RSSFeed, config *config.GrueConfig) {
	// if account.UserAgent != nil {
	// 	feed.SetUserAgent(*account.UserAgent)
	// }
	now := time.Now()
	if account.NextQuery > now.Unix() {
		<-fp.sem
		fp.finished <- 1
		return
	}
	feed, err := fp.parser.ParseURL(account.config.URI)
	account.LastQueried = now.Unix()
	if err != nil {
		if account.Tries > 0 {
			account.NextQuery = now.Add(time.Duration(math.Exp2(float64(account.Tries+4))) * time.Minute).Unix()
		}
		account.Tries++
		fmt.Printf("Caught error when parsing %s: %s\n", account.config.URI, err)
		<-fp.sem
		fp.finished <- 1
		return
	}
	account.NextQuery = 0
	account.Tries = 0
	guids := account.GUIDList
	if float64(len(guids)) > 1.2*float64(len(feed.Items)) {
		account.GUIDList = make(map[string]struct{})
	}
	for _, item := range feed.Items {
		date, newer := hasNewerDate(item, account.LastFetched)
		if newer == DateNewer {
			e := createEmail(feedName, feed.Title, item, date, account.config, config)
			err = e.Send(fp.messages)
			if err != nil {
				break
			}
		} else if newer == NoDate {
			_, exists := guids[item.GUID]
			if !exists {
				e := createEmail(feedName, feed.Title, item, date, account.config, config)
				err = e.Send(fp.messages)
			}
			if err == nil {
				account.GUIDList[item.GUID] = struct{}{}
			} else {
				break
			}
		}
	}
	if err == nil {
		account.LastFetched = account.LastQueried
	}

	<-fp.sem
	fp.finished <- 1
}

func fetchFeeds(ret chan error, conf *config.GrueConfig, init bool) {
	defer close(ret)
	hist, err := ReadHistory()
	if err != nil {
		ret <- err
		return
	}
	ch := make(chan *Email)
	defer close(ch)
	if !init {
		s, err := setupDialer(conf)
		if err != nil {
			ret <- err
			return
		}
		go startDialing(s, ch, ret)
	} else {
		go func() {
			for m := range ch {
				m.ret <- nil
			}
		}()
	}

	fp := FeedParser{parser: gofeed.NewParser(), messages: ch, sem: make(chan int, 10), finished: make(chan int)}
	go func() {
		for name, accountConfig := range conf.Accounts {
			fp.sem <- 1
			account, exist := hist.Feeds[name]
			if !exist {
				account = new(RSSFeed)
				account.GUIDList = make(map[string]struct{})
				hist.Feeds[name] = account
			} else if len(account.GUIDList) == 0 {
				account.GUIDList = make(map[string]struct{})
			}
			account.config = accountConfig
			go fetchFeed(fp, name, account, conf)
		}
	}()
	for range conf.Accounts {
		<-fp.finished
	}
	ret <- hist.Write()
}
<|endoftext|>


<commit_before>package main

/*
loadTinzenite loads an existing Tinzenite directory and runs it.
*/import (
	"fmt"
	"github.com/tinzenite/bootstrap"
	"github.com/tinzenite/core"
	"github.com/tinzenite/shared"
	"log"
	"os"
	"os/signal"
	"time"
)

func bootstrapTinzenite(path string) {
	var boot *bootstrap.Bootstrap
	var err error
	done := make(chan bool)
	// if tinzenite OR encrypted we can just load the previous bootstrap
	if shared.IsTinzenite(path) || shared.IsEncrypted(path) {
		boot, err = bootstrap.Load(path, func() {
			done <- true
		})
		if err != nil {
			logMain("Bootstrap load error:", err.Error())
			return
		}
	} else {
		// ask whether this is supposed to be a trusted peer
		question := shared.CreateYesNo("Is this a TRUSTED peer?")
		trusted := question.Ask() > 0
		// get peer name
		peerName := shared.GetString("Enter the peer name for this Tinzenite directory:")
		// get address to connect to BEFORE starting boot to avoid terminal clutter
		address := shared.GetString("Please enter the address of the peer to connect to:")
		// build object
		boot, err = bootstrap.Create(path, peerName, trusted, func() {
			done <- true
		})
		if err != nil {
			logMain("Bootstrap create error:", err.Error())
			return
		}
		// connect to:
		err = boot.Start(address)
		if err != nil {
			logMain("Bootstrap start error:", err.Error())
			// return because we don't want to store a faulty bootstrap
			return
		}
		// if everything ok, try storing this bootstrap
		err = boot.Store()
		if err != nil {
			logMain("Bootstrap store error:", err.Error())
		}
	}
	// print information
	address, _ := boot.Address()
	fmt.Printf("Bootstrapping.\nID: %s\n", address)
	// wait for successful bootstrap
	<-done
	log.Println("Closing Bootstrap.")
	// this is required before closing boot because ToxCore may still need to
	// notify the other client that the file transfers are complete - this can
	// take a few iterations, so we delay for a second to give it time to do that.
	<-time.Tick(1 * time.Second)
	// manually close boot if we're done! It won't close itself!
	boot.Close()
	// IF w bootstrapped an encrypted peer, write that to the log and quit.
	if !boot.IsTrusted() {
		fmt.Println("Bootstrapping was successful. Run server to start encrypted peer.")
		return
	}
	// continue to executing the directory
	fmt.Println("Bootstrapping was successful. Loading Tinzenite.")
	// load tinzenite with password
	password := getPassword()
	loadTinzenite(path, password)
}

func createTinzenite(path string) {
	if shared.IsTinzenite(path) {
		logMain("Directory is already a valid Tinzenite directory!")
		return
	}
	// get options
	peerName := shared.GetString("Enter the peer name for this Tinzenite directory:")
	userName := shared.GetString("Enter your username:")
	password := shared.GetString("Enter a password for this Tinzenite network:")
	relPath := shared.CreatePathRoot(path)
	tinzenite, err := core.CreateTinzenite(relPath.LastElement(), relPath.FullPath(), peerName, userName, password)
	if err != nil {
		logMain("Creation error:", err.Error())
		return
	}
	err = tinzenite.SyncLocal()
	if err != nil {
		logMain("Initial model sync error:", err.Error())
	}
	// run tinzenite until killed
	runTinzenite(tinzenite)
}

func loadTinzenite(path, password string) {
	if !shared.IsTinzenite(path) {
		logMain("Directory is not a valid Tinzenite directory!")
		return
	}
	tinzenite, err := core.LoadTinzenite(path, password)
	if err != nil {
		// TODO catch wrong password and allow retry
		logMain("Loading error:", err.Error())
		return
	}
	err = tinzenite.SyncLocal()
	if err != nil {
		logMain("Initial model sync error:", err.Error())
	}
	// run tinzenite until killed
	runTinzenite(tinzenite)
}

/*
runTinzenite runs the given Tinzenite instance.
*/
func runTinzenite(t *core.Tinzenite) {
	// do this here so that it is guaranteed to be set
	t.RegisterPeerValidation(func(address string, wantsTrust bool) {
		var allow bool
		if wantsTrust {
			question := shared.CreateYesNo("Add peer " + address[:8] + " as TRUSTED peer?")
			allow = question.Ask() > 0
		} else {
			question := shared.CreateYesNo("Add peer " + address[:8] + " as ENCRYPTED peer?")
			allow = question.Ask() > 0
		}
		if !allow {
			log.Println("Tin: will not add peer, as requested.")
			return
		}
		// allow peer
		err := t.AllowPeer(address)
		if err != nil {
			log.Println("Tinzenite: failed to allow peer:", err)
		}
		log.Println("Tin: will allow peer, as requested.")
	})
	// print important info
	address, _ := t.Address()
	fmt.Printf("Running peer <%s>.\nID: %s\n", t.Name(), address)
	// build ticks only once instead of every time
	// FIXME: for now using prime numbers to keep them from all ticking at the same time
	tickUpdate := time.Tick(time.Duration(7) * time.Second)
	tickRemote := time.Tick(time.Duration(29) * time.Second)
	tickEncrypted := time.Tick(time.Duration(53) * time.Second)
	// prepare quitting via ctrl-c
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	// loop until close
	for {
		select {
		case <-tickUpdate:
			err := t.SyncLocal()
			if err != nil {
				logMain("SyncLocal error:", err.Error())
			}
		case <-tickRemote:
			err := t.SyncRemote()
			if err != nil {
				logMain("SyncRemote error:", err.Error())
			}
		case <-tickEncrypted:
			err := t.SyncEncrypted()
			if err != nil {
				logMain("SyncEncrypted error:", err.Error())
			}
		case <-c:
			// on interrupt close tinzenite
			t.Close()
			return
		} // select
	} // for
}
<commit_msg>small improvement to bootstrap behavior<commit_after>package main

/*
loadTinzenite loads an existing Tinzenite directory and runs it.
*/import (
	"fmt"
	"github.com/tinzenite/bootstrap"
	"github.com/tinzenite/core"
	"github.com/tinzenite/shared"
	"log"
	"os"
	"os/signal"
	"time"
)

func bootstrapTinzenite(path string) {
	var boot *bootstrap.Bootstrap
	var err error
	// done and onSuccess are used to determine when a bootstrap has completed
	done := make(chan bool, 1)
	onSuccess := func() { done <- true }
	// if tinzenite OR encrypted we can just load the previous bootstrap
	if shared.IsTinzenite(path) || shared.IsEncrypted(path) {
		boot, err = bootstrap.Load(path, onSuccess)
		if err != nil {
			logMain("Bootstrap load error:", err.Error())
			return
		}
	} else {
		// ask whether this is supposed to be a trusted peer
		question := shared.CreateYesNo("Is this a TRUSTED peer?")
		trusted := question.Ask() > 0
		// get peer name
		peerName := shared.GetString("Enter the peer name for this Tinzenite directory:")
		// get address to connect to BEFORE starting boot to avoid terminal clutter
		address := shared.GetString("Please enter the address of the peer to connect to:")
		// build object
		boot, err = bootstrap.Create(path, peerName, trusted, onSuccess)
		if err != nil {
			logMain("Bootstrap create error:", err.Error())
			return
		}
		// connect to:
		err = boot.Start(address)
		if err != nil {
			logMain("Bootstrap start error:", err.Error())
			// return because we don't want to store a faulty bootstrap
			return
		}
		// if everything ok, try storing this bootstrap
		err = boot.Store()
		if err != nil {
			logMain("Bootstrap store error:", err.Error())
		}
	}
	// print information
	address, _ := boot.Address()
	fmt.Printf("Bootstrapping.\nID: %s\n", address)
	// wait for successful bootstrap
	<-done
	log.Println("Closing Bootstrap.")
	// this is required before closing boot because ToxCore may still need to
	// notify the other client that the file transfers are complete - this can
	// take a few iterations, so we delay for a second to give it time to do that.
	<-time.Tick(1 * time.Second)
	// manually close boot if we're done! It won't close itself!
	boot.Close()
	// IF w bootstrapped an encrypted peer, write that to the log and quit.
	if !boot.IsTrusted() {
		fmt.Println("Bootstrapping was successful. Run server to start encrypted peer.")
		return
	}
	// continue to executing the directory
	fmt.Println("Bootstrapping was successful. Loading Tinzenite.")
	// load tinzenite with password
	password := getPassword()
	loadTinzenite(path, password)
}

func createTinzenite(path string) {
	if shared.IsTinzenite(path) {
		logMain("Directory is already a valid Tinzenite directory!")
		return
	}
	// get options
	peerName := shared.GetString("Enter the peer name for this Tinzenite directory:")
	userName := shared.GetString("Enter your username:")
	password := shared.GetString("Enter a password for this Tinzenite network:")
	relPath := shared.CreatePathRoot(path)
	tinzenite, err := core.CreateTinzenite(relPath.LastElement(), relPath.FullPath(), peerName, userName, password)
	if err != nil {
		logMain("Creation error:", err.Error())
		return
	}
	err = tinzenite.SyncLocal()
	if err != nil {
		logMain("Initial model sync error:", err.Error())
	}
	// run tinzenite until killed
	runTinzenite(tinzenite)
}

func loadTinzenite(path, password string) {
	if !shared.IsTinzenite(path) {
		logMain("Directory is not a valid Tinzenite directory!")
		return
	}
	tinzenite, err := core.LoadTinzenite(path, password)
	if err != nil {
		// TODO catch wrong password and allow retry
		logMain("Loading error:", err.Error())
		return
	}
	err = tinzenite.SyncLocal()
	if err != nil {
		logMain("Initial model sync error:", err.Error())
	}
	// run tinzenite until killed
	runTinzenite(tinzenite)
}

/*
runTinzenite runs the given Tinzenite instance.
*/
func runTinzenite(t *core.Tinzenite) {
	// do this here so that it is guaranteed to be set
	t.RegisterPeerValidation(func(address string, wantsTrust bool) {
		var allow bool
		if wantsTrust {
			question := shared.CreateYesNo("Add peer " + address[:8] + " as TRUSTED peer?")
			allow = question.Ask() > 0
		} else {
			question := shared.CreateYesNo("Add peer " + address[:8] + " as ENCRYPTED peer?")
			allow = question.Ask() > 0
		}
		if !allow {
			log.Println("Tin: will not add peer, as requested.")
			return
		}
		// allow peer
		err := t.AllowPeer(address)
		if err != nil {
			log.Println("Tinzenite: failed to allow peer:", err)
		}
		log.Println("Tin: will allow peer, as requested.")
	})
	// print important info
	address, _ := t.Address()
	fmt.Printf("Running peer <%s>.\nID: %s\n", t.Name(), address)
	// build ticks only once instead of every time
	// FIXME: for now using prime numbers to keep them from all ticking at the same time
	tickUpdate := time.Tick(time.Duration(7) * time.Second)
	tickRemote := time.Tick(time.Duration(29) * time.Second)
	tickEncrypted := time.Tick(time.Duration(53) * time.Second)
	// prepare quitting via ctrl-c
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	// loop until close
	for {
		select {
		case <-tickUpdate:
			err := t.SyncLocal()
			if err != nil {
				logMain("SyncLocal error:", err.Error())
			}
		case <-tickRemote:
			err := t.SyncRemote()
			if err != nil {
				logMain("SyncRemote error:", err.Error())
			}
		case <-tickEncrypted:
			err := t.SyncEncrypted()
			if err != nil {
				logMain("SyncEncrypted error:", err.Error())
			}
		case <-c:
			// on interrupt close tinzenite
			t.Close()
			return
		} // select
	} // for
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/client"
	"github.com/loadimpact/speedboat/lib"
	"github.com/loadimpact/speedboat/simple"
	"gopkg.in/urfave/cli.v1"
	"math"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	TypeAuto = "auto"
	TypeURL  = "url"
	TypeJS   = "js"
)

var (
	ErrUnknownType = errors.New("Unable to infer type from argument; specify with -t/--type")
	ErrInvalidType = errors.New("Invalid type specified, see --help")
)

var commandRun = cli.Command{
	Name:      "run",
	Usage:     "Starts running a load test",
	ArgsUsage: "url|filename",
	Flags: []cli.Flag{
		cli.Int64Flag{
			Name:  "vus, u",
			Usage: "virtual users to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "test duration, 0 to run until cancelled",
			Value: 10 * time.Second,
		},
		cli.Int64Flag{
			Name:  "prepare, p",
			Usage: "VUs to prepare (but not start)",
			Value: 0,
		},
		cli.StringFlag{
			Name:  "type, t",
			Usage: "input type, one of: auto, url, js",
			Value: "auto",
		},
	},
	Action: actionRun,
}

func guessType(filename string) string {
	switch {
	case strings.Contains(filename, "://"):
		return TypeURL
	case strings.HasSuffix(filename, ".js"):
		return TypeJS
	default:
		return ""
	}
}

func makeRunner(filename, t string) (lib.Runner, error) {
	if t == TypeAuto {
		t = guessType(filename)
	}

	switch t {
	case TypeAuto:
		return makeRunner(filename, t)
	case "":
		return nil, ErrUnknownType
	case TypeURL:
		return simple.New(filename)
	default:
		return nil, ErrInvalidType
	}
}

func actionRun(cc *cli.Context) error {
	wg := sync.WaitGroup{}

	args := cc.Args()
	if len(args) != 1 {
		return cli.NewExitError("Wrong number of arguments!", 1)
	}

	// Collect arguments
	addr := cc.GlobalString("address")

	duration := cc.Duration("duration")
	if duration == 0 {
		duration = time.Duration(math.MaxInt64)
	}

	vus := cc.Int64("vus")

	prepared := cc.Int64("prepare")
	if prepared == 0 {
		prepared = vus
	}

	// Make the Runner
	filename := args[0]
	runnerType := cc.String("type")
	runner, err := makeRunner(filename, runnerType)
	if err != nil {
		log.WithError(err).Error("Couldn't create a runner")
	}

	// Make the Engine
	engine := &lib.Engine{
		Runner: runner,
	}
	engineC, cancelEngine := context.WithCancel(context.Background())

	// Make the API Server
	api := &APIServer{
		Engine: engine,
		Cancel: cancelEngine,
		Info: lib.Info{
			Version: cc.App.Version,
		},
	}
	apiC, cancelAPI := context.WithCancel(context.Background())

	// Make the Client
	cl, err := client.New(addr)
	if err != nil {
		log.WithError(err).Error("Couldn't make a client; is the address valid?")
		return err
	}

	// Run the engine and API server in the background
	wg.Add(2)
	go func() {
		defer func() {
			log.Debug("Engine terminated")
			wg.Done()
		}()
		log.WithField("prepared", prepared).Debug("Starting engine...")
		if err := engine.Run(engineC, prepared); err != nil {
			log.WithError(err).Error("Engine Error")
		}
	}()
	go func() {
		defer func() {
			log.Debug("API Server terminated")
			wg.Done()
		}()
		log.WithField("addr", addr).Debug("API Server starting...")
		api.Run(apiC, addr)
	}()

	// Wait for the API server to come online
	startTime := time.Now()
	for {
		if err := cl.Ping(); err != nil {
			if time.Since(startTime) < 1*time.Second {
				log.WithError(err).Debug("Waiting for API server to start...")
				time.Sleep(1 * time.Millisecond)
			} else {
				log.WithError(err).Warn("Connection to API server failed; retrying...")
				time.Sleep(1 * time.Second)
			}
			continue
		}
		break
	}

	// Scale the test up to the desired VU count
	if vus > 0 {
		log.WithField("vus", vus).Debug("Scaling test...")
		if err := cl.Scale(vus); err != nil {
			log.WithError(err).Error("Couldn't scale test")
		}
	}

	// Wait for a signal or timeout before shutting down
	quit := make(chan os.Signal)
	signal.Notify(quit, os.Interrupt, syscall.SIGTERM)

	log.Debug("Waiting for test to finish")
	select {
	case <-time.After(duration):
		log.Debug("Duration expired; shutting down...")
	case sig := <-quit:
		log.WithField("signal", sig).Debug("Signal received; shutting down...")
	}

	// Shut down the API server and engine, wait for them to terminate before exiting
	cancelAPI()
	cancelEngine()
	wg.Wait()

	return nil
}
<commit_msg>[feat] Flag to quit on test completion<commit_after>package main

import (
	"context"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/client"
	"github.com/loadimpact/speedboat/lib"
	"github.com/loadimpact/speedboat/simple"
	"gopkg.in/urfave/cli.v1"
	"math"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	TypeAuto = "auto"
	TypeURL  = "url"
	TypeJS   = "js"
)

var (
	ErrUnknownType = errors.New("Unable to infer type from argument; specify with -t/--type")
	ErrInvalidType = errors.New("Invalid type specified, see --help")
)

var commandRun = cli.Command{
	Name:      "run",
	Usage:     "Starts running a load test",
	ArgsUsage: "url|filename",
	Flags: []cli.Flag{
		cli.Int64Flag{
			Name:  "vus, u",
			Usage: "virtual users to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "test duration, 0 to run until cancelled",
			Value: 10 * time.Second,
		},
		cli.Int64Flag{
			Name:  "prepare, p",
			Usage: "VUs to prepare (but not start)",
			Value: 0,
		},
		cli.StringFlag{
			Name:  "type, t",
			Usage: "input type, one of: auto, url, js",
			Value: "auto",
		},
		cli.BoolFlag{
			Name:  "quit, q",
			Usage: "quit immediately on test completion",
		},
	},
	Action: actionRun,
}

func guessType(filename string) string {
	switch {
	case strings.Contains(filename, "://"):
		return TypeURL
	case strings.HasSuffix(filename, ".js"):
		return TypeJS
	default:
		return ""
	}
}

func makeRunner(filename, t string) (lib.Runner, error) {
	if t == TypeAuto {
		t = guessType(filename)
	}

	switch t {
	case TypeAuto:
		return makeRunner(filename, t)
	case "":
		return nil, ErrUnknownType
	case TypeURL:
		return simple.New(filename)
	default:
		return nil, ErrInvalidType
	}
}

func actionRun(cc *cli.Context) error {
	wg := sync.WaitGroup{}

	args := cc.Args()
	if len(args) != 1 {
		return cli.NewExitError("Wrong number of arguments!", 1)
	}

	// Collect arguments
	addr := cc.GlobalString("address")

	duration := cc.Duration("duration")
	if duration == 0 {
		duration = time.Duration(math.MaxInt64)
	}

	vus := cc.Int64("vus")

	prepared := cc.Int64("prepare")
	if prepared == 0 {
		prepared = vus
	}

	quit := cc.Bool("quit")

	// Make the Runner
	filename := args[0]
	runnerType := cc.String("type")
	runner, err := makeRunner(filename, runnerType)
	if err != nil {
		log.WithError(err).Error("Couldn't create a runner")
	}

	// Make the Engine
	engine := &lib.Engine{
		Runner: runner,
	}
	engineC, cancelEngine := context.WithCancel(context.Background())

	// Make the API Server
	api := &APIServer{
		Engine: engine,
		Cancel: cancelEngine,
		Info: lib.Info{
			Version: cc.App.Version,
		},
	}
	apiC, cancelAPI := context.WithCancel(context.Background())

	// Make the Client
	cl, err := client.New(addr)
	if err != nil {
		log.WithError(err).Error("Couldn't make a client; is the address valid?")
		return err
	}

	// Run the engine and API server in the background
	wg.Add(2)
	go func() {
		defer func() {
			log.Debug("Engine terminated")
			if quit {
				log.Debug("Quit requested; terminating API server...")
				cancelAPI()
			}
			wg.Done()
		}()
		log.WithField("prepared", prepared).Debug("Starting engine...")
		if err := engine.Run(engineC, prepared); err != nil {
			log.WithError(err).Error("Engine Error")
		}
	}()
	go func() {
		defer func() {
			log.Debug("API Server terminated")
			wg.Done()
		}()
		log.WithField("addr", addr).Debug("API Server starting...")
		api.Run(apiC, addr)
	}()

	// Wait for the API server to come online
	startTime := time.Now()
	for {
		if err := cl.Ping(); err != nil {
			if time.Since(startTime) < 1*time.Second {
				log.WithError(err).Debug("Waiting for API server to start...")
				time.Sleep(1 * time.Millisecond)
			} else {
				log.WithError(err).Warn("Connection to API server failed; retrying...")
				time.Sleep(1 * time.Second)
			}
			continue
		}
		break
	}

	// Scale the test up to the desired VU count
	if vus > 0 {
		log.WithField("vus", vus).Debug("Scaling test...")
		if err := cl.Scale(vus); err != nil {
			log.WithError(err).Error("Couldn't scale test")
		}
	}

	// Wait for a signal or timeout before shutting down
	signals := make(chan os.Signal)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM)

	log.Debug("Waiting for test to finish")
	select {
	case <-time.After(duration):
		log.Debug("Duration expired; shutting down...")
	case sig := <-signals:
		log.WithField("signal", sig).Debug("Signal received; shutting down...")
	}

	// Shut down the API server and engine, wait for them to terminate before exiting
	cancelAPI()
	cancelEngine()
	wg.Wait()

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Hajime Hoshi
//
// 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 ebiten

import (
	"sync/atomic"

	"github.com/hajimehoshi/ebiten/internal/clock"
	"github.com/hajimehoshi/ebiten/internal/loop"
	"github.com/hajimehoshi/ebiten/internal/ui"
)

// FPS represents how many times game updating happens in a second (60).
const FPS = clock.FPS

// CurrentFPS returns the current number of frames per second of rendering.
//
// This function is concurrent-safe.
//
// This value represents how many times rendering happens in 1/60 second and
// NOT how many times logical game updating (a passed function to Run) happens.
// Note that logical game updating is assured to happen 60 times in a second
// as long as the screen is active.
func CurrentFPS() float64 {
	return loop.CurrentFPS()
}

var (
	isRunningSlowly = int32(0)
)

func setRunningSlowly(slow bool) {
	v := int32(0)
	if slow {
		v = 1
	}
	atomic.StoreInt32(&isRunningSlowly, v)
}

// IsRunningSlowly returns true if the game is running too slowly to keep 60 FPS of rendering.
// The game screen is not updated when IsRunningSlowly is true.
// It is recommended to skip heavy processing, especially drawing, when IsRunningSlowly is true.
//
// This function is concurrent-safe.
func IsRunningSlowly() bool {
	return atomic.LoadInt32(&isRunningSlowly) != 0
}

var theGraphicsContext atomic.Value

// Run runs the game.
// f is a function which is called at every frame.
// The argument (*Image) is the render target that represents the screen.
//
// Run must be called from the main thread.
// Note that ebiten bounds the main goroutine to the main OS thread by runtime.LockOSThread.
//
// The given function f is guaranteed to be called 60 times a second
// even if a rendering frame is skipped.
// f is not called when the screen is not shown.
//
// The given scale is ignored on fullscreen mode.
//
// Run returns error when 1) OpenGL error happens, or 2) f returns error.
// In the case of 2), Run returns the same error.
//
// The size unit is device-independent pixel.
func Run(f func(*Image) error, width, height int, scale float64, title string) error {
	ch := make(chan error)
	go func() {
		g := newGraphicsContext(f)
		theGraphicsContext.Store(g)
		if err := loop.Run(g, width, height, scale, title); err != nil {
			ch <- err
		}
		close(ch)
	}()
	// TODO: Use context in Go 1.7?
	if err := ui.RunMainThreadLoop(ch); err != nil {
		return err
	}
	return nil
}

// RunWithoutMainLoop runs the game, but don't call the loop on the main (UI) thread.
// Different from Run, this function returns immediately.
//
// Typically, Ebiten users don't have to call this directly.
// Instead, functions in github.com/hajimehoshi/ebiten/mobile module call this.
//
// The size unit is device-independent pixel.
func RunWithoutMainLoop(f func(*Image) error, width, height int, scale float64, title string) <-chan error {
	ch := make(chan error)
	go func() {
		g := newGraphicsContext(f)
		theGraphicsContext.Store(g)
		if err := loop.Run(g, width, height, scale, title); err != nil {
			ch <- err
		}
		close(ch)
	}()
	return ch
}

// SetScreenSize changes the (logical) size of the screen.
// This doesn't affect the current scale of the screen.
//
// Unit is device-independent pixel.
//
// This function is concurrent-safe.
func SetScreenSize(width, height int) {
	if width <= 0 || height <= 0 {
		panic("ebiten: width and height must be positive")
	}
	ui.SetScreenSize(width, height)
}

// SetScreenScale changes the scale of the screen.
//
// This function is concurrent-safe.
func SetScreenScale(scale float64) {
	if scale <= 0 {
		panic("ebiten: scale must be positive")
	}
	ui.SetScreenScale(scale)
}

// ScreenScale returns the current screen scale.
//
// If Run is not called, this returns 0.
//
// This function is concurrent-safe.
func ScreenScale() float64 {
	return ui.ScreenScale()
}

// SetCursorVisibility changes the state of cursor visiblity.
//
// This function is concurrent-safe.
func SetCursorVisibility(visible bool) {
	ui.SetCursorVisibility(visible)
}

// IsScreen returns a boolean value indicating whether
// the current mode is fullscreen or not.
//
// This function is concurrent-safe.
func IsFullscreen() bool {
	return ui.IsFullscreen()
}

// SetFullscreen changes the current mode to fullscreen or not.
//
// On fullscreen mode, the game screen is automatically enlarged
// to fit with the monitor. The current scale value is ignored.
//
// Ebiten uses 'windowed' fullscreen mode, which doesn't change
// your monitor's resolution.
//
// This function is concurrent-safe.
func SetFullscreen(fullscreen bool) {
	ui.SetFullscreen(fullscreen)
}
<commit_msg>ui: Fix comments<commit_after>// Copyright 2014 Hajime Hoshi
//
// 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 ebiten

import (
	"sync/atomic"

	"github.com/hajimehoshi/ebiten/internal/clock"
	"github.com/hajimehoshi/ebiten/internal/loop"
	"github.com/hajimehoshi/ebiten/internal/ui"
)

// FPS represents how many times game updating happens in a second (60).
const FPS = clock.FPS

// CurrentFPS returns the current number of frames per second of rendering.
//
// This function is concurrent-safe.
//
// This value represents how many times rendering happens in 1/60 second and
// NOT how many times logical game updating (a passed function to Run) happens.
// Note that logical game updating is assured to happen 60 times in a second
// as long as the screen is active.
func CurrentFPS() float64 {
	return loop.CurrentFPS()
}

var (
	isRunningSlowly = int32(0)
)

func setRunningSlowly(slow bool) {
	v := int32(0)
	if slow {
		v = 1
	}
	atomic.StoreInt32(&isRunningSlowly, v)
}

// IsRunningSlowly returns true if the game is running too slowly to keep 60 FPS of rendering.
// The game screen is not updated when IsRunningSlowly is true.
// It is recommended to skip heavy processing, especially drawing, when IsRunningSlowly is true.
//
// This function is concurrent-safe.
func IsRunningSlowly() bool {
	return atomic.LoadInt32(&isRunningSlowly) != 0
}

var theGraphicsContext atomic.Value

// Run runs the game.
// f is a function which is called at every frame.
// The argument (*Image) is the render target that represents the screen.
//
// Run must be called from the main thread.
// Note that ebiten bounds the main goroutine to the main OS thread by runtime.LockOSThread.
//
// The given function f is guaranteed to be called 60 times a second
// even if a rendering frame is skipped.
// f is not called when the screen is not shown.
//
// The given scale is ignored on fullscreen mode.
//
// Run returns error when 1) OpenGL error happens, or 2) f returns error.
// In the case of 2), Run returns the same error.
//
// The size unit is device-independent pixel.
func Run(f func(*Image) error, width, height int, scale float64, title string) error {
	ch := make(chan error)
	go func() {
		g := newGraphicsContext(f)
		theGraphicsContext.Store(g)
		if err := loop.Run(g, width, height, scale, title); err != nil {
			ch <- err
		}
		close(ch)
	}()
	// TODO: Use context in Go 1.7?
	if err := ui.RunMainThreadLoop(ch); err != nil {
		return err
	}
	return nil
}

// RunWithoutMainLoop runs the game, but don't call the loop on the main (UI) thread.
// Different from Run, this function returns immediately.
//
// Typically, Ebiten users don't have to call this directly.
// Instead, functions in github.com/hajimehoshi/ebiten/mobile module call this.
//
// The size unit is device-independent pixel.
func RunWithoutMainLoop(f func(*Image) error, width, height int, scale float64, title string) <-chan error {
	ch := make(chan error)
	go func() {
		g := newGraphicsContext(f)
		theGraphicsContext.Store(g)
		if err := loop.Run(g, width, height, scale, title); err != nil {
			ch <- err
		}
		close(ch)
	}()
	return ch
}

// SetScreenSize changes the (logical) size of the screen.
// This doesn't affect the current scale of the screen.
//
// Unit is device-independent pixel.
//
// This function is concurrent-safe.
func SetScreenSize(width, height int) {
	if width <= 0 || height <= 0 {
		panic("ebiten: width and height must be positive")
	}
	ui.SetScreenSize(width, height)
}

// SetScreenScale changes the scale of the screen.
//
// This function is concurrent-safe.
func SetScreenScale(scale float64) {
	if scale <= 0 {
		panic("ebiten: scale must be positive")
	}
	ui.SetScreenScale(scale)
}

// ScreenScale returns the current screen scale.
//
// If Run is not called, this returns 0.
//
// This function is concurrent-safe.
func ScreenScale() float64 {
	return ui.ScreenScale()
}

// SetCursorVisibility changes the state of cursor visiblity.
//
// This function is concurrent-safe.
func SetCursorVisibility(visible bool) {
	ui.SetCursorVisibility(visible)
}

// IsScreen returns a boolean value indicating whether
// the current mode is fullscreen or not.
//
// This function is concurrent-safe.
func IsFullscreen() bool {
	return ui.IsFullscreen()
}

// SetFullscreen changes the current mode to fullscreen or not.
//
// On fullscreen mode, the game screen is automatically enlarged
// to fit with the monitor. The current scale value is ignored.
//
// On desktops, Ebiten uses 'windowed' fullscreen mode, which doesn't change
// your monitor's resolution.
//
// On browsers, the game screen is resized to fit with the body element (client) size.
// Additionally, the game screen is automatically resized when the body element is resized.
//
// SetFullscreen doesn't work on mobiles.
//
// This function is concurrent-safe.
func SetFullscreen(fullscreen bool) {
	ui.SetFullscreen(fullscreen)
}
<|endoftext|>


<commit_before>// Author:          Michael John
// Copyright:       2014-2017 Crown Equipment Corp.

// Package semaphore provides an interface to named userspace semaphores.
package semaphore

import (
	"errors"
	"syscall"
	"time"
	"unsafe"
)

// #cgo LDFLAGS: -pthread
// #include <stdlib.h>
// #include <fcntl.h>
// #include <sys/stat.h>
// #include <sys/types.h>
// #include <semaphore.h>
// #include <time.h>
// sem_t* Go_sem_open(const char *name, int oflag, mode_t mode, unsigned int value)
// {
//		return sem_open(name, oflag, mode, value);
// }
import "C"

type Semaphore struct {
	sem  *C.sem_t //semaphore returned by sem_open
	name string   //name of semaphore
}

func (s *Semaphore) isSemaphoreInitialized() (bool, error) {
	if s.sem == nil {
		return false, errors.New("Not a valid semaphore")
	}
	return true, nil
}

// Open creates a new POSIX semaphore or opens an existing semaphore.
// The semaphore is identified by name. The mode argument specifies the permissions
// to be placed on the new semaphore. The value argument specifies the initial
// value for the new semaphore. If the named semaphore already exist, mode and
// value are ignored.
// For details see sem_overview(7).
func (s *Semaphore) Open(name string, mode, value uint32) error {
	s.name = name
	n := C.CString(name)

	var err error
	s.sem, err = C.Go_sem_open(n, syscall.O_CREAT, C.mode_t(mode), C.uint(value))
	C.free(unsafe.Pointer(n))
	if s.sem == nil {
		return err
	}
	return nil
}

// Close closes the named semaphore, allowing any resources that the system has
// allocated to the calling process for this semaphore to be freed.
func (s *Semaphore) Close() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_close(s.sem)
	if ret != 0 {
		return err
	}
	s.sem = nil
	return nil
}

// GetValue returns the current value of the semaphore.
func (s *Semaphore) GetValue() (int, error) {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return 0, err
	}

	var val C.int
	ret, err := C.sem_getvalue(s.sem, &val)
	if ret != 0 {
		return int(ret), err
	}
	return int(val), nil
}

// Post increments the semaphore.
func (s *Semaphore) Post() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_post(s.sem)
	if ret != 0 {
		return err
	}
	return nil
}

// Wait decrements the semaphore. If the semaphore's value is greater than zero,
// then the decrement proceeds, and the function returns, immediately. If the
// semaphore currently has the value zero, then the call blocks until either
// it becomes possible to perform the decrement, or a signal interrupts the call.
func (s *Semaphore) Wait() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_wait(s.sem)
	if ret != 0 {
		return err
	}
	return nil
}

// TryWait is the same as Wait(), except that if the decrement cannot be immediately
// performed, then the call returns an error instead of blocking.
func (s *Semaphore) TryWait() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_trywait(s.sem)
	if ret != 0 {
		return err
	}
	return nil
}

// TimedWait is the same as Wait(), except that d specifies a limit on the
// amount of time that the call should block if the decrement cannot be
// immediately performed.
func (s *Semaphore) TimedWait(d time.Duration) error {
	if err := s.TryWait(); err == nil {
		// success
		return nil
	}
	time.Sleep(d)
	if err := s.TryWait(); err == nil {
		// success
		return nil
	}
	return errors.New("The call timed out before the semaphore could be locked")
}

// Unlink removes the named semaphore. The semaphore name is removed immediately.
// The semaphore is destroyed once all other processes that have the semaphore
// open close it.
func (s *Semaphore) Unlink() error {
	name := C.CString(s.name)
	ret, err := C.sem_unlink(name)
	C.free(unsafe.Pointer(name))
	if ret != 0 {
		return err
	}
	return nil
}
<commit_msg>Minor updates.<commit_after>// Project:         C1515
// Author:          Michael John
// Copyright:       2014-2018 Crown Equipment Corp.

// +build linux,cgo

// Package semaphore provides an interface to named userspace semaphores.
package semaphore

import (
	"errors"
	"syscall"
	"time"
	"unsafe"
)

// #cgo LDFLAGS: -pthread
// #include <stdlib.h>
// #include <fcntl.h>
// #include <sys/stat.h>
// #include <sys/types.h>
// #include <semaphore.h>
// #include <time.h>
// #ifndef GO_SEM_LIB_
// #define GO_SEM_LIB_
// sem_t* Go_sem_open(const char *name, int oflag, mode_t mode, unsigned int value)
// {
//		return sem_open(name, oflag, mode, value);
// }
// #endif
import "C"

type Semaphore struct {
	sem  *C.sem_t //semaphore returned by sem_open
	name string   //name of semaphore
}

func (s *Semaphore) isSemaphoreInitialized() (bool, error) {
	if s.sem == nil {
		return false, errors.New("Not a valid semaphore")
	}
	return true, nil
}

// Open creates a new POSIX semaphore or opens an existing semaphore.
// The semaphore is identified by name. The mode argument specifies the permissions
// to be placed on the new semaphore. The value argument specifies the initial
// value for the new semaphore. If the named semaphore already exist, mode and
// value are ignored.
// For details see sem_overview(7).
func (s *Semaphore) Open(name string, mode, value uint32) error {
	s.name = name
	n := C.CString(name)

	var err error
	s.sem, err = C.Go_sem_open(n, syscall.O_CREAT, C.mode_t(mode), C.uint(value))
	C.free(unsafe.Pointer(n))
	if s.sem == nil {
		return err
	}
	return nil
}

// Close closes the named semaphore, allowing any resources that the system has
// allocated to the calling process for this semaphore to be freed.
func (s *Semaphore) Close() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_close(s.sem)
	if ret != 0 {
		return err
	}
	s.sem = nil
	return nil
}

// GetValue returns the current value of the semaphore.
func (s *Semaphore) GetValue() (int, error) {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return 0, err
	}

	var val C.int
	ret, err := C.sem_getvalue(s.sem, &val)
	if ret != 0 {
		return int(ret), err
	}
	return int(val), nil
}

// Post increments the semaphore.
func (s *Semaphore) Post() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_post(s.sem)
	if ret != 0 {
		return err
	}
	return nil
}

// Wait decrements the semaphore. If the semaphore's value is greater than zero,
// then the decrement proceeds, and the function returns, immediately. If the
// semaphore currently has the value zero, then the call blocks until either
// it becomes possible to perform the decrement, or a signal interrupts the call.
func (s *Semaphore) Wait() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_wait(s.sem)
	if ret != 0 {
		return err
	}
	return nil
}

// TryWait is the same as Wait(), except that if the decrement cannot be immediately
// performed, then the call returns an error instead of blocking.
func (s *Semaphore) TryWait() error {
	if ok, err := s.isSemaphoreInitialized(); !ok {
		return err
	}

	ret, err := C.sem_trywait(s.sem)
	if ret != 0 {
		return err
	}
	return nil
}

// TimedWait is the same as Wait(), except that d specifies a limit on the
// amount of time that the call should block if the decrement cannot be
// immediately performed.
func (s *Semaphore) TimedWait(d time.Duration) error {
	if err := s.TryWait(); err == nil {
		// success
		return nil
	}
	time.Sleep(d)
	if err := s.TryWait(); err == nil {
		// success
		return nil
	}
	return errors.New("The call timed out before the semaphore could be locked")
}

// Unlink removes the named semaphore. The semaphore name is removed immediately.
// The semaphore is destroyed once all other processes that have the semaphore
// open close it.
func (s *Semaphore) Unlink() error {
	name := C.CString(s.name)
	ret, err := C.sem_unlink(name)
	C.free(unsafe.Pointer(name))
	if ret != 0 {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package sshctl

import (
	"regexp"
)

// Set represents all the hosts read from a particular file.
type Set struct {
	groups map[string]*Group
	file   string // File from which set is read
}

// New reurns a new Group as read from the specified file.  If file is "",
// an empty group will be returned.  If file has no hosts, an empty set is
// returned.  Any changes to the set will be written to the file.
func New(file string) Set {
}

// Group returns a pre-defined group of hosts.
func (s Set) Group(name string) *Group {
}

// AddGroup makes a new group of hosts based on a regular expression and
// returns a pointer to that group.  The group will be saved for later
// retrieval with Group()
func (s *Set) NewGroup(re regexp.Regexp) *Group {
}

//
<commit_msg>Functions to make a new set and host<commit_after>package sshctl

import (
	"os"
	"regexp"
)

// Set represents all the hosts read from a particular file.
type Set struct {
	groups map[string]*Group
	file   string // File from which set is read
}

// New reurns a new Group as read from the specified file.  If file is "",
// an empty group will be returned.  If file has no hosts, an empty set is
// returned.  Any changes to the set will be written to the file.
func New(file string) (Set, error) {
	/* Try to open the file */
	f, err := os.Open(file)
	if err != nil {
		return nil, error
	}
}

// Group returns a pre-defined group of hosts.
func (s Set) Group(name string) *Group {
}

// AddGroup makes a new group of hosts based on a regular expression and
// returns a pointer to that group.  The group will be saved for later
// retrieval with Group()
func (s *Set) NewGroup(re regexp.Regexp) *Group {
}

// New makes a new host, and adds it to the specified groups
func (s *Set) NewHost(user, address string, port uint16, groups []string) Host {
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"runtime"
	"strconv"
	"strings"
	"unicode/utf8"
)

var (
	CACHE = map[string]string{}
	bind  = flag.String("bind", "127.0.0.1:11211", "Address:port to bind to")
	db    = flag.String("db", "talon.db", "path to database")
)

type CacheItem struct {
	Key   string
	Value []byte
}

func main() {
	// NUmber of cpu's to use
	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()

	ioChannel := make(chan CacheItem)
	go ioHandler(ioChannel)

	listener, err := net.Listen("tcp", *bind)
	if err != nil {
		log.Printf("net.Listen error")
		panic("Error listening on 11211: " + err.Error())
	}

	CACHE = make(map[string]string)

	log.Printf("\x1b[32m [*] Listening on:\x1b[0m %s", *bind)

	// Load the cache saved on disk
	loadCache()

	for {
		netconn, err := listener.Accept()
		if err != nil {
			log.Printf("Listener.Accept() error")
			panic("Accept error: " + err.Error())
		}

		go handleConn(netconn, ioChannel)
	}

}

func ioHandler(cs chan CacheItem) {
	for {
		item := <-cs
		CACHE[item.Key] = string(item.Value)
	}
}

func loadCache() {
	n, err := ioutil.ReadFile(*db)
	if err != nil {
		return
	}

	p := bytes.NewBuffer(n)

	dec := gob.NewDecoder(p)
	err = dec.Decode(&CACHE)
	if err != nil {
		return
	}
}

func syncCache() {
	b := new(bytes.Buffer)
	enc := gob.NewEncoder(b)
	err := enc.Encode(CACHE)
	if err != nil {
		log.Printf("Error detected while encoding: %v", err)
		return
	}

	// Write gob object to file

	// open output file
	fo, err := os.Create(*db)
	if err != nil {
		log.Printf("Saving error")
		panic(err)
	}
	// close fo on exit and check for its returned error
	defer func() {
		log.Printf(" [*] Done saving to disk")
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	if _, err := fo.Write(b.Bytes()); err != nil {
		log.Printf("Error writing bytes SAVE")
		panic(err)
	}
}

func handleConn(conn net.Conn, ioHandler chan CacheItem) {
	defer conn.Close()

	reader := bufio.NewReader(conn)
	for {

		// Fetch
		content, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		} else if err != nil {
			fmt.Println(err)
			return
		}

		content = content[:len(content)-2] // Chop \r\n

		// Handle
		parts := strings.Split(content, " ")
		cmd := parts[0]

		switch cmd {
		case "get":
			key := parts[1]
			value, ok := CACHE[key]
			if ok {
				_, err := conn.Write([]uint8("VALUE " + key + " " + string(value) + "\r\n\r\n"))
				if err != nil {
					return
				}
			} else {
				_, err = conn.Write([]uint8("VALUE nil"))
				if err != nil {
					return
				}

			}

			conn.Write([]uint8("\r\n"))
			return

		case "set":
			key := parts[1]

			length := utf8.RuneCountInString(parts[2]) + 120
			val := make([]byte, length)
			val = []byte(parts[2])

			kv := CacheItem{key, val}
			ioHandler <- kv

			_, err := conn.Write([]uint8("STORED\r\n"))
			if err != nil {
				conn.Write([]uint8("ERROR"))
				return
			}
			return

		case "save":
			log.Printf(" [*] Writing CACHE to disk")
			go syncCache()

		case "delete":
			key := parts[1]
			delete(CACHE, key)
			log.Printf(" [*] Deleted [%v] from CACHE", key)
			return

		case "stats":
			stats := strconv.Itoa(len(CACHE))
			_, err = conn.Write([]uint8("KEYS " + stats))
			if err != nil {
				return
			}

		case "exit":
			return
		}
	}
}
<commit_msg>Added autoclose option. As requested<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"runtime"
	"strconv"
	"strings"
	"unicode/utf8"
)

var (
	CACHE     = map[string]string{}
	bind      = flag.String("bind", "127.0.0.1:11211", "Address:port to bind to")
	db        = flag.String("db", "talon.db", "path to database")
	autoclose = flag.Bool("autoclose", false, "Autoclose every connection [true/false]")
)

type CacheItem struct {
	Key   string
	Value []byte
}

func main() {
	// NUmber of cpu's to use
	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()

	ioChannel := make(chan CacheItem)
	go ioHandler(ioChannel)

	listener, err := net.Listen("tcp", *bind)
	if err != nil {
		log.Printf("net.Listen error")
		panic("Error listening on 11211: " + err.Error())
	}

	CACHE = make(map[string]string)

	log.Printf("\x1b[32m [*] Listening on:\x1b[0m %s", *bind)

	// Load the cache saved on disk
	loadCache()

	for {
		netconn, err := listener.Accept()
		if err != nil {
			log.Printf("Listener.Accept() error")
			panic("Accept error: " + err.Error())
		}

		go handleConn(netconn, ioChannel)
	}

}

func ioHandler(cs chan CacheItem) {
	for {
		item := <-cs
		CACHE[item.Key] = string(item.Value)
	}
}

func loadCache() {
	n, err := ioutil.ReadFile(*db)
	if err != nil {
		return
	}

	p := bytes.NewBuffer(n)

	dec := gob.NewDecoder(p)
	err = dec.Decode(&CACHE)
	if err != nil {
		return
	}
}

func syncCache() {
	b := new(bytes.Buffer)
	enc := gob.NewEncoder(b)
	err := enc.Encode(CACHE)
	if err != nil {
		log.Printf("Error detected while encoding: %v", err)
		return
	}

	// Write gob object to file

	// open output file
	fo, err := os.Create(*db)
	if err != nil {
		log.Printf("Saving error")
		panic(err)
	}
	// close fo on exit and check for its returned error
	defer func() {
		log.Printf(" [*] Done saving to disk")
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	if _, err := fo.Write(b.Bytes()); err != nil {
		log.Printf("Error writing bytes SAVE")
		panic(err)
	}
}

func handleConn(conn net.Conn, ioHandler chan CacheItem) {
	defer conn.Close()

	reader := bufio.NewReader(conn)
	for {

		// Fetch
		content, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		} else if err != nil {
			fmt.Println(err)
			return
		}

		content = content[:len(content)-2] // Chop \r\n

		// Handle
		parts := strings.Split(content, " ")
		cmd := parts[0]

		switch cmd {
		case "get":
			key := parts[1]
			value, ok := CACHE[key]
			if ok {
				_, err := conn.Write([]uint8("VALUE " + key + " " + string(value) + "\r\n\r\n"))
				if err != nil {
					return
				}
			} else {
				_, err = conn.Write([]uint8("VALUE nil"))
				if err != nil {
					return
				}

			}

			conn.Write([]uint8("\r\n"))

			if *autoclose == true {
				return
			}

		case "set":
			key := parts[1]

			length := utf8.RuneCountInString(parts[2]) + 120
			val := make([]byte, length)
			val = []byte(parts[2])

			kv := CacheItem{key, val}
			ioHandler <- kv

			_, err := conn.Write([]uint8("STORED\r\n"))
			if err != nil {
				conn.Write([]uint8("ERROR"))
				return
			}

			if *autoclose == true {
				return
			}

		case "save":
			log.Printf(" [*] Writing CACHE to disk")
			go syncCache()

		case "delete":
			key := parts[1]
			delete(CACHE, key)
			log.Printf(" [*] Deleted [%v] from CACHE", key)

			if *autoclose == true {
				return
			}

		case "stats":
			stats := strconv.Itoa(len(CACHE))
			_, err = conn.Write([]uint8("KEYS " + stats))
			if err != nil {
				return
			}

		case "exit":
			return
		}
	}
}
<|endoftext|>


<commit_before>package build

import (
	"fmt"
	"github.com/fatih/color"
	"golang.org/x/crypto/ssh/terminal"
	"os"
	"strings"
	"unicode/utf8"
)

// Flag that tells if we print on console without color
var Grey = false

// Color definitions
var red = color.New(color.FgRed, color.Bold).SprintFunc()
var yellow = color.New(color.FgYellow).SprintFunc()
var green = color.New(color.FgGreen, color.Bold).SprintFunc()

// Print a message
func Message(text string, args ...interface{}) {
	printGrey(text, args...)
}

// Print a title
func Title(text string) {
	length := termWidth() - (4 + utf8.RuneCountInString(text))
	if length < 2 {
		length = 2
	}
	message := fmt.Sprintf("%s %s --", strings.Repeat("-", length), text)
	if Grey {
		printGrey(message)
	} else {
		printColor(yellow(message))
	}
}

// Print OK
func PrintOk() {
	if Grey {
		printGrey("OK")
	} else {
		printColor(green("OK"))
	}
}

// Print ERROR
func PrintError(text string) {
	if Grey {
		printGrey("ERROR %s", text)
	} else {
		printColor("%s %s", red("ERROR"), text)
	}
}

// Print string with arguments in given color
func printColor(format string, fields ...interface{}) {
	fmt.Fprintf(color.Output, format, fields...)
	fmt.Println()
}

// Print string with arguments in grey
func printGrey(format string, fields ...interface{}) {
	fmt.Printf(format, fields...)
	fmt.Println()
}

// Get terminal width
func termWidth() int {
	width, _, err := terminal.GetSize(int(os.Stdout.Fd()))
	if err != nil {
		return 80
	}
	if width <= 0 {
		width = 80
	}
	return width
}
<commit_msg>Using constant for default width<commit_after>package build

import (
	"fmt"
	"github.com/fatih/color"
	"golang.org/x/crypto/ssh/terminal"
	"os"
	"strings"
	"unicode/utf8"
)

const (
	DEFAULT_WIDTH = 80
)

// Flag that tells if we print on console without color
var Grey = false

// Color definitions
var red = color.New(color.FgRed, color.Bold).SprintFunc()
var yellow = color.New(color.FgYellow).SprintFunc()
var green = color.New(color.FgGreen, color.Bold).SprintFunc()

// Print a message
func Message(text string, args ...interface{}) {
	printGrey(text, args...)
}

// Print a title
func Title(text string) {
	length := termWidth() - (4 + utf8.RuneCountInString(text))
	if length < 2 {
		length = 2
	}
	message := fmt.Sprintf("%s %s --", strings.Repeat("-", length), text)
	if Grey {
		printGrey(message)
	} else {
		printColor(yellow(message))
	}
}

// Print OK
func PrintOk() {
	if Grey {
		printGrey("OK")
	} else {
		printColor(green("OK"))
	}
}

// Print ERROR
func PrintError(text string) {
	if Grey {
		printGrey("ERROR %s", text)
	} else {
		printColor("%s %s", red("ERROR"), text)
	}
}

// Print string with arguments in given color
func printColor(format string, fields ...interface{}) {
	fmt.Fprintf(color.Output, format, fields...)
	fmt.Println()
}

// Print string with arguments in grey
func printGrey(format string, fields ...interface{}) {
	fmt.Printf(format, fields...)
	fmt.Println()
}

// Get terminal width
func termWidth() int {
	width, _, err := terminal.GetSize(int(os.Stdout.Fd()))
	if err != nil {
		return DEFAULT_WIDTH
	}
	if width <= 0 {
		width = DEFAULT_WIDTH
	}
	return width
}
<|endoftext|>


<commit_before>package appdeploy

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/eapache/go-resiliency/retrier"
	"github.com/rubenv/kube-appdeploy/kubectl"
)

var CleanTypes = []string{
	"deployment",
	"service",
	"cronjob",
}

type Target interface {
	Prepare(vars *ProcessVariables) error
	Apply(m Manifest, data []byte) error
	Cleanup(items []Manifest) error
}

// ---------- Folder ----------

type FolderTarget struct {
	Path string
}

var _ Target = &FolderTarget{}

func NewFolderTarget(path string) *FolderTarget {
	return &FolderTarget{
		Path: path,
	}
}

func (t *FolderTarget) Prepare(vars *ProcessVariables) error {
	return os.MkdirAll(t.Path, 0755)
}

func (t *FolderTarget) Apply(m Manifest, data []byte) error {
	return ioutil.WriteFile(m.Filename(t.Path), data, 0644)
}

func (t *FolderTarget) Cleanup(items []Manifest) error {
	files, err := ioutil.ReadDir(t.Path)
	if err != nil {
		return err
	}

	filenames := make([]string, 0, len(items))
	for _, item := range items {
		filenames = append(filenames, item.Filename(""))
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		name := file.Name()
		sep := strings.Index(name, "--")
		if sep < 0 {
			continue
		}

		prefix := name[0:sep]
		found := false
		for _, t := range CleanTypes {
			if t == prefix {
				found = true
			}
		}
		if !found {
			continue
		}

		known := false
		for _, f := range filenames {
			if f == name {
				known = true
			}
		}

		if !known {
			err = os.Remove(path.Join(t.Path, name))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// ---------- Kubernetes ----------

type KubernetesTarget struct {
	config         *rest.Config
	client         *kubernetes.Clientset
	kubectl        *kubectl.KubeCtl
	namespace      string
	manageCronjobs bool
}

var _ Target = &KubernetesTarget{}

func NewKubernetesTarget(config *rest.Config) *KubernetesTarget {
	return &KubernetesTarget{
		config: config,
	}
}

func (t *KubernetesTarget) Prepare(vars *ProcessVariables) error {
	client, err := kubernetes.NewForConfig(t.config)
	if err != nil {
		return err
	}

	t.client = client

	// Copy some vars
	t.namespace = vars.Namespace
	t.kubectl = kubectl.NewKubeCtl(t.config, t.namespace)
	t.manageCronjobs = vars.ManageCronjobs

	// Ensure we have the needed namespace
	nsClient := t.client.Core().Namespaces()

	create := false
	_, err = nsClient.Get(t.namespace, metav1.GetOptions{})
	if err != nil {
		ignore := false
		if e, ok := err.(*errors.StatusError); ok {
			if e.ErrStatus.Reason == "NotFound" {
				ignore = true
				create = true
			}
		}
		if !ignore {
			return err
		}
	}
	if create {
		_, err = nsClient.Create(&v1.Namespace{
			ObjectMeta: metav1.ObjectMeta{
				Name: t.namespace,
			},
		})
		if err != nil {
			return err
		}
	}

	// Add the image pull secrets
	if len(vars.ImagePullSecrets) > 0 {
		saClient := t.client.Core().ServiceAccounts(t.namespace)

		var sa *v1.ServiceAccount
		// Account isn't always available right away, but it gets created in the end, just wait for it
		r := retrier.New(retrier.ConstantBackoff(10, 1*time.Second), nil)
		err := r.Run(func() error {
			s, err := saClient.Get("default", metav1.GetOptions{})
			if err != nil {
				return err
			}
			sa = s
			return nil
		})

		secrets := make([]v1.LocalObjectReference, 0)
		for _, s := range vars.ImagePullSecrets {
			secrets = append(secrets, v1.LocalObjectReference{
				Name: s,
			})
		}

		sa.ImagePullSecrets = secrets
		_, err = saClient.Update(sa)
		if err != nil {
			return err
		}
	}

	return nil
}

func (t *KubernetesTarget) Apply(m Manifest, data []byte) error {
	// Temporary fix for https://github.com/kubernetes/kubernetes/issues/35149
	// If a cronjob is applied, an error occurs
	// Thus we delete the cronjob first if it exists
	if m.Kind == "CronJob" {
		out, err := t.runKubeCtl(nil, "get", "cronjob", "-o", "name")
		if err != nil {
			return err
		}
		lines := strings.Split(strings.TrimSpace(out), "\n")
		searchline := fmt.Sprintf("cronjob/%s", m.Metadata.Name)
		for _, line := range lines {
			if line == searchline {
				_, err := t.runKubeCtl(nil, "delete", "cronjob", m.Metadata.Name)
				if err != nil {
					return err
				}
				break
			}
		}
	}
	_, err := t.runKubeCtl(data, "apply", "-f", "-")
	return err
}

func (t *KubernetesTarget) Cleanup(items []Manifest) error {
	for _, ct := range CleanTypes {
		if ct == "cronjob" && !t.manageCronjobs {
			continue
		}
		err := t.cleanType(items, ct)
		if err != nil {
			return err
		}
	}
	return nil
}

func (t *KubernetesTarget) cleanType(items []Manifest, ct string) error {
	out, err := t.runKubeCtl(nil, "get", ct, "-o", "name")
	if err != nil {
		return err
	}

	known := []string{}
	for _, m := range items {
		if strings.ToLower(m.Kind) == ct {
			known = append(known, fmt.Sprintf("%s/%s", ct, m.Metadata.Name))
		}
	}

	lines := strings.Split(strings.TrimSpace(out), "\n")
	for _, line := range lines {
		if line == "" {
			continue
		}

		found := false
		for _, k := range known {
			if line == k {
				found = true
			}
		}

		if !found {
			_, err := t.runKubeCtl(nil, "delete", line)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (t *KubernetesTarget) runKubeCtl(stdin []byte, args ...string) (string, error) {
	return t.kubectl.Run(stdin, args...)
}
<commit_msg>Handle initially-missing service account<commit_after>package appdeploy

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/eapache/go-resiliency/retrier"
	"github.com/rubenv/kube-appdeploy/kubectl"
)

var CleanTypes = []string{
	"deployment",
	"service",
	"cronjob",
}

type Target interface {
	Prepare(vars *ProcessVariables) error
	Apply(m Manifest, data []byte) error
	Cleanup(items []Manifest) error
}

// ---------- Folder ----------

type FolderTarget struct {
	Path string
}

var _ Target = &FolderTarget{}

func NewFolderTarget(path string) *FolderTarget {
	return &FolderTarget{
		Path: path,
	}
}

func (t *FolderTarget) Prepare(vars *ProcessVariables) error {
	return os.MkdirAll(t.Path, 0755)
}

func (t *FolderTarget) Apply(m Manifest, data []byte) error {
	return ioutil.WriteFile(m.Filename(t.Path), data, 0644)
}

func (t *FolderTarget) Cleanup(items []Manifest) error {
	files, err := ioutil.ReadDir(t.Path)
	if err != nil {
		return err
	}

	filenames := make([]string, 0, len(items))
	for _, item := range items {
		filenames = append(filenames, item.Filename(""))
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		name := file.Name()
		sep := strings.Index(name, "--")
		if sep < 0 {
			continue
		}

		prefix := name[0:sep]
		found := false
		for _, t := range CleanTypes {
			if t == prefix {
				found = true
			}
		}
		if !found {
			continue
		}

		known := false
		for _, f := range filenames {
			if f == name {
				known = true
			}
		}

		if !known {
			err = os.Remove(path.Join(t.Path, name))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// ---------- Kubernetes ----------

type KubernetesTarget struct {
	config         *rest.Config
	client         *kubernetes.Clientset
	kubectl        *kubectl.KubeCtl
	namespace      string
	manageCronjobs bool
}

var _ Target = &KubernetesTarget{}

func NewKubernetesTarget(config *rest.Config) *KubernetesTarget {
	return &KubernetesTarget{
		config: config,
	}
}

func (t *KubernetesTarget) Prepare(vars *ProcessVariables) error {
	client, err := kubernetes.NewForConfig(t.config)
	if err != nil {
		return err
	}

	t.client = client

	// Copy some vars
	t.namespace = vars.Namespace
	t.kubectl = kubectl.NewKubeCtl(t.config, t.namespace)
	t.manageCronjobs = vars.ManageCronjobs

	// Ensure we have the needed namespace
	nsClient := t.client.Core().Namespaces()

	create := false
	_, err = nsClient.Get(t.namespace, metav1.GetOptions{})
	if err != nil {
		ignore := false
		if e, ok := err.(*errors.StatusError); ok {
			if e.ErrStatus.Reason == "NotFound" {
				ignore = true
				create = true
			}
		}
		if !ignore {
			return err
		}
	}
	if create {
		_, err = nsClient.Create(&v1.Namespace{
			ObjectMeta: metav1.ObjectMeta{
				Name: t.namespace,
			},
		})
		if err != nil {
			return err
		}
	}

	// Add the image pull secrets
	if len(vars.ImagePullSecrets) > 0 {
		saClient := t.client.Core().ServiceAccounts(t.namespace)

		var sa *v1.ServiceAccount
		// Account isn't always available right away, but it gets created in the end, just wait for it
		r := retrier.New(retrier.ConstantBackoff(10, 1*time.Second), nil)
		err := r.Run(func() error {
			s, err := saClient.Get("default", metav1.GetOptions{})
			if err != nil {
				return err
			}
			if s == nil {
				return fmt.Errorf("Service account not found (yet)")
			}
			sa = s
			return nil
		})

		secrets := make([]v1.LocalObjectReference, 0)
		for _, s := range vars.ImagePullSecrets {
			secrets = append(secrets, v1.LocalObjectReference{
				Name: s,
			})
		}

		sa.ImagePullSecrets = secrets
		_, err = saClient.Update(sa)
		if err != nil {
			return err
		}
	}

	return nil
}

func (t *KubernetesTarget) Apply(m Manifest, data []byte) error {
	// Temporary fix for https://github.com/kubernetes/kubernetes/issues/35149
	// If a cronjob is applied, an error occurs
	// Thus we delete the cronjob first if it exists
	if m.Kind == "CronJob" {
		out, err := t.runKubeCtl(nil, "get", "cronjob", "-o", "name")
		if err != nil {
			return err
		}
		lines := strings.Split(strings.TrimSpace(out), "\n")
		searchline := fmt.Sprintf("cronjob/%s", m.Metadata.Name)
		for _, line := range lines {
			if line == searchline {
				_, err := t.runKubeCtl(nil, "delete", "cronjob", m.Metadata.Name)
				if err != nil {
					return err
				}
				break
			}
		}
	}
	_, err := t.runKubeCtl(data, "apply", "-f", "-")
	return err
}

func (t *KubernetesTarget) Cleanup(items []Manifest) error {
	for _, ct := range CleanTypes {
		if ct == "cronjob" && !t.manageCronjobs {
			continue
		}
		err := t.cleanType(items, ct)
		if err != nil {
			return err
		}
	}
	return nil
}

func (t *KubernetesTarget) cleanType(items []Manifest, ct string) error {
	out, err := t.runKubeCtl(nil, "get", ct, "-o", "name")
	if err != nil {
		return err
	}

	known := []string{}
	for _, m := range items {
		if strings.ToLower(m.Kind) == ct {
			known = append(known, fmt.Sprintf("%s/%s", ct, m.Metadata.Name))
		}
	}

	lines := strings.Split(strings.TrimSpace(out), "\n")
	for _, line := range lines {
		if line == "" {
			continue
		}

		found := false
		for _, k := range known {
			if line == k {
				found = true
			}
		}

		if !found {
			_, err := t.runKubeCtl(nil, "delete", line)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (t *KubernetesTarget) runKubeCtl(stdin []byte, args ...string) (string, error) {
	return t.kubectl.Run(stdin, args...)
}
<|endoftext|>