<commit_before>package multio_test

import (
	"fmt"
	"github.com/creack/multio"
	"io"
	"log"
)

func ExampleMultiplexer() {
	r, w := io.Pipe()
	r2, w2 := io.Pipe()
	ch := make(chan struct{})
	go func() {
		defer close(ch)
		m, err := multio.NewMultiplexer(r, w2)
		if err != nil {
			log.Fatal(err)
		}
		wr := m.NewWriter(0)
		fmt.Fprintf(wr, "Hello World!!!")

		wr2 := m.NewWriter(1)
		fmt.Fprintf(wr2, "Hello the World!!!")
	}()
	m, _ := multio.NewMultiplexer(r2, w)

	var (
		buf = make([]byte, multio.PageSize)
		n   int
		err error
	)

	rd := m.NewReader(0)
	n, err = rd.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])

	rd2 := m.NewReader(1)
	n, err = rd2.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])
	<-ch

	// output:
	// Hello World!!!
	// Hello the World!!!
}
<commit_msg>Update example<commit_after>package multio_test

import (
	"fmt"
	"github.com/creack/multio"
	"io"
	"log"
	"os"
)

func ExampleMultiplexer() {
	// Create two pipe pair to have both way read/write
	// We could use socketpair, socket or other, but let's
	// stick with pipes for now
	r, w := io.Pipe()
	r2, w2 := io.Pipe()

	// Channel to wait on the goroutine before exit
	ch := make(chan struct{})
	go func() {
		// When the goroutine finished, unlock (close chan)
		defer close(ch)

		// Create a side of the multiplexer
		m, err := multio.NewMultiplexer(r, w2)
		if err != nil {
			log.Fatal(err)
		}

		// Create two writers
		wr := m.NewWriter(0)
		wr2 := m.NewWriter(1)

		// Send data and close.
		// In this example, we use io.Copy which blocks until close
		// If we do not close before write to wr2, the other goroutine
		// will never read so it will block forever.
		fmt.Fprintf(wr, "Hello World!!!\n")
		wr.Close()
		fmt.Fprintf(wr2, "Hello the World!!!\n")
		wr2.Close()
	}()

	// Create the other side of the multiplexer
	m, err := multio.NewMultiplexer(r2, w)
	if err != nil {
		log.Fatal(err)
	}

	// Create two readers corresponding to our writers
	rd := m.NewReader(0)
	rd2 := m.NewReader(1)

	// Copy everything (until close) to stdout
	io.Copy(os.Stdout, rd)
	io.Copy(os.Stdout, rd2)

	// Wait on the goroutine to finish
	<-ch

	// output:
	// Hello World!!!
	// Hello the World!!!
}
<|endoftext|>


<commit_before>package cachex_test

import (
	"fmt"
	"sync/atomic"
	"time"

	"github.com/alicebob/miniredis"
	"github.com/wencan/cachex"
	"github.com/wencan/cachex/lrucache"
	"github.com/wencan/cachex/rdscache"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
)

var (
	db *sqlx.DB

	rds *miniredis.Miniredis
)

func init() {
	db, _ = sqlx.Open("sqlite3", ":memory:")

	rds, _ = miniredis.Run()
}

func Example_WithLRUCache() {
	type DateTime struct {
		ID   int    `db:"id"`
		Date string `db:"date"`
		Time string `db:"time"`
	}

	var incrementer uint64
	query := func(key, value interface{}) error {
		dt := value.(*DateTime)

		id := atomic.AddUint64(&incrementer, 1)
		err := db.Get(dt, fmt.Sprintf("SELECT %d AS id, '2019-08-25' AS date, '10:54:35' AS time", id))
		if err != nil {
			fmt.Println(err)
			return err
		}
		return nil
	}

	s := lrucache.NewLRUCache(1000, time.Second)
	cache := cachex.NewCachex(s, cachex.QueryFunc(query))

	var dt DateTime
	err := cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	err = cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	// Output:
	// 1 2019-08-25 10:54:35
	// 1 2019-08-25 10:54:35
}

func Example_WithRdsCache() {
	type DateTime struct {
		ID   string `db:"id" msgpack:"i"`
		Date string `db:"date" msgpack:"d"`
		Time string `db:"time" msgpack:"t"`
	}

	var incrementer uint64
	query := func(key, value interface{}) error {
		dt := value.(*DateTime)

		id := atomic.AddUint64(&incrementer, 1)
		err := db.Get(dt, fmt.Sprintf("SELECT %d AS id, '2019-08-25' AS date, '10:54:35' AS time", id))
		if err != nil {
			fmt.Println(err)
			return err
		}
		return nil
	}

	s := rdscache.NewRdsCache("tcp", rds.Addr(), rdscache.PoolConfig{DB: 1}, rdscache.RdsKeyPrefixOption("cache"))
	cache := cachex.NewCachex(s, cachex.QueryFunc(query))

	var dt DateTime
	err := cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	err = cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	// Output:
	// 1 2019-08-25 10:54:35
	// 1 2019-08-25 10:54:35
}
<commit_msg>update example<commit_after>package cachex_test

import (
	"fmt"
	"sync/atomic"
	"time"

	"github.com/alicebob/miniredis"
	"github.com/wencan/cachex"
	"github.com/wencan/cachex/lrucache"
	"github.com/wencan/cachex/rdscache"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
)

var (
	db *sqlx.DB

	rds *miniredis.Miniredis
)

func init() {
	db, _ = sqlx.Open("sqlite3", ":memory:")

	rds, _ = miniredis.Run()
}

func ExampleCachex_LRUCache() {
	type DateTime struct {
		ID   int    `db:"id"`
		Date string `db:"date"`
		Time string `db:"time"`
	}

	var incrementer uint64
	query := func(key, value interface{}) error {
		dt := value.(*DateTime)

		id := atomic.AddUint64(&incrementer, 1)
		err := db.Get(dt, fmt.Sprintf("SELECT %d AS id, '2019-08-25' AS date, '10:54:35' AS time", id))
		if err != nil {
			fmt.Println(err)
			return err
		}
		return nil
	}

	s := lrucache.NewLRUCache(1000, time.Second)
	cache := cachex.NewCachex(s, cachex.QueryFunc(query))

	var dt DateTime
	err := cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	err = cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	// Output:
	// 1 2019-08-25 10:54:35
	// 1 2019-08-25 10:54:35
}

func ExampleCachex_RdsCache() {
	type DateTime struct {
		ID   string `db:"id" msgpack:"i"`
		Date string `db:"date" msgpack:"d"`
		Time string `db:"time" msgpack:"t"`
	}

	var incrementer uint64
	query := func(key, value interface{}) error {
		dt := value.(*DateTime)

		id := atomic.AddUint64(&incrementer, 1)
		err := db.Get(dt, fmt.Sprintf("SELECT %d AS id, '2019-08-25' AS date, '10:54:35' AS time", id))
		if err != nil {
			fmt.Println(err)
			return err
		}
		return nil
	}

	s := rdscache.NewRdsCache("tcp", rds.Addr(), rdscache.PoolConfig{DB: 1}, rdscache.RdsKeyPrefixOption("cache"))
	cache := cachex.NewCachex(s, cachex.QueryFunc(query))

	var dt DateTime
	err := cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	err = cache.Get(time.Now().Second(), &dt)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(dt.ID, dt.Date, dt.Time)

	// Output:
	// 1 2019-08-25 10:54:35
	// 1 2019-08-25 10:54:35
}
<|endoftext|>


<commit_before>package nlp

import "fmt"

func Example() {
	testCorpus := []string{
		"The quick brown fox jumped over the lazy dog",
		"hey diddle diddle, the cat and the fiddle",
		"the cow jumped over the moon",
		"the little dog laughed to see such fun",
		"and the dish ran away with the spoon",
	}

	query := "the brown fox ran around the dog"

	tokeniser := NewCountVectoriser()
	transformer := NewTfidfTransformer()

	// set k (the number of dimensions following truncation) to 4
	reducer := NewTruncatedSVD(4)

	// Transform the corpus into an LSI fitting the model to the documents in the process
	mat, _ := tokeniser.FitTransform(testCorpus...)
	mat, _ = transformer.FitTransform(mat)
	lsi, _ := reducer.FitTransform(mat)

	// run the query through the same pipeline that was fitted to the corpus and
	// to project it into the same dimensional space
	mat, _ = tokeniser.Transform(query)
	mat, _ = transformer.Transform(mat)
	queryVector, _ := reducer.Transform(mat)

	// iterate over document feature vectors (columns) in the LSI and compare with the
	// query vector for similarity.  Similarity is determined by the difference between
	// the angles of the vectors known as the cosine similarity
	highestSimilarity := -1.0
	var matched int
	_, docs := lsi.Dims()
	for i := 0; i < docs; i++ {
		similarity := CosineSimilarity(queryVector.ColView(0), lsi.ColView(i))
		fmt.Printf("Comparing '%s' = %f\n", testCorpus[i], similarity)
		if similarity > highestSimilarity {
			matched = i
			highestSimilarity = similarity
		}
	}

	fmt.Printf("Matched '%s' with a Cosine Similarity of %f\n", testCorpus[matched], highestSimilarity)
}
<commit_msg>added expected output<commit_after>package nlp

import "fmt"

func Example() {
	testCorpus := []string{
		"The quick brown fox jumped over the lazy dog",
		"hey diddle diddle, the cat and the fiddle",
		"the cow jumped over the moon",
		"the little dog laughed to see such fun",
		"and the dish ran away with the spoon",
	}

	query := "the brown fox ran around the dog"

	tokeniser := NewCountVectoriser()
	transformer := NewTfidfTransformer()

	// set k (the number of dimensions following truncation) to 4
	reducer := NewTruncatedSVD(4)

	// Transform the corpus into an LSI fitting the model to the documents in the process
	mat, _ := tokeniser.FitTransform(testCorpus...)
	mat, _ = transformer.FitTransform(mat)
	lsi, _ := reducer.FitTransform(mat)

	// run the query through the same pipeline that was fitted to the corpus and
	// to project it into the same dimensional space
	mat, _ = tokeniser.Transform(query)
	mat, _ = transformer.Transform(mat)
	queryVector, _ := reducer.Transform(mat)

	// iterate over document feature vectors (columns) in the LSI and compare with the
	// query vector for similarity.  Similarity is determined by the difference between
	// the angles of the vectors known as the cosine similarity
	highestSimilarity := -1.0
	var matched int
	_, docs := lsi.Dims()
	for i := 0; i < docs; i++ {
		similarity := CosineSimilarity(queryVector.ColView(0), lsi.ColView(i))
		if similarity > highestSimilarity {
			matched = i
			highestSimilarity = similarity
		}
	}

	fmt.Printf("Matched '%s'", testCorpus[matched])
	// Output: Matched 'The quick brown fox jumped over the lazy dog'
}
<|endoftext|>


<commit_before>package mpb_test

import (
	"io"
	"math/rand"
	"net/http"
	"time"

	"github.com/vbauerster/mpb"
	"github.com/vbauerster/mpb/decor"
)

func Example() {
	p := mpb.New(
		// override default (80) width
		mpb.WithWidth(100),
		// override default "[=>-]" format
		mpb.WithFormat("╢▌▌░╟"),
		// override default 120ms refresh rate
		mpb.WithRefreshRate(100*time.Millisecond),
	)

	total := 100
	name := "Single Bar:"
	// adding a single bar
	bar := p.AddBar(int64(total),
		mpb.PrependDecorators(
			// Display our static name with one space on the right
			decor.StaticName(name, len(name)+1, decor.DidentRight),
			// ETA decorator with width reservation of 3 runes
			decor.ETA(3, 0),
		),
		mpb.AppendDecorators(
			// Percentage decorator with width reservation of 5 runes
			decor.Percentage(5, 0),
		),
	)

	// simulating some work
	max := 100 * time.Millisecond
	for i := 0; i < total; i++ {
		time.Sleep(time.Duration(rand.Intn(10)+1) * max / 10)
		// increment by 1 (there is bar.IncrBy(int) method, if needed)
		bar.Increment()
	}
	// wait for our bar to complete and flush
	p.Wait()
}

func ExampleBar_Completed() {
	p := mpb.New()
	bar := p.AddBar(100)

	max := 100 * time.Millisecond
	for !bar.Completed() {
		time.Sleep(time.Duration(rand.Intn(10)+1) * max / 10)
		bar.Increment()
	}

	p.Wait()
}

func ExampleBar_ProxyReader() {
	p := mpb.New()
	// make http get request
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	// Assuming ContentLength > 0
	bar := p.AddBar(resp.ContentLength,
		decor.CountersKibiByte("%6.1f / %6.1f", 12, 0),
	)

	// create proxy reader
	reader := bar.ProxyReader(resp.Body)

	// and copy from reader, ignoring errors
	io.Copy(dest, reader)

	p.Wait()
}
<commit_msg>AppendDecorators<commit_after>package mpb_test

import (
	"io"
	"math/rand"
	"net/http"
	"time"

	"github.com/vbauerster/mpb"
	"github.com/vbauerster/mpb/decor"
)

func Example() {
	p := mpb.New(
		// override default (80) width
		mpb.WithWidth(100),
		// override default "[=>-]" format
		mpb.WithFormat("╢▌▌░╟"),
		// override default 120ms refresh rate
		mpb.WithRefreshRate(100*time.Millisecond),
	)

	total := 100
	name := "Single Bar:"
	// adding a single bar
	bar := p.AddBar(int64(total),
		mpb.PrependDecorators(
			// Display our static name with one space on the right
			decor.StaticName(name, len(name)+1, decor.DidentRight),
			// ETA decorator with width reservation of 3 runes
			decor.ETA(3, 0),
		),
		mpb.AppendDecorators(
			// Percentage decorator with width reservation of 5 runes
			decor.Percentage(5, 0),
		),
	)

	// simulating some work
	max := 100 * time.Millisecond
	for i := 0; i < total; i++ {
		time.Sleep(time.Duration(rand.Intn(10)+1) * max / 10)
		// increment by 1 (there is bar.IncrBy(int) method, if needed)
		bar.Increment()
	}
	// wait for our bar to complete and flush
	p.Wait()
}

func ExampleBar_Completed() {
	p := mpb.New()
	bar := p.AddBar(100)

	max := 100 * time.Millisecond
	for !bar.Completed() {
		time.Sleep(time.Duration(rand.Intn(10)+1) * max / 10)
		bar.Increment()
	}

	p.Wait()
}

func ExampleBar_ProxyReader() {
	p := mpb.New()
	// make http get request
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	// Assuming ContentLength > 0
	bar := p.AddBar(resp.ContentLength,
		mpb.AppendDecorators(
			decor.CountersKibiByte("%6.1f / %6.1f", 12, 0),
		),
	)

	// create proxy reader
	reader := bar.ProxyReader(resp.Body)

	// and copy from reader, ignoring errors
	io.Copy(dest, reader)

	p.Wait()
}
<|endoftext|>


<commit_before>package ipState

import (
	"bytes"
	"encoding/json"
	"net/http"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/nitro/memberlist"
)

type peersIpStatus struct {
	IpAddr  string
	IpsView map[string]bool
}

type GlobalMap struct {
	IpLive       map[string]bool           // true if an IP is still alive and false otherwise
	peersIpView  map[string]*peersIpStatus //map of monitored IP addresses to each peers' views of them
	MinAgreement int
	Mutex        *sync.RWMutex
	Gossip       *memberlist.Memberlist
	Peers        []string
	LocalAddr    string
	Clustered    bool
	Members      *memberlist.Memberlist
	DryRun       bool
}

type HttpUpdate struct {
	Peer      string
	IpAddress string
	Live      bool
}

var Gm *GlobalMap

//Initialize the Global Map struct
func InitGM(ipAddresses []string, dryRun bool) {
	m := make(map[string]bool)
	for _, ip := range ipAddresses {
		m[ip] = true
	}

	Gm = &GlobalMap{IpLive: m, Mutex: &sync.RWMutex{}, DryRun: dryRun}
}

func newPeersViewsMap() *peersIpStatus {
	return &peersIpStatus{IpsView: make(map[string]bool)}
}

//Initialize the state of Peers' IP views to be true
func InitPeersIpViews() {
	Gm.peersIpView = make(map[string]*peersIpStatus)
	for ipAddr, _ := range Gm.IpLive {
		Gm.peersIpView[ipAddr] = &peersIpStatus{IpAddr: ipAddr, IpsView: make(map[string]bool)}
		for _, peer := range Gm.Peers {
			Gm.peersIpView[ipAddr].IpsView[peer] = true
		}
	}
}

//notifyPeers will update all peers on an IP address state change
func notifyPeers(ipAddress string, live bool) error {
	for _, peer := range Gm.Peers {
		log.Debugln("Send update to peer:", peer)
		upd := &HttpUpdate{Peer: Gm.LocalAddr, IpAddress: ipAddress, Live: live}
		data, err := json.Marshal(upd)
		if err != nil {
			log.Errorln("Error Marshaling", err)
			return err
		}
		buff := bytes.NewBuffer(data)
		req, err := http.NewRequest("POST", "http://"+peer, buff)
		req.Close = true
		req.Header.Set("Content-Type", "application/json")
		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			log.Errorln("HTTP Post Error", err)
			return err
		}
		defer resp.Body.Close()
	}
	log.Debugln("Sent update to Peers", Gm.Peers)
	return nil
}

//Update ipstate in Peers' views based on http update
func UpdateGlobalState(ipAddress string, live bool, peer string) (liveCheck int) {
	Gm.Mutex.Lock()
	Gm.peersIpView[ipAddress].IpsView[peer] = live
	liveCheck = checkGlobalState(ipAddress)
	Gm.Mutex.Unlock()
	return
}

func checkGlobalState(ipAddress string) (liveCheck int) {
	liveCheck = 0
	check := Gm.peersIpView[ipAddress]
	for _, v := range check.IpsView {
		if v {
			liveCheck += 1
		} else {
			liveCheck -= 1
		}
	}
	return
}

func NotifyIpState(ipAddress string, live bool, peerUpdate bool) error {
	Gm.Mutex.RLock()
	if Gm.IpLive[ipAddress] == live && peerUpdate == true { //if we have seen the failure and received a notification from a peer that they've seen the failure
		Gm.Mutex.RUnlock()
		log.Debugln("Received agreement from Peer.  Updating IP Pool")
		if live == false {
			err := Master.Dns.RemoveIp(ipAddress, Gm.DryRun)
			if err != nil {
				log.Errorln("Error Removing IP: ", ipAddress, err)
			}
		} else {
			err := Master.Dns.AddIp(ipAddress, Gm.DryRun)
			if err != nil {
				log.Errorln("Error Adding IP: ", ipAddress, err)
			}
		}
	} else if peerUpdate == false { //we haven't received a message from a peer yet so just update the state and notify peers
		Gm.Mutex.RUnlock()
		Gm.Mutex.Lock()
		Gm.IpLive[ipAddress] = live
		Gm.Mutex.Unlock()
		log.Debugln("Notifying peers: ", ipAddress, live)
		err := notifyPeers(ipAddress, live)
		if err != nil {
			log.Errorln("Error notifying Peers", err)
		}
	} else {
		Gm.Mutex.RUnlock()
	}

	return nil
}
<commit_msg>add error logging for http req<commit_after>package ipState

import (
	"bytes"
	"encoding/json"
	"net/http"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/nitro/memberlist"
)

type peersIpStatus struct {
	IpAddr  string
	IpsView map[string]bool
}

type GlobalMap struct {
	IpLive       map[string]bool           // true if an IP is still alive and false otherwise
	peersIpView  map[string]*peersIpStatus //map of monitored IP addresses to each peers' views of them
	MinAgreement int
	Mutex        *sync.RWMutex
	Gossip       *memberlist.Memberlist
	Peers        []string
	LocalAddr    string
	Clustered    bool
	Members      *memberlist.Memberlist
	DryRun       bool
}

type HttpUpdate struct {
	Peer      string
	IpAddress string
	Live      bool
}

var Gm *GlobalMap

//Initialize the Global Map struct
func InitGM(ipAddresses []string, dryRun bool) {
	m := make(map[string]bool)
	for _, ip := range ipAddresses {
		m[ip] = true
	}

	Gm = &GlobalMap{IpLive: m, Mutex: &sync.RWMutex{}, DryRun: dryRun}
}

func newPeersViewsMap() *peersIpStatus {
	return &peersIpStatus{IpsView: make(map[string]bool)}
}

//Initialize the state of Peers' IP views to be true
func InitPeersIpViews() {
	Gm.peersIpView = make(map[string]*peersIpStatus)
	for ipAddr, _ := range Gm.IpLive {
		Gm.peersIpView[ipAddr] = &peersIpStatus{IpAddr: ipAddr, IpsView: make(map[string]bool)}
		for _, peer := range Gm.Peers {
			Gm.peersIpView[ipAddr].IpsView[peer] = true
		}
	}
}

//notifyPeers will update all peers on an IP address state change
func notifyPeers(ipAddress string, live bool) error {
	for _, peer := range Gm.Peers {
		log.Debugln("Send update to peer:", peer)
		upd := &HttpUpdate{Peer: Gm.LocalAddr, IpAddress: ipAddress, Live: live}
		data, err := json.Marshal(upd)
		if err != nil {
			log.Errorln("Error Marshaling", err)
			return err
		}
		buff := bytes.NewBuffer(data)
		req, err := http.NewRequest("POST", "http://"+peer, buff)
		if err != nil {
			log.Errorln("Error forming new request", err)
			return err
		}
		req.Close = true
		req.Header.Set("Content-Type", "application/json")
		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			log.Errorln("HTTP Post Error", err)
			return err
		}
		defer resp.Body.Close()
	}
	log.Debugln("Sent update to Peers", Gm.Peers)
	return nil
}

//Update ipstate in Peers' views based on http update
func UpdateGlobalState(ipAddress string, live bool, peer string) (liveCheck int) {
	Gm.Mutex.Lock()
	Gm.peersIpView[ipAddress].IpsView[peer] = live
	liveCheck = checkGlobalState(ipAddress)
	Gm.Mutex.Unlock()
	return
}

func checkGlobalState(ipAddress string) (liveCheck int) {
	liveCheck = 0
	check := Gm.peersIpView[ipAddress]
	for _, v := range check.IpsView {
		if v {
			liveCheck += 1
		} else {
			liveCheck -= 1
		}
	}
	return
}

func NotifyIpState(ipAddress string, live bool, peerUpdate bool) error {
	Gm.Mutex.RLock()
	if Gm.IpLive[ipAddress] == live && peerUpdate == true { //if we have seen the failure and received a notification from a peer that they've seen the failure
		Gm.Mutex.RUnlock()
		log.Debugln("Received agreement from Peer.  Updating IP Pool")
		if live == false {
			err := Master.Dns.RemoveIp(ipAddress, Gm.DryRun)
			if err != nil {
				log.Errorln("Error Removing IP: ", ipAddress, err)
			}
		} else {
			err := Master.Dns.AddIp(ipAddress, Gm.DryRun)
			if err != nil {
				log.Errorln("Error Adding IP: ", ipAddress, err)
			}
		}
	} else if peerUpdate == false { //we haven't received a message from a peer yet so just update the state and notify peers
		Gm.Mutex.RUnlock()
		Gm.Mutex.Lock()
		Gm.IpLive[ipAddress] = live
		Gm.Mutex.Unlock()
		log.Debugln("Notifying peers: ", ipAddress, live)
		err := notifyPeers(ipAddress, live)
		if err != nil {
			log.Errorln("Error notifying Peers", err)
		}
	} else {
		Gm.Mutex.RUnlock()
	}

	return nil
}
<|endoftext|>


<commit_before>// select random value from source

// TODO: load up the source

func LoadSource(sourcePath string) (*SourceFile, error) {
}

<commit_msg>structs<commit_after>// select random value from source

type JargonSource struct {
    Source  string `json:"source,omitempty"`
    Jargon  string `json:"jargons,omitempty"`
}

type JargonFile struct {
    Jargons  map[string]JargonSource `json:"jargons,omitempty"`
    filePath string
}

// TODO: load up the source
func LoadSource(filePath string) (*JargonFile, error) {
}

<|endoftext|>


<commit_before>package client

import (
    "path/filepath"
    "fmt"
    "io"
    "log"
    "net"
    "os"
    "testing"
    "time"
)

func testXMLClient(xmlGlob string) *XMLClient {
    streamFiles, err := filepath.Glob(xmlGlob)
    if err != nil {
        panic(fmt.Errorf("filepath.Glob: %v\n", err))
    }

    // setup mock server
    tcpListener, err := net.ListenTCP("tcp", nil)
    if err != nil {
        panic(err)
    }

    go func(tcpListener *net.TCPListener) {
        for {
            tcpConn, err := tcpListener.AcceptTCP()
            if err != nil {
                panic(err)
            }

            // write each .xml file into the testConn end of the pipe, for XMLClient to read
            go func(tcpConn *net.TCPConn) {
                defer tcpConn.CloseWrite()

                for _, xmlPath := range streamFiles {
                    xmlFile, err := os.Open(xmlPath)
                    if err != nil {
                        panic(fmt.Errorf("os.Open: %v\n", err))
                    }
                    defer xmlFile.Close()

                    if write, err := io.Copy(tcpConn, xmlFile); err != nil {
                        panic(fmt.Errorf("io.Copy: %v\n", err))
                    } else {
                        log.Printf("Send %v: %d bytes\n", xmlPath, write)
                    }

                    time.Sleep(1 * time.Second)
                }

                log.Printf("Send: done\n")
            }(tcpConn)

            go func(tcpConn *net.TCPConn) {
                defer tcpConn.Close()

                for {
                    buf := make([]byte, 1500)

                    if read, err := tcpConn.Read(buf); err != nil {
                        log.Printf("Recv error: %v\n", err)
                        break
                    } else {
                        log.Printf("Recv %d bytes: %#v\n", read, string(buf[:read]))
                    }
                }

                log.Printf("Recv: done\n")
            }(tcpConn)
        }
    }(tcpListener)

    // connect client to test server
    listenAddr := tcpListener.Addr().(*net.TCPAddr)

    tcpConn, err := net.DialTCP("tcp", nil, listenAddr)
    if err != nil {
        panic(err)
    }

    xmlClient := XMLClient{
        timeout:    10 * time.Second,
        conn:       tcpConn,
    }

    if err := xmlClient.start(); err != nil {
        panic(err)
    }

    return &xmlClient
}

func TestXmlRead(t *testing.T) {
    xmlClient := testXMLClient("./test-xml/test1-*.xml")

    listenChan, err := xmlClient.Listen()
    if err != nil {
        t.Fatalf("xmlClient.Listen: %v", err)
    }

    // read System state updates, and updates this to be the final System state after all cumulative updates
    var listenSystem System

    for listenSystem = range listenChan {
        // read it
        _ = listenSystem.String()
    }

    if err := xmlClient.ListenError(); err != nil {
        t.Fatalf("xmlClient.ListenError: %v\n", err)
    } else {
        t.Logf("End of Listen\n")
    }

    // check resulting system state
    if source0, exists := listenSystem.SrcMgr.SourceCol[0]; !exists {
        t.Errorf("Source #0 does not exist")
    } else {
        t.Logf("Source #0: %#v\n", source0)

        if source0.Name != "PC 1-1" {
            t.Errorf("Source #0 Name: %v", source0.Name)
        }
    }

    if screen0, exists := listenSystem.DestMgr.ScreenDestCol[0]; !exists {
        t.Errorf("Screen #0 does not exist")
    } else {
        t.Logf("Screen #0: %#v\n", screen0)

        if screen0.Name != "ScreenDest1" {
            t.Errorf("Screen #0 Name: %v", screen0.Name)
        }

        if screen0.IsActive != 1 {
            t.Errorf("Screen #0: IsActive=%v", screen0.IsActive)
        }

        if layer0, exists := screen0.LayerCollection[0]; !exists {
            t.Errorf("Layer #0 does not exist")
        } else {
            t.Logf("Layer #0: %#v\n", layer0)

            if layer0.LastSrcIdx != 0 {
                t.Errorf("Layer #0: LastSrcIdx=%v", layer0.LastSrcIdx)
            }

            if layer0.PgmMode != 1 || layer0.PvwMode != 0 {
                t.Errorf("Layer #0: PgmMode=%v PvwMode=%v", layer0.PgmMode, layer0.PvwMode)
            }
        }

        if layer1, exists := screen0.LayerCollection[1]; !exists {
            t.Errorf("Layer #1 does not exist")
        } else {
            t.Logf("Layer #1: %#v\n", layer1)

            if layer1.LastSrcIdx != 0 {
                t.Errorf("Layer #1: LastSrcIdx=%v", layer1.LastSrcIdx)
            }

            if layer1.PgmMode != 0 || layer1.PvwMode != 1 {
                t.Errorf("Layer #0: PgmMode=%v PvwMode=%v", layer1.PgmMode, layer1.PvwMode)
            }
        }
    }

    if preset0, exists := listenSystem.PresetMgr.Preset[0]; !exists {
        t.Errorf("Preset #0 does not exist")
    } else {
        t.Logf("Preset #0: %#v\n", preset0)

        if preset0.Name != "Test 1" {
            t.Errorf("Preset #0 Name: %v", preset0.Name)
        }
    }
}

// Adding new outputs/inputs from initial empty state
func TestXmlAdd(t *testing.T) {
    xmlClient := testXMLClient("./test-xml/test2-*.xml")

    listenChan, err := xmlClient.Listen()
    if err != nil {
        t.Fatalf("xmlClient.Listen: %v", err)
    }

    // read System state updates, and updates this to be the final System state after all cumulative updates
    var listenSystem System

    for listenSystem = range listenChan {
        // read it
        _ = listenSystem.String()
    }

    log.Printf("End of Listen\n")

    // check resulting system state
    if source0, exists := listenSystem.SrcMgr.SourceCol[0]; !exists {
        t.Errorf("Source #0 does not exist")
    } else {
        t.Logf("Source #0: %#v\n", source0)

        if source0.Name != "ScreenDest1_PGM-1" {
            t.Errorf("Source #0 Name: %v", source0.Name)
        }
        if !(source0.InputCfgIndex == -1 && source0.StillIndex == -1 && source0.DestIndex == 0) {
            t.Errorf("Source 0: InputCfgIndex=%v StillIndex=%v DestIndex=%v", source0.InputCfgIndex, source0.StillIndex, source0.DestIndex)
        }
    }
    if source1, exists := listenSystem.SrcMgr.SourceCol[1]; !exists {
        t.Errorf("Source #1 does not exist")
    } else {
        t.Logf("Source #1: %#v\n", source1)

        if source1.Name != "Input1-2" {
            t.Errorf("Source #1 Name: %v", source1.Name)
        }
        if !(source1.InputCfgIndex == 0 && source1.StillIndex == -1 && source1.DestIndex == -1) {
            t.Errorf("Source 1: InputCfgIndex=%v StillIndex=%v DestIndex=%v", source1.InputCfgIndex, source1.StillIndex, source1.DestIndex)
        }
    }
    if source2, exists := listenSystem.SrcMgr.SourceCol[2]; !exists {
        t.Logf("Source #2: removed\n")
    } else {
        t.Logf("Source #2: still exists: %#v\n", source2)
    }

    if screen0, exists := listenSystem.DestMgr.ScreenDestCol[0]; !exists {
        t.Errorf("Screen #0 does not exist")
    } else {
        t.Logf("Screen #0: %#v\n", screen0)

        if screen0.Name != "ScreenDest1" {
            t.Errorf("Screen #0 Name: %v", screen0.Name)
        }

        if screen0.IsActive != 0 {
            t.Errorf("Screen #0: IsActive=%v", screen0.IsActive)
        }
    }
}
<commit_msg>client: refactor testXMLClientRead()<commit_after>package client

import (
    "path/filepath"
    "fmt"
    "io"
    "log"
    "net"
    "os"
    "testing"
    "time"
)

type testXML struct {
    fileGlob    string
    fileDelay   time.Duration
}

// Create a new XMLClient, connected to a local TCP server that serves up the globbed .xml files in order, using a 1
func testXMLClient(test testXML) *XMLClient {
    streamFiles, err := filepath.Glob(test.fileGlob)
    if err != nil {
        panic(fmt.Errorf("filepath.Glob: %v\n", err))
    }

    // setup mock server
    tcpListener, err := net.ListenTCP("tcp", nil)
    if err != nil {
        panic(err)
    }

    go func(tcpListener *net.TCPListener) {
        for {
            tcpConn, err := tcpListener.AcceptTCP()
            if err != nil {
                panic(err)
            }

            // write each .xml file into the testConn end of the pipe, for XMLClient to read
            go func(tcpConn *net.TCPConn) {
                defer tcpConn.CloseWrite()

                for _, xmlPath := range streamFiles {
                    xmlFile, err := os.Open(xmlPath)
                    if err != nil {
                        panic(fmt.Errorf("os.Open: %v\n", err))
                    }
                    defer xmlFile.Close()

                    if write, err := io.Copy(tcpConn, xmlFile); err != nil {
                        panic(fmt.Errorf("io.Copy: %v\n", err))
                    } else {
                        log.Printf("Send %v: %d bytes\n", xmlPath, write)
                    }

                    if test.fileDelay > 0 {
                        time.Sleep(test.fileDelay)
                    }
                }

                log.Printf("Send: done\n")
            }(tcpConn)

            go func(tcpConn *net.TCPConn) {
                defer tcpConn.Close()

                for {
                    buf := make([]byte, 1500)

                    if read, err := tcpConn.Read(buf); err != nil {
                        log.Printf("Recv error: %v\n", err)
                        break
                    } else {
                        log.Printf("Recv %d bytes: %#v\n", read, string(buf[:read]))
                    }
                }

                log.Printf("Recv: done\n")
            }(tcpConn)
        }
    }(tcpListener)

    // connect client to test server
    listenAddr := tcpListener.Addr().(*net.TCPAddr)

    tcpConn, err := net.DialTCP("tcp", nil, listenAddr)
    if err != nil {
        panic(err)
    }

    xmlClient := XMLClient{
        timeout:    10 * time.Second,
        conn:       tcpConn,
    }

    if err := xmlClient.start(); err != nil {
        panic(err)
    }

    return &xmlClient
}

// Create an XMLClient, and return the final System state after reading all the .xml files
func testXMLClientRead(test testXML) System {
    xmlClient := testXMLClient(test)

    var system System

    for {
        if readSystem, err := xmlClient.Read(); err == io.EOF {
            return system
        } else if err != nil {
            panic(err)
        } else {
            system = readSystem
        }
    }
}

func TestXmlRead(t *testing.T) {
    // test with some minor concurrency
    xmlClient := testXMLClient(testXML{
        fileGlob:   "./test-xml/test1-*.xml",
        fileDelay:  10 * time.Millisecond,
    })

    listenChan, err := xmlClient.Listen()
    if err != nil {
        t.Fatalf("xmlClient.Listen: %v", err)
    }

    // read System state updates, and updates this to be the final System state after all cumulative updates
    var listenSystem System

    for listenSystem = range listenChan {
        // read it
        _ = listenSystem.String()
    }

    if err := xmlClient.ListenError(); err != nil {
        t.Fatalf("xmlClient.ListenError: %v\n", err)
    } else {
        t.Logf("End of Listen\n")
    }

    // check resulting system state
    if source0, exists := listenSystem.SrcMgr.SourceCol[0]; !exists {
        t.Errorf("Source #0 does not exist")
    } else {
        t.Logf("Source #0: %#v\n", source0)

        if source0.Name != "PC 1-1" {
            t.Errorf("Source #0 Name: %v", source0.Name)
        }
    }

    if screen0, exists := listenSystem.DestMgr.ScreenDestCol[0]; !exists {
        t.Errorf("Screen #0 does not exist")
    } else {
        t.Logf("Screen #0: %#v\n", screen0)

        if screen0.Name != "ScreenDest1" {
            t.Errorf("Screen #0 Name: %v", screen0.Name)
        }

        if screen0.IsActive != 1 {
            t.Errorf("Screen #0: IsActive=%v", screen0.IsActive)
        }

        if layer0, exists := screen0.LayerCollection[0]; !exists {
            t.Errorf("Layer #0 does not exist")
        } else {
            t.Logf("Layer #0: %#v\n", layer0)

            if layer0.LastSrcIdx != 0 {
                t.Errorf("Layer #0: LastSrcIdx=%v", layer0.LastSrcIdx)
            }

            if layer0.PgmMode != 1 || layer0.PvwMode != 0 {
                t.Errorf("Layer #0: PgmMode=%v PvwMode=%v", layer0.PgmMode, layer0.PvwMode)
            }
        }

        if layer1, exists := screen0.LayerCollection[1]; !exists {
            t.Errorf("Layer #1 does not exist")
        } else {
            t.Logf("Layer #1: %#v\n", layer1)

            if layer1.LastSrcIdx != 0 {
                t.Errorf("Layer #1: LastSrcIdx=%v", layer1.LastSrcIdx)
            }

            if layer1.PgmMode != 0 || layer1.PvwMode != 1 {
                t.Errorf("Layer #0: PgmMode=%v PvwMode=%v", layer1.PgmMode, layer1.PvwMode)
            }
        }
    }

    if preset0, exists := listenSystem.PresetMgr.Preset[0]; !exists {
        t.Errorf("Preset #0 does not exist")
    } else {
        t.Logf("Preset #0: %#v\n", preset0)

        if preset0.Name != "Test 1" {
            t.Errorf("Preset #0 Name: %v", preset0.Name)
        }
    }
}

// Adding new outputs/inputs from initial empty state
func TestXmlAdd(t *testing.T) {
    system := testXMLClientRead(testXML{
        fileGlob:   "./test-xml/test2-*.xml",
    })

    // check resulting system state
    if source0, exists := system.SrcMgr.SourceCol[0]; !exists {
        t.Errorf("Source #0 does not exist")
    } else {
        if source0.Name != "ScreenDest1_PGM-1" {
            t.Errorf("Source #0 Name: %v", source0.Name)
        }
        if !(source0.InputCfgIndex == -1 && source0.StillIndex == -1 && source0.DestIndex == 0) {
            t.Errorf("Source 0: InputCfgIndex=%v StillIndex=%v DestIndex=%v", source0.InputCfgIndex, source0.StillIndex, source0.DestIndex)
        }
    }
    if source1, exists := system.SrcMgr.SourceCol[1]; !exists {
        t.Errorf("Source #1 does not exist")
    } else {
        if source1.Name != "Input1-2" {
            t.Errorf("Source #1 Name: %v", source1.Name)
        }
        if !(source1.InputCfgIndex == 0 && source1.StillIndex == -1 && source1.DestIndex == -1) {
            t.Errorf("Source 1: InputCfgIndex=%v StillIndex=%v DestIndex=%v", source1.InputCfgIndex, source1.StillIndex, source1.DestIndex)
        }
    }
    if source2, exists := system.SrcMgr.SourceCol[2]; !exists {

    } else {
        t.Errorf("Source #2: still exists: %#v\n", source2)
    }

    if screen0, exists := system.DestMgr.ScreenDestCol[0]; !exists {
        t.Errorf("Screen #0 does not exist")
    } else {
        if screen0.Name != "ScreenDest1" {
            t.Errorf("Screen #0 Name: %v", screen0.Name)
        }

        if screen0.IsActive != 0 {
            t.Errorf("Screen #0: IsActive=%v", screen0.IsActive)
        }
    }
}
<|endoftext|>


<commit_before>package cloudbrain

import (
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/go-multierror"
	"github.com/travis-ci/cloud-brain/background"
	"github.com/travis-ci/cloud-brain/cbcontext"
	"github.com/travis-ci/cloud-brain/cloud"
	"github.com/travis-ci/cloud-brain/database"
	"golang.org/x/crypto/scrypt"
	"golang.org/x/net/context"
)

const MaxCreateRetries = 10

type Core struct {
	db database.DB
	bb background.Backend

	cloudProvidersMutex sync.Mutex
	cloudProviders      map[string]cloud.Provider
}

func NewCore(db database.DB, bb background.Backend) *Core {
	return &Core{
		db: db,
		bb: bb,
	}
}

// GetInstance gets the instance information stored in the database for a given
// instance ID.
func (c *Core) GetInstance(ctx context.Context, id string) (*Instance, error) {
	instance, err := c.db.GetInstance(id)

	if err == database.ErrInstanceNotFound {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	return &Instance{
		ID:           instance.ID,
		ProviderName: instance.ProviderName,
		Image:        instance.Image,
		State:        instance.State,
		IPAddress:    instance.IPAddress,
	}, nil
}

// CreateInstanceAttributes contains attributes needed to start an instance
type CreateInstanceAttributes struct {
	ImageName    string
	InstanceType string
	PublicSSHKey string
}

// CreateInstance creates an instance in the database and queues off the cloud
// create job in the background.
func (c *Core) CreateInstance(ctx context.Context, providerName string, attr CreateInstanceAttributes) (*Instance, error) {
	id, err := c.db.CreateInstance(database.Instance{
		ProviderName: providerName,
		Image:        attr.ImageName,
		InstanceType: attr.InstanceType,
		PublicSSHKey: attr.PublicSSHKey,
		State:        "creating",
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":        err,
			"provider":   providerName,
			"image_name": attr.ImageName,
		}).Error("error creating instance in database")
		return nil, err
	}

	err = c.bb.Enqueue(background.Job{
		Context:    ctx,
		Payload:    []byte(id),
		Queue:      "create",
		MaxRetries: MaxCreateRetries,
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error enqueueing 'create' job in the background")
		// TODO(henrikhodne): Delete the record in the database?
		return nil, err
	}

	return &Instance{
		ID:           id,
		ProviderName: providerName,
		Image:        attr.ImageName,
		State:        "creating",
	}, nil
}

func (c *Core) ProviderCreateInstance(ctx context.Context, byteID []byte) error {
	id := string(byteID)

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
	}).Info("creating instance")

	dbInstance, err := c.db.GetInstance(id)
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error fetching instance from DB")
		return err
	}

	cloudProvider, err := c.cloudProvider(dbInstance.ProviderName)
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"instance_id":   id,
			"provider_name": dbInstance.ProviderName,
			"err":           err,
		}).Error("couldn't find provider with given name")
		return err
	}

	instance, err := cloudProvider.Create(id, cloud.CreateAttributes{
		ImageName:    dbInstance.Image,
		InstanceType: cloud.InstanceType(dbInstance.InstanceType),
		PublicSSHKey: dbInstance.PublicSSHKey,
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error creating instance")
		return err
	}

	dbInstance.State = "starting"

	err = c.db.UpdateInstance(dbInstance)
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
			"provider_id": instance.ID,
		}).Error("couldn't update instance in DB")
		return err
	}

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
		"provider_id": instance.ID,
	}).Info("created instance")

	return nil
}

func (c *Core) ProviderRefresh(ctx context.Context) error {
	c.refreshProviders()
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()

	var result error

	for providerName, cloudProvider := range c.cloudProviders {
		instances, err := cloudProvider.List()
		if err != nil {
			result = multierror.Append(result, err)
			continue
		}

		for _, instance := range instances {
			dbInstance, err := c.db.GetInstance(instance.ID)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider_name": providerName,
					"provider_id":   instance.ID,
				}).Error("failed fetching instance from database")
				continue
			}

			dbInstance.IPAddress = instance.IPAddress
			dbInstance.State = string(instance.State)

			err = c.db.UpdateInstance(dbInstance)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider":    providerName,
					"provider_id": instance.ID,
					"db_id":       dbInstance.ID,
				}).Error("failed to update instance in database")
			}
		}

		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"provider":       providerName,
			"instance_count": len(instances),
		}).Info("refreshed instances")
	}

	return result
}

func (c *Core) CheckToken(tokenID uint64, token string) (bool, error) {
	salt, hash, err := c.db.GetSaltAndHashForTokenID(tokenID)
	if err != nil {
		return false, err
	}

	decodedToken, err := hex.DecodeString(token)
	if err != nil {
		return false, err
	}

	generatedHash, err := scrypt.Key(decodedToken, salt, 16384, 8, 1, 32)
	if err != nil {
		return false, err
	}

	return subtle.ConstantTimeCompare(generatedHash, hash) == 1, nil
}

func (c *Core) cloudProvider(name string) (cloud.Provider, error) {
	c.cloudProvidersMutex.Lock()
	cloudProvider, ok := c.cloudProviders[name]
	c.cloudProvidersMutex.Unlock()
	if !ok {
		err := c.refreshProviders()
		if err != nil {
			return nil, err
		}

		c.cloudProvidersMutex.Lock()
		cloudProvider, ok = c.cloudProviders[name]
		c.cloudProvidersMutex.Unlock()
		if !ok {
			// This really shouldn't happen, since the database should ensure
			// that a provider with a matching name exists.
			return nil, fmt.Errorf("couldn't find a provider with that name")
		}
	}

	return cloudProvider, nil
}

func (c *Core) refreshProviders() error {
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()
	dbCloudProviders, err := c.db.ListProviders()
	if err != nil {
		return err
	}

	cloudProviders := make(map[string]cloud.Provider)

	for _, dbCloudProvider := range dbCloudProviders {
		cloudProvider, err := cloud.NewProvider(dbCloudProvider.Type, dbCloudProvider.Config)
		if err != nil {
			return err
		}

		cloudProviders[dbCloudProvider.Name] = cloudProvider
	}

	c.cloudProviders = cloudProviders

	return nil
}

type Instance struct {
	ID           string
	ProviderName string
	Image        string
	State        string
	IPAddress    string
}
<commit_msg>core: document all exported and some unexported names<commit_after>package cloudbrain

import (
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/go-multierror"
	"github.com/travis-ci/cloud-brain/background"
	"github.com/travis-ci/cloud-brain/cbcontext"
	"github.com/travis-ci/cloud-brain/cloud"
	"github.com/travis-ci/cloud-brain/database"
	"golang.org/x/crypto/scrypt"
	"golang.org/x/net/context"
)

// MaxCreateRetries is the number of times the "create" job will be retried.
const MaxCreateRetries = 10

// Core is used as a central manager for all Cloud Brain functionality. The HTTP
// API and the background workers are just frontends for the Core, and calls
// methods on Core for functionality.
type Core struct {
	db database.DB
	bb background.Backend

	cloudProvidersMutex sync.Mutex
	cloudProviders      map[string]cloud.Provider
}

// NewCore is used to create a new Core backed by the given database and
// background Backend.
func NewCore(db database.DB, bb background.Backend) *Core {
	return &Core{
		db: db,
		bb: bb,
	}
}

// GetInstance gets the instance information stored in the database for a given
// instance ID.
func (c *Core) GetInstance(ctx context.Context, id string) (*Instance, error) {
	instance, err := c.db.GetInstance(id)

	if err == database.ErrInstanceNotFound {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	return &Instance{
		ID:           instance.ID,
		ProviderName: instance.ProviderName,
		Image:        instance.Image,
		State:        instance.State,
		IPAddress:    instance.IPAddress,
	}, nil
}

// CreateInstanceAttributes contains attributes needed to start an instance
type CreateInstanceAttributes struct {
	ImageName    string
	InstanceType string
	PublicSSHKey string
}

// CreateInstance creates an instance in the database and queues off the cloud
// create job in the background.
func (c *Core) CreateInstance(ctx context.Context, providerName string, attr CreateInstanceAttributes) (*Instance, error) {
	id, err := c.db.CreateInstance(database.Instance{
		ProviderName: providerName,
		Image:        attr.ImageName,
		InstanceType: attr.InstanceType,
		PublicSSHKey: attr.PublicSSHKey,
		State:        "creating",
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":        err,
			"provider":   providerName,
			"image_name": attr.ImageName,
		}).Error("error creating instance in database")
		return nil, err
	}

	err = c.bb.Enqueue(background.Job{
		Context:    ctx,
		Payload:    []byte(id),
		Queue:      "create",
		MaxRetries: MaxCreateRetries,
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error enqueueing 'create' job in the background")
		// TODO(henrikhodne): Delete the record in the database?
		return nil, err
	}

	return &Instance{
		ID:           id,
		ProviderName: providerName,
		Image:        attr.ImageName,
		State:        "creating",
	}, nil
}

// ProviderCreateInstance is used to schedule the creation of the instance with
// the given ID on the provider selected for that instance.
func (c *Core) ProviderCreateInstance(ctx context.Context, byteID []byte) error {
	id := string(byteID)

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
	}).Info("creating instance")

	dbInstance, err := c.db.GetInstance(id)
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error fetching instance from DB")
		return err
	}

	cloudProvider, err := c.cloudProvider(dbInstance.ProviderName)
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"instance_id":   id,
			"provider_name": dbInstance.ProviderName,
			"err":           err,
		}).Error("couldn't find provider with given name")
		return err
	}

	instance, err := cloudProvider.Create(id, cloud.CreateAttributes{
		ImageName:    dbInstance.Image,
		InstanceType: cloud.InstanceType(dbInstance.InstanceType),
		PublicSSHKey: dbInstance.PublicSSHKey,
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error creating instance")
		return err
	}

	dbInstance.State = "starting"

	err = c.db.UpdateInstance(dbInstance)
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
			"provider_id": instance.ID,
		}).Error("couldn't update instance in DB")
		return err
	}

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
		"provider_id": instance.ID,
	}).Info("created instance")

	return nil
}

// ProviderRefresh is used to synchronize the data on all the cloud providers
// with the data in our database.
func (c *Core) ProviderRefresh(ctx context.Context) error {
	c.refreshProviders()
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()

	var result error

	for providerName, cloudProvider := range c.cloudProviders {
		instances, err := cloudProvider.List()
		if err != nil {
			result = multierror.Append(result, err)
			continue
		}

		for _, instance := range instances {
			dbInstance, err := c.db.GetInstance(instance.ID)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider_name": providerName,
					"provider_id":   instance.ID,
				}).Error("failed fetching instance from database")
				continue
			}

			dbInstance.IPAddress = instance.IPAddress
			dbInstance.State = string(instance.State)

			err = c.db.UpdateInstance(dbInstance)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider":    providerName,
					"provider_id": instance.ID,
					"db_id":       dbInstance.ID,
				}).Error("failed to update instance in database")
			}
		}

		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"provider":       providerName,
			"instance_count": len(instances),
		}).Info("refreshed instances")
	}

	return result
}

// CheckToken is used to check whether a given tokenID+token is in the database.
// Returns (true, nil) iff the token is valid, (false, nil) if the token is
// invalid, and (false, err) if an error occurred while fetching the token.
func (c *Core) CheckToken(tokenID uint64, token string) (bool, error) {
	salt, hash, err := c.db.GetSaltAndHashForTokenID(tokenID)
	if err != nil {
		return false, err
	}

	decodedToken, err := hex.DecodeString(token)
	if err != nil {
		return false, err
	}

	generatedHash, err := scrypt.Key(decodedToken, salt, 16384, 8, 1, 32)
	if err != nil {
		return false, err
	}

	return subtle.ConstantTimeCompare(generatedHash, hash) == 1, nil
}

// cloudProvider is used to get the provider implementation for the cloud
// provider with a given name. Return an error if no cloud provider with the
// given name exists, or if an error occurred refreshing the configuration from
// the database.
func (c *Core) cloudProvider(name string) (cloud.Provider, error) {
	c.cloudProvidersMutex.Lock()
	cloudProvider, ok := c.cloudProviders[name]
	c.cloudProvidersMutex.Unlock()
	if !ok {
		err := c.refreshProviders()
		if err != nil {
			return nil, err
		}

		c.cloudProvidersMutex.Lock()
		cloudProvider, ok = c.cloudProviders[name]
		c.cloudProvidersMutex.Unlock()
		if !ok {
			// This really shouldn't happen, since the database should ensure
			// that a provider with a matching name exists.
			return nil, fmt.Errorf("couldn't find a provider with that name")
		}
	}

	return cloudProvider, nil
}

// refreshProviders is used to regenerate the c.cloudProviders map with the
// configurations stored in the database.
func (c *Core) refreshProviders() error {
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()
	dbCloudProviders, err := c.db.ListProviders()
	if err != nil {
		return err
	}

	cloudProviders := make(map[string]cloud.Provider)

	for _, dbCloudProvider := range dbCloudProviders {
		cloudProvider, err := cloud.NewProvider(dbCloudProvider.Type, dbCloudProvider.Config)
		if err != nil {
			return err
		}

		cloudProviders[dbCloudProvider.Name] = cloudProvider
	}

	c.cloudProviders = cloudProviders

	return nil
}

// Instance is a single compute instance.
type Instance struct {
	ID           string
	ProviderName string
	Image        string
	State        string
	IPAddress    string
}
<|endoftext|>


<commit_before>// This file implements a cluster state machine.  It relies on a cluster
// wide key-value store for coordinating the state of the cluster.
// It also stores the state of the cluster in this key-value store.
package cluster

import (
	"container/list"
	"encoding/gob"
	"errors"
	"net"
	"time"

	"github.com/samalba/dockerclient"

	"github.com/libopenstorage/gossip"
	gossiptypes "github.com/libopenstorage/gossip/types"
	"github.com/libopenstorage/openstorage/api"

	log "github.com/Sirupsen/logrus"

	kv "github.com/portworx/kvdb"
	"github.com/portworx/systemutils"
)

const (
	dockerHost   = "unix:///var/run/docker.sock"
	heartbeatKey = "heartbeat"
)

type ClusterManager struct {
	listeners *list.List
	config    Config
	kv        kv.Kvdb
	status    api.Status
	nodeCache map[string]api.Node // Cached info on the nodes in the cluster.
	docker    *dockerclient.DockerClient
	g         gossip.Gossiper
}

func externalIp() (string, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return "", err
	}
	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // loopback interface
		}
		addrs, err := iface.Addrs()
		if err != nil {
			return "", err
		}
		for _, addr := range addrs {
			var ip net.IP
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
			case *net.IPAddr:
				ip = v.IP
			}
			if ip == nil || ip.IsLoopback() {
				continue
			}
			ip = ip.To4()
			if ip == nil {
				continue // not an ipv4 address
			}
			return ip.String(), nil
		}
	}

	return "", errors.New("Node not connected to the network.")
}

func (c *ClusterManager) AddEventListener(listener ClusterListener) error {
	log.Printf("Adding cluster event listener: %s", listener.String())
	c.listeners.PushBack(listener)
	return nil
}

func (c *ClusterManager) getSelf() *api.Node {
	var node = api.Node{}

	// Get physical node info.
	node.Id = c.config.NodeId
	node.Status = api.StatusOk
	node.Ip, _ = externalIp()
	node.Timestamp = time.Now()

	return &node
}

func (c *ClusterManager) getCurrentState() *api.Node {
	node := c.getSelf()
	s := systemutils.New()

	node.Cpu, _, _ = s.CpuUsage()
	node.Memory = s.MemUsage()
	node.Luns = s.Luns()

	node.Timestamp = time.Now()

	// Get containers running on this system.
	node.Containers, _ = c.docker.ListContainers(true, true, "")

	return node
}

func (c *ClusterManager) initNode(db *Database) (*api.Node, bool) {
	node := c.getSelf()
	c.nodeCache[node.Id] = *node

	_, exists := db.NodeEntries[node.Id]

	// Add us into the database.
	db.NodeEntries[c.config.NodeId] = NodeEntry{Id: node.Id, Ip: node.Ip}

	log.Infof("Node %s joining cluster... \n\tCluster ID: %s\n\tIP: %s",
		c.config.NodeId, c.config.ClusterId, node.Ip)

	return node, exists
}

// Initialize node and alert listeners that we are joining the cluster.
func (c *ClusterManager) joinCluster(db *Database, self *api.Node, exist bool) error {
	var err error

	// If I am already in the cluster map, don't add me again.
	if exist {
		goto found
	}

	// Alert all listeners that we are a new node joining an existing cluster.
	for e := c.listeners.Front(); e != nil; e = e.Next() {
		err = e.Value.(ClusterListener).Init(self, db)
		if err != nil {
			log.Warnf("Failed to initialize %s: %v",
				e.Value.(ClusterListener).String(), err)
			goto done
		}
	}

found:
	// Alert all listeners that we are joining the cluster.
	for e := c.listeners.Front(); e != nil; e = e.Next() {
		err = e.Value.(ClusterListener).Join(self, db)
		if err != nil {
			log.Warnf("Failed to initialize %s: %v",
				e.Value.(ClusterListener).String(), err)
			goto done
		}
	}

	for id, n := range db.NodeEntries {
		if id != c.config.NodeId {
			// Check to see if the IP is the same.  If it is, then we have a stale entry.
			if n.Ip == self.Ip {
				log.Warnf("Warning, Detected node %s with the same IP %s in the database.  Will not connect to this node.",
					id, n.Ip)
			} else {
				log.Infof("Connecting to node %s with IP %s.", id, n.Ip)
				// Gossip with this node.
				c.g.AddNode(n.Ip + ":9002")

				// Assume this node is OK.  We will catch any problems during heartbeating.
				c.nodeCache[id] = api.Node{Status: api.StatusOk, Timestamp: time.Now()}
			}
		}
	}

done:
	return err
}

func (c *ClusterManager) initCluster(db *Database, self *api.Node, exist bool) error {
	err := error(nil)

	// Alert all listeners that we are initializing a new cluster.
	for e := c.listeners.Front(); e != nil; e = e.Next() {
		err = e.Value.(ClusterListener).ClusterInit(self, db)
		if err != nil {
			log.Printf("Failed to initialize %s",
				e.Value.(ClusterListener).String())
			goto done
		}
	}

	err = c.joinCluster(db, self, exist)
	if err != nil {
		log.Printf("Failed to join new cluster")
		goto done
	}

done:
	return err
}

func (c *ClusterManager) heartBeat() {
	for {
		node := c.getCurrentState()
		c.nodeCache[node.Id] = *node

		c.g.UpdateSelf(gossiptypes.StoreKey(heartbeatKey+c.config.ClusterId), *node)

		// Process heartbeats from other nodes...
		gossipValues := c.g.GetStoreKeyValue(gossiptypes.StoreKey(heartbeatKey + c.config.ClusterId))

		for _, nodeInfo := range gossipValues {
			n, ok := nodeInfo.Value.(api.Node)

			if !ok {
				log.Warn("Received a bad broadcast packet: %v", nodeInfo.Value)
				continue
			}

			if n.Id == node.Id {
				continue
			}

			_, ok = c.nodeCache[n.Id]
			if ok {
				if n.Status != api.StatusOk {
					log.Warn("Detected node ", n.Id, " to be unhealthy.")

					for e := c.listeners.Front(); e != nil; e = e.Next() {
						err := e.Value.(ClusterListener).Update(&n)
						if err != nil {
							log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
						}
					}

					delete(c.nodeCache, n.Id)
				} else if time.Since(n.Timestamp) > 60*time.Second {
					log.Warn("Detected node ", n.Id, " to be offline due to inactivity.")

					n.Status = api.StatusOffline

					/*
						for e := c.listeners.Front(); e != nil; e = e.Next() {
							err := e.Value.(ClusterListener).Update(&n)
							if err != nil {
								log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
							}
						}
					*/

					delete(c.nodeCache, n.Id)
				}
			} else if time.Since(n.Timestamp) <= 60*time.Second {
				// A node joined the cluster.
				log.Warn("Detected node ", n.Id, " to be online.")

				c.nodeCache[n.Id] = n
				for e := c.listeners.Front(); e != nil; e = e.Next() {
					err := e.Value.(ClusterListener).Update(&n)
					if err != nil {
						log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
					}
				}
			}
		}

		// Process stale entries in our local cache.
		for id, n := range c.nodeCache {
			if id == node.Id {
				continue
			}

			if time.Since(n.Timestamp) > 60*time.Second {
				log.Warn("Detected node ", n.Id, " to be offline due to a lack of heartbeat.")

				n.Status = api.StatusOffline

				/*
					for e := c.listeners.Front(); e != nil; e = e.Next() {
						 err := e.Value.(ClusterListener).Update(&n)
						if err != nil {
							log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
						}
					}
				*/

				delete(c.nodeCache, id)
			}
		}

		time.Sleep(2 * time.Second)
	}
}

func (c *ClusterManager) Start() error {
	log.Info("Cluster manager starting...")
	kvdb := kv.Instance()

	// Start the gossip protocol.
	// XXX make the port configurable.
	gob.Register(api.Node{})
	c.g = gossip.New("0.0.0.0:9002", gossiptypes.NodeId(c.config.NodeId))
	c.g.SetGossipInterval(2 * time.Second)

	kvlock, err := kvdb.Lock("cluster/lock", 60)
	if err != nil {
		log.Panic("Fatal, Unable to obtain cluster lock.", err)
	}

	db, err := readDatabase()
	if err != nil {
		log.Panic(err)
	}

	if db.Status == api.StatusInit {
		log.Info("Will initialize a new cluster.")

		c.status = api.StatusOk
		db.Status = api.StatusOk
		self, _ := c.initNode(&db)

		err = c.initCluster(&db, self, false)
		if err != nil {
			kvdb.Unlock(kvlock)
			log.Errorf("Failed to initialize the cluster.", err)
			log.Panic(err)
		}

		// Update the new state of the cluster in the KV Database
		err = writeDatabase(&db)
		if err != nil {
			log.Errorf("Failed to save the database.", err)
			log.Panic(err)
		}

		err = kvdb.Unlock(kvlock)
		if err != nil {
			log.Panic("Fatal, unable to unlock cluster... Did something take too long to initialize?", err)
		}
	} else if db.Status&api.StatusOk > 0 {
		log.Info("Cluster state is OK... Joining the cluster.")

		c.status = api.StatusOk
		self, exist := c.initNode(&db)

		err = c.joinCluster(&db, self, exist)
		if err != nil {
			kvdb.Unlock(kvlock)
			log.Panic(err)
		}

		err = writeDatabase(&db)
		if err != nil {
			log.Panic(err)
		}

		err = kvdb.Unlock(kvlock)
		if err != nil {
			log.Panic("Fatal, unable to unlock cluster... Did something take too long to initialize?", err)
		}
	} else {
		kvdb.Unlock(kvlock)
		err = errors.New("Fatal, Cluster is in an unexpected state.")
		log.Panic(err)
	}

	// Start heartbeating to other nodes.
	go c.heartBeat()

	return nil
}

func (c *ClusterManager) Init() error {
	docker, err := dockerclient.NewDockerClient(dockerHost, nil)
	if err != nil {
		log.Printf("Fatal, could not connect to Docker.")
		return err
	}

	c.listeners = list.New()
	c.nodeCache = make(map[string]api.Node)
	c.docker = docker

	return nil
}

func (c *ClusterManager) Enumerate() (api.Cluster, error) {
	i := 0

	cluster := api.Cluster{Id: c.config.ClusterId, Status: c.status}
	cluster.Nodes = make([]api.Node, len(c.nodeCache))
	for _, n := range c.nodeCache {
		cluster.Nodes[i] = n
		i++
	}

	return cluster, nil
}

func (c *ClusterManager) Remove(nodes []api.Node) error {
	// TODO
	return nil
}

func (c *ClusterManager) Shutdown(cluster bool, nodes []api.Node) error {
	// TODO
	return nil
}
<commit_msg>Enable node down events<commit_after>// This file implements a cluster state machine.  It relies on a cluster
// wide key-value store for coordinating the state of the cluster.
// It also stores the state of the cluster in this key-value store.
package cluster

import (
	"container/list"
	"encoding/gob"
	"errors"
	"net"
	"time"

	"github.com/samalba/dockerclient"

	"github.com/libopenstorage/gossip"
	gossiptypes "github.com/libopenstorage/gossip/types"
	"github.com/libopenstorage/openstorage/api"

	log "github.com/Sirupsen/logrus"

	kv "github.com/portworx/kvdb"
	"github.com/portworx/systemutils"
)

const (
	dockerHost   = "unix:///var/run/docker.sock"
	heartbeatKey = "heartbeat"
)

type ClusterManager struct {
	listeners *list.List
	config    Config
	kv        kv.Kvdb
	status    api.Status
	nodeCache map[string]api.Node // Cached info on the nodes in the cluster.
	docker    *dockerclient.DockerClient
	g         gossip.Gossiper
}

func externalIp() (string, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return "", err
	}
	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // loopback interface
		}
		addrs, err := iface.Addrs()
		if err != nil {
			return "", err
		}
		for _, addr := range addrs {
			var ip net.IP
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
			case *net.IPAddr:
				ip = v.IP
			}
			if ip == nil || ip.IsLoopback() {
				continue
			}
			ip = ip.To4()
			if ip == nil {
				continue // not an ipv4 address
			}
			return ip.String(), nil
		}
	}

	return "", errors.New("Node not connected to the network.")
}

func (c *ClusterManager) AddEventListener(listener ClusterListener) error {
	log.Printf("Adding cluster event listener: %s", listener.String())
	c.listeners.PushBack(listener)
	return nil
}

func (c *ClusterManager) getSelf() *api.Node {
	var node = api.Node{}

	// Get physical node info.
	node.Id = c.config.NodeId
	node.Status = api.StatusOk
	node.Ip, _ = externalIp()
	node.Timestamp = time.Now()

	return &node
}

func (c *ClusterManager) getCurrentState() *api.Node {
	node := c.getSelf()
	s := systemutils.New()

	node.Cpu, _, _ = s.CpuUsage()
	node.Memory = s.MemUsage()
	node.Luns = s.Luns()

	node.Timestamp = time.Now()

	// Get containers running on this system.
	node.Containers, _ = c.docker.ListContainers(true, true, "")

	return node
}

func (c *ClusterManager) initNode(db *Database) (*api.Node, bool) {
	node := c.getSelf()
	c.nodeCache[node.Id] = *node

	_, exists := db.NodeEntries[node.Id]

	// Add us into the database.
	db.NodeEntries[c.config.NodeId] = NodeEntry{Id: node.Id, Ip: node.Ip}

	log.Infof("Node %s joining cluster... \n\tCluster ID: %s\n\tIP: %s",
		c.config.NodeId, c.config.ClusterId, node.Ip)

	return node, exists
}

// Initialize node and alert listeners that we are joining the cluster.
func (c *ClusterManager) joinCluster(db *Database, self *api.Node, exist bool) error {
	var err error

	// If I am already in the cluster map, don't add me again.
	if exist {
		goto found
	}

	// Alert all listeners that we are a new node joining an existing cluster.
	for e := c.listeners.Front(); e != nil; e = e.Next() {
		err = e.Value.(ClusterListener).Init(self, db)
		if err != nil {
			log.Warnf("Failed to initialize %s: %v",
				e.Value.(ClusterListener).String(), err)
			goto done
		}
	}

found:
	// Alert all listeners that we are joining the cluster.
	for e := c.listeners.Front(); e != nil; e = e.Next() {
		err = e.Value.(ClusterListener).Join(self, db)
		if err != nil {
			log.Warnf("Failed to initialize %s: %v",
				e.Value.(ClusterListener).String(), err)
			goto done
		}
	}

	for id, n := range db.NodeEntries {
		if id != c.config.NodeId {
			// Check to see if the IP is the same.  If it is, then we have a stale entry.
			if n.Ip == self.Ip {
				log.Warnf("Warning, Detected node %s with the same IP %s in the database.  Will not connect to this node.",
					id, n.Ip)
			} else {
				log.Infof("Connecting to node %s with IP %s.", id, n.Ip)
				// Gossip with this node.
				c.g.AddNode(n.Ip + ":9002")

				// Assume this node is OK.  We will catch any problems during heartbeating.
				c.nodeCache[id] = api.Node{Status: api.StatusOk, Timestamp: time.Now()}
			}
		}
	}

done:
	return err
}

func (c *ClusterManager) initCluster(db *Database, self *api.Node, exist bool) error {
	err := error(nil)

	// Alert all listeners that we are initializing a new cluster.
	for e := c.listeners.Front(); e != nil; e = e.Next() {
		err = e.Value.(ClusterListener).ClusterInit(self, db)
		if err != nil {
			log.Printf("Failed to initialize %s",
				e.Value.(ClusterListener).String())
			goto done
		}
	}

	err = c.joinCluster(db, self, exist)
	if err != nil {
		log.Printf("Failed to join new cluster")
		goto done
	}

done:
	return err
}

func (c *ClusterManager) heartBeat() {
	for {
		node := c.getCurrentState()
		c.nodeCache[node.Id] = *node

		c.g.UpdateSelf(gossiptypes.StoreKey(heartbeatKey+c.config.ClusterId), *node)

		// Process heartbeats from other nodes...
		gossipValues := c.g.GetStoreKeyValue(gossiptypes.StoreKey(heartbeatKey + c.config.ClusterId))

		for _, nodeInfo := range gossipValues {
			n, ok := nodeInfo.Value.(api.Node)

			if !ok {
				log.Warn("Received a bad broadcast packet: %v", nodeInfo.Value)
				continue
			}

			if n.Id == node.Id {
				continue
			}

			_, ok = c.nodeCache[n.Id]
			if ok {
				if n.Status != api.StatusOk {
					log.Warn("Detected node ", n.Id, " to be unhealthy.")

					for e := c.listeners.Front(); e != nil; e = e.Next() {
						err := e.Value.(ClusterListener).Update(&n)
						if err != nil {
							log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
						}
					}

					delete(c.nodeCache, n.Id)
				} else if time.Since(n.Timestamp) > 60*time.Second {
					log.Warn("Detected node ", n.Id, " to be offline due to inactivity.")

					n.Status = api.StatusOffline
					for e := c.listeners.Front(); e != nil; e = e.Next() {
						err := e.Value.(ClusterListener).Update(&n)
						if err != nil {
							log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
						}
					}

					delete(c.nodeCache, n.Id)
				}
			} else if time.Since(n.Timestamp) <= 60*time.Second {
				// A node joined the cluster.
				log.Warn("Detected node ", n.Id, " to be online.")

				c.nodeCache[n.Id] = n
				for e := c.listeners.Front(); e != nil; e = e.Next() {
					err := e.Value.(ClusterListener).Update(&n)
					if err != nil {
						log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
					}
				}
			}
		}

		// Process stale entries in our local cache.
		for id, n := range c.nodeCache {
			if id == node.Id {
				continue
			}

			if time.Since(n.Timestamp) > 60*time.Second {
				log.Warn("Detected node ", n.Id, " to be offline due to a lack of heartbeat.")

				n.Status = api.StatusOffline

				for e := c.listeners.Front(); e != nil; e = e.Next() {
					err := e.Value.(ClusterListener).Update(&n)
					if err != nil {
						log.Warn("Failed to notify ", e.Value.(ClusterListener).String())
					}
				}

				delete(c.nodeCache, id)
			}
		}

		time.Sleep(2 * time.Second)
	}
}

func (c *ClusterManager) Start() error {
	log.Info("Cluster manager starting...")
	kvdb := kv.Instance()

	// Start the gossip protocol.
	// XXX make the port configurable.
	gob.Register(api.Node{})
	c.g = gossip.New("0.0.0.0:9002", gossiptypes.NodeId(c.config.NodeId))
	c.g.SetGossipInterval(2 * time.Second)

	kvlock, err := kvdb.Lock("cluster/lock", 60)
	if err != nil {
		log.Panic("Fatal, Unable to obtain cluster lock.", err)
	}

	db, err := readDatabase()
	if err != nil {
		log.Panic(err)
	}

	if db.Status == api.StatusInit {
		log.Info("Will initialize a new cluster.")

		c.status = api.StatusOk
		db.Status = api.StatusOk
		self, _ := c.initNode(&db)

		err = c.initCluster(&db, self, false)
		if err != nil {
			kvdb.Unlock(kvlock)
			log.Errorf("Failed to initialize the cluster.", err)
			log.Panic(err)
		}

		// Update the new state of the cluster in the KV Database
		err = writeDatabase(&db)
		if err != nil {
			log.Errorf("Failed to save the database.", err)
			log.Panic(err)
		}

		err = kvdb.Unlock(kvlock)
		if err != nil {
			log.Panic("Fatal, unable to unlock cluster... Did something take too long to initialize?", err)
		}
	} else if db.Status&api.StatusOk > 0 {
		log.Info("Cluster state is OK... Joining the cluster.")

		c.status = api.StatusOk
		self, exist := c.initNode(&db)

		err = c.joinCluster(&db, self, exist)
		if err != nil {
			kvdb.Unlock(kvlock)
			log.Panic(err)
		}

		err = writeDatabase(&db)
		if err != nil {
			log.Panic(err)
		}

		err = kvdb.Unlock(kvlock)
		if err != nil {
			log.Panic("Fatal, unable to unlock cluster... Did something take too long to initialize?", err)
		}
	} else {
		kvdb.Unlock(kvlock)
		err = errors.New("Fatal, Cluster is in an unexpected state.")
		log.Panic(err)
	}

	// Start heartbeating to other nodes.
	go c.heartBeat()

	return nil
}

func (c *ClusterManager) Init() error {
	docker, err := dockerclient.NewDockerClient(dockerHost, nil)
	if err != nil {
		log.Printf("Fatal, could not connect to Docker.")
		return err
	}

	c.listeners = list.New()
	c.nodeCache = make(map[string]api.Node)
	c.docker = docker

	return nil
}

func (c *ClusterManager) Enumerate() (api.Cluster, error) {
	i := 0

	cluster := api.Cluster{Id: c.config.ClusterId, Status: c.status}
	cluster.Nodes = make([]api.Node, len(c.nodeCache))
	for _, n := range c.nodeCache {
		cluster.Nodes[i] = n
		i++
	}

	return cluster, nil
}

func (c *ClusterManager) Remove(nodes []api.Node) error {
	// TODO
	return nil
}

func (c *ClusterManager) Shutdown(cluster bool, nodes []api.Node) error {
	// TODO
	return nil
}
<|endoftext|>


<commit_before>package main

func main() {
	Execute()
}
<commit_msg>Initialize random seed<commit_after>package main

import (
	"math/rand"
	"time"
)

func main() {
	rand.Seed(time.Now().UTC().UnixNano())

	Execute()
}
<|endoftext|>


<commit_before>package build

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"html/template"
	"io"
	"math"
	"net"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
)

const F = 0xffffffff

type Node struct {
	IP, Mask, Mask2 uint32
}

type Graph []Node

func (g Graph) Len() int           { return len(g) }
func (g Graph) Swap(i, j int)      { g[i], g[j] = g[j], g[i] }
func (g Graph) Less(i, j int) bool { return g[i].IP < g[j].IP }

var (
	EOL    = byte('\n')
	URL    = "http://ftp.apnic.net/apnic/stats/apnic/delegated-apnic-latest"
	CNIPV4 = regexp.MustCompile(`apnic\|(CN|cn)\|ipv4\|([0-9\.]+)\|([0-9]+)\|([0-9]+)\|a.*`)
)

func ip2long(ipstr string) uint32 {
	ip := net.ParseIP(ipstr)
	if ip == nil {
		return 0
	}
	ip = ip.To4()
	return binary.BigEndian.Uint32(ip)
}

func long2ip(ipLong uint32) string {
	ipByte := make([]byte, 4)
	binary.BigEndian.PutUint32(ipByte, ipLong)
	ip := net.IP(ipByte)
	return ip.String()
}

func fetchIPData(results *Graph) (err error) {
	var (
		n   int
		buf string
		r   *bufio.Reader

		startIP, prevIP, smask string
		numIP, imask           int
	)

	fmt.Println("Fetching data from apnic.net, it might take a few minutes, please wait...")

	res, err := http.Get(URL)
	if err != nil {
		panic(err)
	}
	r = bufio.NewReader(res.Body)
	defer res.Body.Close()
	for {
		// read line by line
		buf, err = r.ReadString(EOL)
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		matches := CNIPV4.FindStringSubmatch(buf)
		if len(matches) > 0 {
			n++
			startIP = matches[2]
			numIP, _ = strconv.Atoi(matches[3])

			imask = F ^ (numIP - 1)
			smask = fmt.Sprintf("%02x", imask)
			mask := [4]string{}
			mask[0] = smask[0:2]
			mask[1] = smask[2:4]
			mask[2] = "0"
			mask[3] = "0"

			for i, s := range mask[:2] {
				num, _ := strconv.ParseInt(s, 16, 10)
				mask[i] = fmt.Sprintf("%d", num)
			}

			mask2 := 32 - uint32(math.Log2(float64(numIP)))

			ip := strings.Split(startIP, ".")
			ip[2] = "0"
			ip[3] = "0"
			startIP = strings.Join(ip, ".")

			maskIP := fmt.Sprintf("%s.%s.%s.%s", mask[0], mask[1], mask[2], mask[3])

			if startIP != prevIP {
				*results = append(*results, Node{ip2long(startIP), ip2long(maskIP), mask2})
				prevIP = startIP
			}
		}
	}
	return nil
}

func Action(c *cli.Context) {
	var pacfile = "go.pac"
	var results = make(Graph, 0)
	results = append(results, Node{ip2long("127.0.0.1"), ip2long("255.0.0.0"), 0})
	results = append(results, Node{ip2long("10.0.0.0"), ip2long("255.0.0.0"), 0})
	results = append(results, Node{ip2long("127.0.0.1"), ip2long("255.240.0.0"), 0})
	results = append(results, Node{ip2long("192.168.0.0"), ip2long("255.255.0.0"), 0})
	fetchIPData(&results)
	sort.Sort(results)

	file, err := os.Create(pacfile)
	if err != nil {
		panic(err)
	}

	t, err := template.ParseFiles("templates/pac.tmpl")
	data := make(map[string]interface{}, 0)
	data["Graph"] = results
	data["Proxy"] = c.String("proxy")
	err = t.Execute(file, data)
	if err != nil {
		panic(err)
	}

	fmt.Printf("Rules: %d items.\n", len(results))
	fmt.Printf("Usage: Use the newly created %s as your web browser's automatic \n", pacfile)
	fmt.Printf("PAC(Proxy auto-config) file.\n")
}
<commit_msg>fix, 386 overflows int<commit_after>package build

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"html/template"
	"io"
	"math"
	"net"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
)

const F = 0xffffff

type Node struct {
	IP, Mask, Mask2 uint32
}

type Graph []Node

func (g Graph) Len() int           { return len(g) }
func (g Graph) Swap(i, j int)      { g[i], g[j] = g[j], g[i] }
func (g Graph) Less(i, j int) bool { return g[i].IP < g[j].IP }

var (
	EOL    = byte('\n')
	URL    = "http://ftp.apnic.net/apnic/stats/apnic/delegated-apnic-latest"
	CNIPV4 = regexp.MustCompile(`apnic\|(CN|cn)\|ipv4\|([0-9\.]+)\|([0-9]+)\|([0-9]+)\|a.*`)
)

func ip2long(ipstr string) uint32 {
	ip := net.ParseIP(ipstr)
	if ip == nil {
		return 0
	}
	ip = ip.To4()
	return binary.BigEndian.Uint32(ip)
}

func long2ip(ipLong uint32) string {
	ipByte := make([]byte, 4)
	binary.BigEndian.PutUint32(ipByte, ipLong)
	ip := net.IP(ipByte)
	return ip.String()
}

func fetchIPData(results *Graph) (err error) {
	var (
		n   int
		buf string
		r   *bufio.Reader

		startIP, prevIP, smask string
		numIP, imask           int
	)

	fmt.Println("Fetching data from apnic.net, it might take a few minutes, please wait...")

	res, err := http.Get(URL)
	if err != nil {
		panic(err)
	}
	r = bufio.NewReader(res.Body)
	defer res.Body.Close()
	for {
		// read line by line
		buf, err = r.ReadString(EOL)
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		matches := CNIPV4.FindStringSubmatch(buf)
		if len(matches) > 0 {
			n++
			startIP = matches[2]
			numIP, _ = strconv.Atoi(matches[3])

			imask = F ^ (numIP - 1)
			smask = fmt.Sprintf("%02x", imask)
			mask := [4]string{}
			mask[0] = smask[0:2]
			mask[1] = smask[2:4]
			mask[2] = "0"
			mask[3] = "0"

			for i, s := range mask[:2] {
				num, _ := strconv.ParseInt(s, 16, 10)
				mask[i] = fmt.Sprintf("%d", num)
			}

			mask2 := 32 - uint32(math.Log2(float64(numIP)))

			ip := strings.Split(startIP, ".")
			ip[2] = "0"
			ip[3] = "0"
			startIP = strings.Join(ip, ".")

			maskIP := fmt.Sprintf("%s.%s.%s.%s", mask[0], mask[1], mask[2], mask[3])

			if startIP != prevIP {
				*results = append(*results, Node{ip2long(startIP), ip2long(maskIP), mask2})
				prevIP = startIP
			}
		}
	}
	return nil
}

func Action(c *cli.Context) {
	var pacfile = "go.pac"
	var results = make(Graph, 0)
	results = append(results, Node{ip2long("127.0.0.1"), ip2long("255.0.0.0"), 0})
	results = append(results, Node{ip2long("10.0.0.0"), ip2long("255.0.0.0"), 0})
	results = append(results, Node{ip2long("127.0.0.1"), ip2long("255.240.0.0"), 0})
	results = append(results, Node{ip2long("192.168.0.0"), ip2long("255.255.0.0"), 0})
	fetchIPData(&results)
	sort.Sort(results)

	file, err := os.Create(pacfile)
	if err != nil {
		panic(err)
	}

	t, err := template.ParseFiles("templates/pac.tmpl")
	data := make(map[string]interface{}, 0)
	data["Graph"] = results
	data["Proxy"] = c.String("proxy")
	err = t.Execute(file, data)
	if err != nil {
		panic(err)
	}

	fmt.Printf("Rules: %d items.\n", len(results))
	fmt.Printf("Usage: Use the newly created %s as your web browser's automatic \n", pacfile)
	fmt.Printf("PAC(Proxy auto-config) file.\n")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"

	"gopkg.in/alecthomas/kingpin.v2"
)

// #cgo windows LDFLAGS: -Wl,--allow-multiple-definition -static
import "C"

var (
	version = "head" // set by command-line on CI release builds
	app     = kingpin.New("butler", "Your very own itch.io helper")

	dlCmd = app.Command("dl", "Download a file (resumes if can, checks hashes)")

	pushCmd = app.Command("push", "Upload a new version of something to itch.io")
)

var appArgs = struct {
	json       *bool
	quiet      *bool
	timestamps *bool
}{
	app.Flag("json", "Enable machine-readable JSON-lines output").Short('j').Bool(),
	app.Flag("quiet", "Hide progress indicators & other extra info").Short('q').Bool(),
	app.Flag("timestamps", "Prefix all output by timestamps (for logging purposes)").Bool(),
}

var dlArgs = struct {
	url  *string
	dest *string
}{
	dlCmd.Arg("url", "Address to download from").Required().String(),
	dlCmd.Arg("dest", "File to write downloaded data to").Required().String(),
}

var pushArgs = struct {
	identity *string
	address  *string
	src      *string
	repo     *string
}{
	pushCmd.Flag("identity", "Path to the private key used for public key authentication.").Default(fmt.Sprintf("%s/%s", os.Getenv("HOME"), ".ssh/id_rsa")).Short('i').ExistingFile(),
	pushCmd.Flag("address", "Specify wharf address (advanced)").Default("wharf.itch.zone").Short('a').Hidden().String(),
	pushCmd.Arg("src", "Directory or zip archive to upload, e.g.").Required().ExistingFileOrDir(),
	pushCmd.Arg("repo", "Repository to push to, e.g. leafo/xmoon:win64").Required().String(),
}

func main() {
	app.HelpFlag.Short('h')
	app.Version(version)
	app.VersionFlag.Short('V')

	cmd, err := app.Parse(os.Args[1:])
	if *appArgs.timestamps {
		log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	} else {
		log.SetFlags(0)
	}

	switch kingpin.MustParse(cmd, err) {
	case dlCmd.FullCommand():
		dl(*dlArgs.url, *dlArgs.dest)

	case pushCmd.FullCommand():
		push(*pushArgs.src, *pushArgs.repo)
	}
}
<commit_msg>statically compile on linux please!<commit_after>package main

import (
	"fmt"
	"log"
	"os"

	"gopkg.in/alecthomas/kingpin.v2"
)

// #cgo windows LDFLAGS: -Wl,--allow-multiple-definition -static
// #cgo linux LDFLAGS: -static
import "C"

var (
	version = "head" // set by command-line on CI release builds
	app     = kingpin.New("butler", "Your very own itch.io helper")

	dlCmd = app.Command("dl", "Download a file (resumes if can, checks hashes)")

	pushCmd = app.Command("push", "Upload a new version of something to itch.io")
)

var appArgs = struct {
	json       *bool
	quiet      *bool
	timestamps *bool
}{
	app.Flag("json", "Enable machine-readable JSON-lines output").Short('j').Bool(),
	app.Flag("quiet", "Hide progress indicators & other extra info").Short('q').Bool(),
	app.Flag("timestamps", "Prefix all output by timestamps (for logging purposes)").Bool(),
}

var dlArgs = struct {
	url  *string
	dest *string
}{
	dlCmd.Arg("url", "Address to download from").Required().String(),
	dlCmd.Arg("dest", "File to write downloaded data to").Required().String(),
}

var pushArgs = struct {
	identity *string
	address  *string
	src      *string
	repo     *string
}{
	pushCmd.Flag("identity", "Path to the private key used for public key authentication.").Default(fmt.Sprintf("%s/%s", os.Getenv("HOME"), ".ssh/id_rsa")).Short('i').ExistingFile(),
	pushCmd.Flag("address", "Specify wharf address (advanced)").Default("wharf.itch.zone").Short('a').Hidden().String(),
	pushCmd.Arg("src", "Directory or zip archive to upload, e.g.").Required().ExistingFileOrDir(),
	pushCmd.Arg("repo", "Repository to push to, e.g. leafo/xmoon:win64").Required().String(),
}

func main() {
	app.HelpFlag.Short('h')
	app.Version(version)
	app.VersionFlag.Short('V')

	cmd, err := app.Parse(os.Args[1:])
	if *appArgs.timestamps {
		log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	} else {
		log.SetFlags(0)
	}

	switch kingpin.MustParse(cmd, err) {
	case dlCmd.FullCommand():
		dl(*dlArgs.url, *dlArgs.dest)

	case pushCmd.FullCommand():
		push(*pushArgs.src, *pushArgs.repo)
	}
}
<|endoftext|>


<commit_before>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Package cmd cmdline client for the Baruwa REST API
package cmd
<commit_msg>TST: Add client tests<commit_after>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Package cmd cmdline client for the Baruwa REST API
package cmd

import (
	"testing"

	"github.com/baruwa-enterprise/baruwa-go/api"
)

func TestGetClientError(t *testing.T) {
	var err error
	*serverURL = ""
	*apiToken = ""

	if _, err = GetClient(); err == nil {
		t.Fatalf("An error should be returned")
	}
	if err.Error() != unsetVarsErr {
		t.Errorf("Expected %s got %s", unsetVarsErr, err.Error())
	}
}

func TestGetClientOK(t *testing.T) {
	var err error
	var c *api.Client
	token := "test-token"
	su := "https://baruwa.example.com"
	*serverURL = su
	*apiToken = token

	if c, err = GetClient(); err != nil {
		t.Fatalf("An error should not be returned")
	}
	if c.BaseURL.String() != su {
		t.Errorf("Expected %s got %s", su, c.BaseURL)
	}
}
<|endoftext|>


<commit_before>// estab exports elasticsearch fields as tab separated values
package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"

	"github.com/belogik/goes"
	"github.com/miku/estab"
)

func main() {

	host := flag.String("host", "localhost", "elasticsearch host")
	port := flag.String("port", "9200", "elasticsearch port")
	indicesString := flag.String("indices", "", "indices to search (or all)")
	fieldsString := flag.String("f", "_id _index", "field or fields space separated")
	timeout := flag.String("timeout", "10m", "scroll timeout")
	size := flag.Int("size", 10000, "scroll batch size")
	nullValue := flag.String("null", "NOT_AVAILABLE", "value for empty fields")
	separator := flag.String("separator", "|", "separator to use for multiple field values")
	delimiter := flag.String("delimiter", "\t", "column delimiter")
	limit := flag.Int("limit", -1, "maximum number of docs to return (return all by default)")
	version := flag.Bool("v", false, "prints current program version")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	queryString := flag.String("query", "", "custom query to run")
	raw := flag.Bool("raw", false, "stream out the raw json records")
	header := flag.Bool("header", false, "output header row with field names")
	singleValue := flag.Bool("1", false, "one value per line (works only with a single column in -f)")
	zeroAsNull := flag.Bool("zero-as-null", false, "treat zero length strings as null values")
	precision := flag.Int("precision", 0, "precision for numeric output")

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if *version {
		fmt.Println(estab.Version)
		os.Exit(0)
	}

	var query map[string]interface{}
	if *queryString == "" {
		query = map[string]interface{}{
			"query": map[string]interface{}{
				"match_all": map[string]interface{}{},
			},
		}
	} else {
		err := json.Unmarshal([]byte(*queryString), &query)
		if err != nil {
			log.Fatal(err)
		}
	}

	indices := strings.Fields(*indicesString)
	fields := strings.Fields(*fieldsString)

	if *raw && *singleValue {
		log.Fatal("-1 xor -raw ")
	}

	if *singleValue && len(fields) > 1 {
		log.Fatalf("-1 works only with a single column, %d given: %s\n", len(fields), strings.Join(fields, " "))
	}

	if !*raw {
		query["fields"] = fields
	}

	conn := goes.NewConnection(*host, *port)
	scanResponse, err := conn.Scan(query, indices, []string{""}, *timeout, *size)
	if err != nil {
		log.Fatal(err)
	}

	w := bufio.NewWriter(os.Stdout)
	defer w.Flush()
	i := 0

	if *header {
		fmt.Fprintln(w, strings.Join(fields, *delimiter))
	}

	for {
		scrollResponse, err := conn.Scroll(scanResponse.ScrollId, *timeout)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		if len(scrollResponse.Hits.Hits) == 0 {
			break
		}
		for _, hit := range scrollResponse.Hits.Hits {
			if i == *limit {
				return
			}
			if *raw {
				b, err := json.Marshal(hit)
				if err != nil {
					log.Fatal(err)
				}
				fmt.Fprintln(w, string(b))
				continue
			}

			var columns []string
			for _, f := range fields {
				var c []string
				switch f {
				case "_id":
					c = append(c, hit.Id)
				case "_index":
					c = append(c, hit.Index)
				case "_type":
					c = append(c, hit.Type)
				case "_score":
					c = append(c, strconv.FormatFloat(hit.Score, 'f', 6, 64))
				default:
					switch value := hit.Fields[f].(type) {
					case nil:
						c = []string{*nullValue}
					case []interface{}:
						for _, e := range value {
							switch e.(type) {
							case string:
								s := e.(string)
								if s == "" && *zeroAsNull {
									c = append(c, *nullValue)
								} else {
									c = append(c, e.(string))
								}
							case float64:
								c = append(c, strconv.FormatFloat(e.(float64), 'f', *precision, 64))
							}
						}
					default:
						log.Fatalf("unknown field type in response: %+v\n", hit.Fields[f])
					}
				}
				if *singleValue {
					for _, value := range c {
						fmt.Fprintln(w, value)
					}
				} else {
					columns = append(columns, strings.Join(c, *separator))
				}
			}
			if !*singleValue {
				fmt.Fprintln(w, strings.Join(columns, *delimiter))
			}
			i++
		}
	}
}
<commit_msg>fix formatting<commit_after>// estab exports elasticsearch fields as tab separated values
package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"

	"github.com/belogik/goes"
	"github.com/miku/estab"
)

func main() {

	host := flag.String("host", "localhost", "elasticsearch host")
	port := flag.String("port", "9200", "elasticsearch port")
	indicesString := flag.String("indices", "", "indices to search (or all)")
	fieldsString := flag.String("f", "_id _index", "field or fields space separated")
	timeout := flag.String("timeout", "10m", "scroll timeout")
	size := flag.Int("size", 10000, "scroll batch size")
	nullValue := flag.String("null", "NOT_AVAILABLE", "value for empty fields")
	separator := flag.String("separator", "|", "separator to use for multiple field values")
	delimiter := flag.String("delimiter", "\t", "column delimiter")
	limit := flag.Int("limit", -1, "maximum number of docs to return (return all by default)")
	version := flag.Bool("v", false, "prints current program version")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	queryString := flag.String("query", "", "custom query to run")
	raw := flag.Bool("raw", false, "stream out the raw json records")
	header := flag.Bool("header", false, "output header row with field names")
	singleValue := flag.Bool("1", false, "one value per line (works only with a single column in -f)")
	zeroAsNull := flag.Bool("zero-as-null", false, "treat zero length strings as null values")
	precision := flag.Int("precision", 0, "precision for numeric output")

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if *version {
		fmt.Println(estab.Version)
		os.Exit(0)
	}

	var query map[string]interface{}
	if *queryString == "" {
		query = map[string]interface{}{
			"query": map[string]interface{}{
				"match_all": map[string]interface{}{},
			},
		}
	} else {
		err := json.Unmarshal([]byte(*queryString), &query)
		if err != nil {
			log.Fatal(err)
		}
	}

	indices := strings.Fields(*indicesString)
	fields := strings.Fields(*fieldsString)

	if *raw && *singleValue {
		log.Fatal("-1 xor -raw ")
	}

	if *singleValue && len(fields) > 1 {
		log.Fatalf("-1 works only with a single column, %d given: %s\n", len(fields), strings.Join(fields, " "))
	}

	if !*raw {
		query["fields"] = fields
	}

	conn := goes.NewConnection(*host, *port)
	scanResponse, err := conn.Scan(query, indices, []string{""}, *timeout, *size)
	if err != nil {
		log.Fatal(err)
	}

	w := bufio.NewWriter(os.Stdout)
	defer w.Flush()
	i := 0

	if *header {
		fmt.Fprintln(w, strings.Join(fields, *delimiter))
	}

	for {
		scrollResponse, err := conn.Scroll(scanResponse.ScrollId, *timeout)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		if len(scrollResponse.Hits.Hits) == 0 {
			break
		}
		for _, hit := range scrollResponse.Hits.Hits {
			if i == *limit {
				return
			}
			if *raw {
				b, err := json.Marshal(hit)
				if err != nil {
					log.Fatal(err)
				}
				fmt.Fprintln(w, string(b))
				continue
			}

			var columns []string
			for _, f := range fields {
				var c []string
				switch f {
				case "_id":
					c = append(c, hit.Id)
				case "_index":
					c = append(c, hit.Index)
				case "_type":
					c = append(c, hit.Type)
				case "_score":
					c = append(c, strconv.FormatFloat(hit.Score, 'f', 6, 64))
				default:
					switch value := hit.Fields[f].(type) {
					case nil:
						c = []string{*nullValue}
					case []interface{}:
						for _, e := range value {
							switch e.(type) {
							case string:
								s := e.(string)
								if s == "" && *zeroAsNull {
									c = append(c, *nullValue)
								} else {
									c = append(c, e.(string))
								}
							case float64:
								c = append(c, strconv.FormatFloat(e.(float64), 'f', *precision, 64))
							}
						}
					default:
						log.Fatalf("unknown field type in response: %+v\n", hit.Fields[f])
					}
				}
				if *singleValue {
					for _, value := range c {
						fmt.Fprintln(w, value)
					}
				} else {
					columns = append(columns, strings.Join(c, *separator))
				}
			}
			if !*singleValue {
				fmt.Fprintln(w, strings.Join(columns, *delimiter))
			}
			i++
		}
	}
}
<|endoftext|>