<commit_before>package rnn

import (
	"bytes"
	"encoding/gob"
	"log"
	"math"
	"math/rand"
	"sync"
	"time"

	"github.com/gonum/matrix/mat64"
	"github.com/kelseyhightower/envconfig"
)

// RNN represents the neural network
// This RNNs parameters are the three mat64rices whh, wxh, why.
// hprev is the last known hidden vector, which is actually the memory of the RNN
// bh, and by are the biais vectors respectivly for the hidden layer and the output layer
type RNN struct {
	sync.Mutex
	whh *mat64.Dense // size is hiddenDimension * hiddenDimension
	wxh *mat64.Dense //
	why *mat64.Dense //
	// This is the last known hidden vector that represents the memory of the RNN
	// This is used only for training
	hprev  []float64
	bh     []float64 // This is the biais
	by     []float64 // This is the biais
	config neuralNetConfig
}

// GobDecode the rnn for restoring
func (rnn *RNN) GobDecode(b []byte) error {
	input := bytes.NewBuffer(b)
	dec := gob.NewDecoder(input) // Will read from network.

	type bkp struct {
		Whh *mat64.Dense // size is hiddenDimension * hiddenDimension
		Wxh *mat64.Dense //
		Why *mat64.Dense //
		// This is the last known hidden vector that represents the memory of the RNN
		// This is used only for training
		Hprev  []float64
		Bh     []float64 // This is the biais
		By     []float64 // This is the biais
		Config neuralNetConfig
	}
	var backup bkp
	err := dec.Decode(&backup)
	rnn.bh = make([]float64, len(backup.Bh))
	rnn.by = make([]float64, len(backup.By))
	rnn.hprev = make([]float64, len(backup.Hprev))
	if err == nil {
		*rnn.whh = *backup.Whh
		*rnn.why = *backup.Why
		*rnn.wxh = *backup.Wxh
		rnn.config = backup.Config
		copy(rnn.bh, backup.Bh)
		copy(rnn.by, backup.By)
		copy(rnn.hprev, backup.Hprev)
	}
	return err
}

// GobEncode the RNN for backup
func (rnn *RNN) GobEncode() ([]byte, error) {
	var output bytes.Buffer // Stand-in for a network connection

	enc := gob.NewEncoder(&output) // Will write to network.

	type bkp struct {
		Whh *mat64.Dense // size is hiddenDimension * hiddenDimension
		Wxh *mat64.Dense //
		Why *mat64.Dense //
		// This is the last known hidden vector that represents the memory of the RNN
		// This is used only for training
		Hprev  []float64
		Bh     []float64 // This is the biais
		By     []float64 // This is the biais
		Config neuralNetConfig
	}
	rnn.Lock()
	err := enc.Encode(bkp{
		rnn.whh,
		rnn.wxh,
		rnn.why,
		rnn.hprev,
		rnn.bh,
		rnn.by,
		rnn.config,
	})
	rnn.Unlock()
	return output.Bytes(), err
}

// NewRNN creates a new RNN with input size of x, outputsize of y and hidden dimension of h
// The hidden state h is initialized with the zero vectornn.
//func newRNN(x, y, h int) *RNN {
func NewRNN(inputNeurons, outputNeurons int) *RNN {
	var conf neuralNetConfig
	if inputNeurons == 0 || outputNeurons == 0 {
		err := envconfig.Usage("RNN", &conf)
		if err != nil {
			log.Fatal(err)
		}
		return nil
	}
	err := envconfig.Process("RNN", &conf)
	if err != nil {
		log.Fatal(err)
	}

	//func NewRNN(config NeuralNetConfig) *RNN {
	var rnn RNN
	conf.InputNeurons = inputNeurons
	conf.OutputNeurons = outputNeurons
	rnn.config = conf
	// Initialize biases/weights.

	rnn.wxh = mat64.NewDense(conf.HiddenNeurons, conf.InputNeurons, nil)
	rnn.whh = mat64.NewDense(conf.HiddenNeurons, conf.HiddenNeurons, nil)
	rnn.why = mat64.NewDense(conf.OutputNeurons, conf.HiddenNeurons, nil)
	rnn.bh = make([]float64, conf.HiddenNeurons)
	rnn.by = make([]float64, conf.OutputNeurons)
	wHiddenRaw := rnn.wxh.RawMatrix().Data
	wHiddenHiddenRaw := rnn.whh.RawMatrix().Data
	wOutRaw := rnn.why.RawMatrix().Data

	for _, param := range [][]float64{
		wHiddenRaw,
		wHiddenHiddenRaw,
		wOutRaw,
	} {
		for i := range param {
			randSource := rand.NewSource(time.Now().UnixNano())
			randGen := rand.New(randSource)
			param[i] = randGen.NormFloat64() * conf.RandomFactor
		}
	}

	// initialise the hidden vector to zero
	rnn.hprev = make([]float64, conf.HiddenNeurons)

	return &rnn
}

// RNNs have a deceptively simple API:
// They accept an input vector x and give you an output vector y.
// However, crucially this output vector’s contents are influenced
// not only by the input you just fed in,
// but also on the entire history of inputs you’ve fed in in the past.
// Written as a class, the RNN’s API consists of a single step function:
func (rnn *RNN) step(x, hprev []float64) (y, h []float64) {
	h = tanh(
		add(
			dot(rnn.wxh, x),
			dot(rnn.whh, hprev),
			rnn.bh,
		))
	y = add(
		dot(rnn.why, h),
		rnn.by)
	return
}

// forwardPass takes a mat64rix of inputs and returns
// the corresponding outputs mat64rix
// and a mat64rix of the hidden states that will be used
// for the backpropagation
func (rnn *RNN) forwardPass(xs [][]float64, hprev []float64) (ys, hs [][]float64) {
	inputSize := len(xs)
	hp := make([]float64, len(hprev))
	copy(hp, hprev)
	// un-normalized log probabilities for next chars
	ys = make([][]float64, inputSize)
	hs = make([][]float64, inputSize)
	for t := 0; t < inputSize; t++ {
		// Initialization of the arrays
		ys[t] = make([]float64, rnn.config.OutputNeurons)
		hs[t] = make([]float64, rnn.config.HiddenNeurons)
		ys[t], hs[t] = rnn.step(xs[t], hp)
		hp = hs[t]
	}
	return
}

// Do a backpropagation of the RNNs and returns the derivates
// xs is the input mat64rix
// ts is the target mat64rices
// ps is the normalized log probability
// hs is a mat64rix of hidden vector
func (rnn *RNN) backPropagation(xs, ps, hs, ts [][]float64) (dwxh, dwhh, dwhy *mat64.Dense, dbh, dby []float64) {
	inputSize := len(xs)
	dwxh = mat64.NewDense(rnn.config.HiddenNeurons, rnn.config.InputNeurons, nil)
	dwhh = mat64.NewDense(rnn.config.HiddenNeurons, rnn.config.HiddenNeurons, nil)
	dwhy = mat64.NewDense(rnn.config.OutputNeurons, rnn.config.HiddenNeurons, nil)
	dhnext := make([]float64, rnn.config.OutputNeurons)
	dbh = make([]float64, rnn.config.HiddenNeurons)
	dby = make([]float64, rnn.config.OutputNeurons)
	dhraw := make([]float64, len(hs[0]))

	for t := inputSize - 1; t >= 0; t-- {
		dy := make([]float64, rnn.config.OutputNeurons)
		for i := range ps[t] {
			dy[i] = ps[t][i] - ts[t][i]
		}
		dwhy.Add(dwhy,
			dotVec(dy, hs[t]),
		)
		dby = add(dby, dy)

		dh := add(
			dot(rnn.why.T(), dy),
			dhnext,
		)

		for i := range hs[t] {
			dhraw[i] = (1 - hs[t][i]*hs[t][i]) * dh[i]
		}

		dbh = add(dbh, dhraw)
		dwxh.Add(dwxh, dotVec(dhraw, xs[t]))
		if t == 0 {
			dwhh.Add(dwhh, dotVec(dhraw, hs[inputSize-1]))
		} else {
			dwhh.Add(dwhh, dotVec(dhraw, hs[t-1]))
		}
		dhnext = dot(rnn.whh.T(), dhraw)
	}

	return
}

// TrainingSet represents an input mat64rix and the expected
// result when passed through a rnn
type TrainingSet struct {
	Inputs  [][]float64
	Targets [][]float64
}

// Train the network.
// The train mechanisme is launched in a seperate go-routine
// it is waiting for an input to be sent in the feeding channel
func (rnn *RNN) Train() (feed chan TrainingSet, info chan float64) {
	feed = make(chan TrainingSet, 1)
	info = make(chan float64, 1)

	adagrad := newAdagrad(rnn.config)
	go func(feed chan TrainingSet, info chan float64) {
		// When we have new data
		for tset := range feed {
			// Forward pass
			xs := make([][]float64, len(tset.Inputs))
			copy(xs, tset.Inputs)
			ts := make([][]float64, len(tset.Targets))
			copy(ts, tset.Targets)
			rnn.Lock()
			ys, hs := rnn.forwardPass(xs, rnn.hprev)
			// Save the last state for future training
			copy(rnn.hprev, hs[len(hs)-1])
			//rnn.hprev = hs[len(hs)-1]
			rnn.Unlock()
			ps := normalizeByRow(ys)
			// Loss evaluation
			loss := float64(0)
			for t := 0; t < len(ps); t++ {
				l := float64(0)
				for i := 0; i < len(ps[t]); i++ {
					l += ps[t][i] * ts[t][i]
				}
				loss -= math.Log(l)
			}
			// Send info on a non blocking channel
			select {
			case info <- loss:
			default:
			}

			// Backpass
			rnn.Lock()
			dwxh, dwhh, dwhy, dbh, dby := rnn.backPropagation(xs, ps, hs, ts)
			rnn.Unlock()
			// Clip to mitigate exploding gradients
			for _, param := range [][]float64{
				dwxh.RawMatrix().Data,
				dwhh.RawMatrix().Data,
				dwhy.RawMatrix().Data,
				dby,
				dbh,
			} {
				func(param []float64) {
					for i := range param {
						if param[i] > 5 {
							param[i] = 5
						}
						if param[i] < -5 {
							param[i] = -5
						}
					}
				}(param)
			}
			// Adaptation
			rnn.Lock()
			adagrad.apply(rnn, dwxh, dwhh, dwhy, dbh, dby)
			rnn.Unlock()
		}
	}(feed, info)
	return feed, info
}

// Predict n element of  output that corresponds to the input xs
// At every iteration, the output is processed by the adapt function
func (rnn *RNN) Predict(xs [][]float64, n int, adapt func([]float64) []float64) [][]float64 {
	ys := make([][]float64, n+len(xs))
	h := make([]float64, len(rnn.hprev))
	y := make([]float64, rnn.config.OutputNeurons)
	for i := 0; i < n+len(xs); i++ {
		x := make([]float64, rnn.config.InputNeurons)
		if i < len(xs) {
			copy(x, xs[i])
		} else {
			copy(x, ys[i-1])
		}

		yr, hr := rnn.step(x, h)
		copy(y, yr)
		copy(h, hr)
		expY := exp(y)
		p := div(expY, sum(expY))
		ys[i] = p
		if i < len(xs) {
			for j := 0; j < len(xs[i]); j++ {
				if xs[i][j] == float64(1) {
					ys[i][j] = float64(1)
				}
			}
		} else {
			ys[i] = adapt(p)
		}
	}
	res := make([][]float64, n)
	copy(res, ys[len(xs):])
	return res
}

// Sample the rnn
func (rnn *RNN) Sample(xs [][]float64, n int, choose func([]float64) int) []int {
	res := make([]int, n+len(xs))
	h := make([]float64, len(rnn.hprev))
	//copy(h, rnn.hprev)
	y := make([]float64, rnn.config.OutputNeurons)
	for i := 0; i < n; i++ {
		x := make([]float64, rnn.config.InputNeurons)
		if i < len(xs) {
			copy(x, xs[i])
		} else {
			x[res[i-1]] = 1
		}

		yr, hr := rnn.step(x, h)
		copy(y, yr)
		copy(h, hr)
		expY := exp(y)
		p := div(expY, sum(expY))
		if i < len(xs) {
			for j := 0; j < len(xs[i]); j++ {
				if xs[i][j] == float64(1) {
					res[i] = j
				}
			}
		} else {
			res[i] = choose(p)
		}
	}

	return res
}
<commit_msg>New function: CopyOf<commit_after>package rnn

import (
	"bytes"
	"encoding/gob"
	"log"
	"math"
	"math/rand"
	"sync"
	"time"

	"github.com/gonum/matrix/mat64"
	"github.com/kelseyhightower/envconfig"
)

// RNN represents the neural network
// This RNNs parameters are the three mat64rices whh, wxh, why.
// hprev is the last known hidden vector, which is actually the memory of the RNN
// bh, and by are the biais vectors respectivly for the hidden layer and the output layer
type RNN struct {
	sync.Mutex
	whh *mat64.Dense // size is hiddenDimension * hiddenDimension
	wxh *mat64.Dense //
	why *mat64.Dense //
	// This is the last known hidden vector that represents the memory of the RNN
	// This is used only for training
	hprev  []float64
	bh     []float64 // This is the biais
	by     []float64 // This is the biais
	config neuralNetConfig
}

// GobDecode the rnn for restoring
func (rnn *RNN) GobDecode(b []byte) error {
	input := bytes.NewBuffer(b)
	dec := gob.NewDecoder(input) // Will read from network.

	type bkp struct {
		Whh *mat64.Dense // size is hiddenDimension * hiddenDimension
		Wxh *mat64.Dense //
		Why *mat64.Dense //
		// This is the last known hidden vector that represents the memory of the RNN
		// This is used only for training
		Hprev  []float64
		Bh     []float64 // This is the biais
		By     []float64 // This is the biais
		Config neuralNetConfig
	}
	var backup bkp
	err := dec.Decode(&backup)
	rnn.bh = make([]float64, len(backup.Bh))
	rnn.by = make([]float64, len(backup.By))
	rnn.hprev = make([]float64, len(backup.Hprev))
	if err == nil {
		*rnn.whh = *backup.Whh
		*rnn.why = *backup.Why
		*rnn.wxh = *backup.Wxh
		rnn.config = backup.Config
		copy(rnn.bh, backup.Bh)
		copy(rnn.by, backup.By)
		copy(rnn.hprev, backup.Hprev)
	}
	return err
}

// GobEncode the RNN for backup
func (rnn *RNN) GobEncode() ([]byte, error) {
	var output bytes.Buffer // Stand-in for a network connection

	enc := gob.NewEncoder(&output) // Will write to network.

	type bkp struct {
		Whh *mat64.Dense // size is hiddenDimension * hiddenDimension
		Wxh *mat64.Dense //
		Why *mat64.Dense //
		// This is the last known hidden vector that represents the memory of the RNN
		// This is used only for training
		Hprev  []float64
		Bh     []float64 // This is the biais
		By     []float64 // This is the biais
		Config neuralNetConfig
	}
	rnn.Lock()
	err := enc.Encode(bkp{
		rnn.whh,
		rnn.wxh,
		rnn.why,
		rnn.hprev,
		rnn.bh,
		rnn.by,
		rnn.config,
	})
	rnn.Unlock()
	return output.Bytes(), err
}

// NewRNN creates a new RNN with input size of x, outputsize of y and hidden dimension of h
// The hidden state h is initialized with the zero vectornn.
//func newRNN(x, y, h int) *RNN {
func NewRNN(inputNeurons, outputNeurons int) *RNN {
	var conf neuralNetConfig
	if inputNeurons == 0 || outputNeurons == 0 {
		err := envconfig.Usage("RNN", &conf)
		if err != nil {
			log.Fatal(err)
		}
		return nil
	}
	err := envconfig.Process("RNN", &conf)
	if err != nil {
		log.Fatal(err)
	}

	//func NewRNN(config NeuralNetConfig) *RNN {
	var rnn RNN
	conf.InputNeurons = inputNeurons
	conf.OutputNeurons = outputNeurons
	rnn.config = conf
	// Initialize biases/weights.

	rnn.wxh = mat64.NewDense(conf.HiddenNeurons, conf.InputNeurons, nil)
	rnn.whh = mat64.NewDense(conf.HiddenNeurons, conf.HiddenNeurons, nil)
	rnn.why = mat64.NewDense(conf.OutputNeurons, conf.HiddenNeurons, nil)
	rnn.bh = make([]float64, conf.HiddenNeurons)
	rnn.by = make([]float64, conf.OutputNeurons)
	wHiddenRaw := rnn.wxh.RawMatrix().Data
	wHiddenHiddenRaw := rnn.whh.RawMatrix().Data
	wOutRaw := rnn.why.RawMatrix().Data

	for _, param := range [][]float64{
		wHiddenRaw,
		wHiddenHiddenRaw,
		wOutRaw,
	} {
		for i := range param {
			randSource := rand.NewSource(time.Now().UnixNano())
			randGen := rand.New(randSource)
			param[i] = randGen.NormFloat64() * conf.RandomFactor
		}
	}

	// initialise the hidden vector to zero
	rnn.hprev = make([]float64, conf.HiddenNeurons)

	return &rnn
}

// RNNs have a deceptively simple API:
// They accept an input vector x and give you an output vector y.
// However, crucially this output vector’s contents are influenced
// not only by the input you just fed in,
// but also on the entire history of inputs you’ve fed in in the past.
// Written as a class, the RNN’s API consists of a single step function:
func (rnn *RNN) step(x, hprev []float64) (y, h []float64) {
	h = tanh(
		add(
			dot(rnn.wxh, x),
			dot(rnn.whh, hprev),
			rnn.bh,
		))
	y = add(
		dot(rnn.why, h),
		rnn.by)
	return
}

// forwardPass takes a mat64rix of inputs and returns
// the corresponding outputs mat64rix
// and a mat64rix of the hidden states that will be used
// for the backpropagation
func (rnn *RNN) forwardPass(xs [][]float64, hprev []float64) (ys, hs [][]float64) {
	inputSize := len(xs)
	hp := make([]float64, len(hprev))
	copy(hp, hprev)
	// un-normalized log probabilities for next chars
	ys = make([][]float64, inputSize)
	hs = make([][]float64, inputSize)
	for t := 0; t < inputSize; t++ {
		// Initialization of the arrays
		ys[t] = make([]float64, rnn.config.OutputNeurons)
		hs[t] = make([]float64, rnn.config.HiddenNeurons)
		ys[t], hs[t] = rnn.step(xs[t], hp)
		hp = hs[t]
	}
	return
}

// Do a backpropagation of the RNNs and returns the derivates
// xs is the input mat64rix
// ts is the target mat64rices
// ps is the normalized log probability
// hs is a mat64rix of hidden vector
func (rnn *RNN) backPropagation(xs, ps, hs, ts [][]float64) (dwxh, dwhh, dwhy *mat64.Dense, dbh, dby []float64) {
	inputSize := len(xs)
	dwxh = mat64.NewDense(rnn.config.HiddenNeurons, rnn.config.InputNeurons, nil)
	dwhh = mat64.NewDense(rnn.config.HiddenNeurons, rnn.config.HiddenNeurons, nil)
	dwhy = mat64.NewDense(rnn.config.OutputNeurons, rnn.config.HiddenNeurons, nil)
	dhnext := make([]float64, rnn.config.OutputNeurons)
	dbh = make([]float64, rnn.config.HiddenNeurons)
	dby = make([]float64, rnn.config.OutputNeurons)
	dhraw := make([]float64, len(hs[0]))

	for t := inputSize - 1; t >= 0; t-- {
		dy := make([]float64, rnn.config.OutputNeurons)
		for i := range ps[t] {
			dy[i] = ps[t][i] - ts[t][i]
		}
		dwhy.Add(dwhy,
			dotVec(dy, hs[t]),
		)
		dby = add(dby, dy)

		dh := add(
			dot(rnn.why.T(), dy),
			dhnext,
		)

		for i := range hs[t] {
			dhraw[i] = (1 - hs[t][i]*hs[t][i]) * dh[i]
		}

		dbh = add(dbh, dhraw)
		dwxh.Add(dwxh, dotVec(dhraw, xs[t]))
		if t == 0 {
			dwhh.Add(dwhh, dotVec(dhraw, hs[inputSize-1]))
		} else {
			dwhh.Add(dwhh, dotVec(dhraw, hs[t-1]))
		}
		dhnext = dot(rnn.whh.T(), dhraw)
	}

	return
}

// TrainingSet represents an input mat64rix and the expected
// result when passed through a rnn
type TrainingSet struct {
	Inputs  [][]float64
	Targets [][]float64
}

// CopyOf the trainingset passed as parameter
func CopyOf(tset TrainingSet) TrainingSet {
	xs := make([][]float64, len(tset.Inputs))
	copy(xs, tset.Inputs)
	ts := make([][]float64, len(tset.Targets))
	copy(ts, tset.Targets)
	return TrainingSet{
		xs,
		ts,
	}
}

// Train the network.
// The train mechanisme is launched in a seperate go-routine
// it is waiting for an input to be sent in the feeding channel
func (rnn *RNN) Train() (feed chan TrainingSet, info chan float64) {
	feed = make(chan TrainingSet, 1)
	info = make(chan float64, 1)

	adagrad := newAdagrad(rnn.config)
	go func(feed chan TrainingSet, info chan float64) {
		// When we have new data
		for tset := range feed {
			// Forward pass
			xs := make([][]float64, len(tset.Inputs))
			copy(xs, tset.Inputs)
			ts := make([][]float64, len(tset.Targets))
			copy(ts, tset.Targets)
			rnn.Lock()
			ys, hs := rnn.forwardPass(xs, rnn.hprev)
			// Save the last state for future training
			copy(rnn.hprev, hs[len(hs)-1])
			//rnn.hprev = hs[len(hs)-1]
			rnn.Unlock()
			ps := normalizeByRow(ys)
			// Loss evaluation
			loss := float64(0)
			for t := 0; t < len(ps); t++ {
				l := float64(0)
				for i := 0; i < len(ps[t]); i++ {
					l += ps[t][i] * ts[t][i]
				}
				loss -= math.Log(l)
			}
			// Send info on a non blocking channel
			select {
			case info <- loss:
			default:
			}

			// Backpass
			rnn.Lock()
			dwxh, dwhh, dwhy, dbh, dby := rnn.backPropagation(xs, ps, hs, ts)
			rnn.Unlock()
			// Clip to mitigate exploding gradients
			for _, param := range [][]float64{
				dwxh.RawMatrix().Data,
				dwhh.RawMatrix().Data,
				dwhy.RawMatrix().Data,
				dby,
				dbh,
			} {
				func(param []float64) {
					for i := range param {
						if param[i] > 5 {
							param[i] = 5
						}
						if param[i] < -5 {
							param[i] = -5
						}
					}
				}(param)
			}
			// Adaptation
			rnn.Lock()
			adagrad.apply(rnn, dwxh, dwhh, dwhy, dbh, dby)
			rnn.Unlock()
		}
	}(feed, info)
	return feed, info
}

// Predict n element of  output that corresponds to the input xs
// At every iteration, the output is processed by the adapt function
func (rnn *RNN) Predict(xs [][]float64, n int, adapt func([]float64) []float64) [][]float64 {
	ys := make([][]float64, n+len(xs))
	h := make([]float64, len(rnn.hprev))
	y := make([]float64, rnn.config.OutputNeurons)
	for i := 0; i < n+len(xs); i++ {
		x := make([]float64, rnn.config.InputNeurons)
		if i < len(xs) {
			copy(x, xs[i])
		} else {
			copy(x, ys[i-1])
		}

		yr, hr := rnn.step(x, h)
		copy(y, yr)
		copy(h, hr)
		expY := exp(y)
		p := div(expY, sum(expY))
		ys[i] = p
		if i < len(xs) {
			for j := 0; j < len(xs[i]); j++ {
				if xs[i][j] == float64(1) {
					ys[i][j] = float64(1)
				}
			}
		} else {
			ys[i] = adapt(p)
		}
	}
	res := make([][]float64, n)
	copy(res, ys[len(xs):])
	return res
}

// Sample the rnn
func (rnn *RNN) Sample(xs [][]float64, n int, choose func([]float64) int) []int {
	res := make([]int, n+len(xs))
	h := make([]float64, len(rnn.hprev))
	//copy(h, rnn.hprev)
	y := make([]float64, rnn.config.OutputNeurons)
	for i := 0; i < n; i++ {
		x := make([]float64, rnn.config.InputNeurons)
		if i < len(xs) {
			copy(x, xs[i])
		} else {
			x[res[i-1]] = 1
		}

		yr, hr := rnn.step(x, h)
		copy(y, yr)
		copy(h, hr)
		expY := exp(y)
		p := div(expY, sum(expY))
		if i < len(xs) {
			for j := 0; j < len(xs[i]); j++ {
				if xs[i][j] == float64(1) {
					res[i] = j
				}
			}
		} else {
			res[i] = choose(p)
		}
	}

	return res
}
<|endoftext|>


<commit_before>package rnn

import (
	"math"
	"math/rand"
	"sync"
	"time"

	"gonum.org/v1/gonum/mat"
)

// RNN represents the neural network
// This RNNs parameters are the three matrices whh, wxh, why.
// hprev is the last known hidden vector, which is actually the memory of the RNN
// bh, and by are the biais vectors respectivly for the hidden layer and the output layer
type RNN struct {
	sync.Mutex
	whh *mat.Dense // size is hiddenDimension * hiddenDimension
	wxh *mat.Dense //
	why *mat.Dense //
	// This is the last known hidden vector that represents the memory of the RNN
	// This is used only for training
	hprev  []float64
	bh     []float64 // This is the biais
	by     []float64 // This is the biais
	config neuralNetConfig
}

// NewRNN creates a new RNN with input size of x, outputsize of y and hidden dimension of h
// The hidden state h is initialized with the zero vector.
//func newRNN(x, y, h int) *RNN {
func NewRNN(inputNeurons, outputNeurons int) *RNN {
	//func NewRNN(config NeuralNetConfig) *RNN {
	var rnn RNN
	conf.inputNeurons = inputNeurons
	conf.outputNeurons = outputNeurons
	rnn.config = conf
	// Initialize biases/weights.

	rnn.wxh = mat.NewDense(conf.HiddenNeurons, conf.inputNeurons, nil)
	rnn.whh = mat.NewDense(conf.HiddenNeurons, conf.HiddenNeurons, nil)
	rnn.why = mat.NewDense(conf.outputNeurons, conf.HiddenNeurons, nil)
	rnn.bh = make([]float64, conf.HiddenNeurons)
	rnn.by = make([]float64, conf.outputNeurons)
	wHiddenRaw := rnn.wxh.RawMatrix().Data
	wHiddenHiddenRaw := rnn.whh.RawMatrix().Data
	wOutRaw := rnn.why.RawMatrix().Data

	for _, param := range [][]float64{
		wHiddenRaw,
		wHiddenHiddenRaw,
		wOutRaw,
	} {
		for i := range param {
			randSource := rand.NewSource(time.Now().UnixNano())
			randGen := rand.New(randSource)
			param[i] = randGen.NormFloat64() * conf.RandomFactor
		}
	}

	// initialise the hidden vector to zero
	rnn.hprev = make([]float64, conf.HiddenNeurons)

	return &rnn
}

// RNNs have a deceptively simple API:
// They accept an input vector x and give you an output vector y.
// However, crucially this output vector’s contents are influenced
// not only by the input you just fed in,
// but also on the entire history of inputs you’ve fed in in the past.
// Written as a class, the RNN’s API consists of a single step function:
func (r *RNN) step(x, hprev []float64) (y, h []float64) {
	r.Lock()
	h = tanh(
		add(
			dot(r.wxh, x),
			dot(r.whh, hprev),
			r.bh,
		))
	y = add(
		dot(r.why, h),
		r.by)
	r.Unlock()
	return
}

// forwardPass takes a matrix of inputs and returns
// the corresponding outputs matrix
// and a matrix of the hidden states that will be used
// for the backpropagation
func (r *RNN) forwardPass(xs [][]float64, hprev []float64) (ys, hs [][]float64) {
	inputSize := len(xs)
	hp := make([]float64, len(hprev))
	r.Lock()
	copy(hp, hprev)
	r.Unlock()
	// un-normalized log probabilities for next chars
	ys = make([][]float64, inputSize)
	hs = make([][]float64, inputSize)
	for t := 0; t < inputSize; t++ {
		// Initialization of the arrays
		ys[t] = make([]float64, r.config.outputNeurons)
		hs[t] = make([]float64, r.config.HiddenNeurons)
		ys[t], hs[t] = r.step(xs[t], hp)
		hp = hs[t]
	}
	return
}

// Do a backpropagation of the RNNs and returns the derivates
// xs is the input matrix
// ts is the target matrices
// ps is the normalized log probability
// hs is a matrix of hidden vector
func (r *RNN) backPropagation(xs, ps, hs, ts [][]float64) (dwxh, dwhh, dwhy *mat.Dense, dbh, dby []float64) {
	inputSize := len(xs)
	dwxh = mat.NewDense(r.config.HiddenNeurons, r.config.inputNeurons, nil)
	dwhh = mat.NewDense(r.config.HiddenNeurons, r.config.HiddenNeurons, nil)
	dwhy = mat.NewDense(r.config.outputNeurons, r.config.HiddenNeurons, nil)
	dhnext := make([]float64, r.config.outputNeurons)
	dbh = make([]float64, r.config.HiddenNeurons)
	dby = make([]float64, r.config.outputNeurons)
	dhraw := make([]float64, len(hs[0]))

	for t := inputSize - 1; t >= 0; t-- {
		dy := make([]float64, r.config.outputNeurons)
		for i := range ps[t] {
			dy[i] = ps[t][i] - ts[t][i]
		}
		dwhy.Add(dwhy,
			dotVec(dy, hs[t]),
		)
		dby = add(dby, dy)

		dh := add(
			dot(r.why.T(), dy),
			dhnext,
		)

		for i := range hs[t] {
			dhraw[i] = (1 - hs[t][i]*hs[t][i]) * dh[i]
		}

		dbh = add(dbh, dhraw)
		dwxh.Add(dwxh, dotVec(dhraw, xs[t]))
		if t == 0 {
			dwhh.Add(dwhh, dotVec(dhraw, hs[inputSize-1]))
		} else {
			dwhh.Add(dwhh, dotVec(dhraw, hs[t-1]))
		}
		dhnext = dot(r.whh.T(), dhraw)
	}

	return
}

// TrainingSet represents an input matrix and the expected
// result when passed through a rnn
type TrainingSet struct {
	Inputs  [][]float64
	Targets [][]float64
}

// Train the network.
// The train mechanisme is launched in a seperate go-routine
// it is waiting for an input to be sent in the feeding channel
func (r *RNN) Train() (feed chan TrainingSet, info chan float64) {
	feed = make(chan TrainingSet, 1)
	info = make(chan float64, 1)

	adagrad := newAdagrad(r.config)
	go func(feed chan TrainingSet, info chan float64) {
		// When we have new data
		for tset := range feed {
			// Forward pass
			xs := make([][]float64, len(tset.Inputs))
			copy(xs, tset.Inputs)
			ts := make([][]float64, len(tset.Targets))
			copy(ts, tset.Targets)
			ys, hs := r.forwardPass(xs, r.hprev)
			// Save the last state for future training
			r.Lock()
			r.hprev = hs[len(hs)-1]
			r.Unlock()
			ps := normalizeByRow(ys)
			// Loss evaluation
			loss := float64(0)
			for t := 0; t < len(ps); t++ {
				l := float64(0)
				for i := 0; i < len(ps[t]); i++ {
					l += ps[t][i] * ts[t][i]
				}
				loss -= math.Log(l)
			}
			// Send info on a non blocking channel
			select {
			case info <- loss:
			default:
			}

			// Backpass
			r.Lock()
			dwxh, dwhh, dwhy, dbh, dby := r.backPropagation(xs, ps, hs, ts)
			r.Unlock()
			// Clip to mitigate exploding gradients
			for _, param := range [][]float64{
				dwxh.RawMatrix().Data,
				dwhh.RawMatrix().Data,
				dwhy.RawMatrix().Data,
				dby,
				dbh,
			} {
				func(param []float64) {
					for i := range param {
						if param[i] > 5 {
							param[i] = 5
						}
						if param[i] < -5 {
							param[i] = -5
						}
					}
				}(param)
			}
			// Adaptation
			r.Lock()
			adagrad.apply(r, dwxh, dwhh, dwhy, dbh, dby)
			r.Unlock()
		}
	}(feed, info)
	return feed, info
}

// Sample the rnn
// the choose function returns the index choosen, this allows to put
// some randomness if needed
// TODO: this function needs to be rewritten
func (r *RNN) Sample(seed, n int, choose func([]float64) int) []int {
	res := make([]int, n)
	h := make([]float64, len(r.hprev))
	//copy(h, r.hprev)
	x := make([]float64, r.config.inputNeurons)
	y := make([]float64, r.config.outputNeurons)
	for i := 0; i < n; i++ {
		if i == 0 {
			x[seed] = 1
		} else {
			x[res[i-1]] = 1
		}

		y, h = r.step(x, h)
		expY := exp(y)
		p := div(expY, sum(expY))
		res[i] = choose(p)
	}

	return res
}
<commit_msg>Bug fix, in the sampling process, and in the hidden vector the input vector was not reset between each iteraction<commit_after>package rnn

import (
	"math"
	"math/rand"
	"sync"
	"time"

	"gonum.org/v1/gonum/mat"
)

// RNN represents the neural network
// This RNNs parameters are the three matrices whh, wxh, why.
// hprev is the last known hidden vector, which is actually the memory of the RNN
// bh, and by are the biais vectors respectivly for the hidden layer and the output layer
type RNN struct {
	sync.Mutex
	whh *mat.Dense // size is hiddenDimension * hiddenDimension
	wxh *mat.Dense //
	why *mat.Dense //
	// This is the last known hidden vector that represents the memory of the RNN
	// This is used only for training
	hprev  []float64
	bh     []float64 // This is the biais
	by     []float64 // This is the biais
	config neuralNetConfig
}

// NewRNN creates a new RNN with input size of x, outputsize of y and hidden dimension of h
// The hidden state h is initialized with the zero vector.
//func newRNN(x, y, h int) *RNN {
func NewRNN(inputNeurons, outputNeurons int) *RNN {
	//func NewRNN(config NeuralNetConfig) *RNN {
	var rnn RNN
	conf.inputNeurons = inputNeurons
	conf.outputNeurons = outputNeurons
	rnn.config = conf
	// Initialize biases/weights.

	rnn.wxh = mat.NewDense(conf.HiddenNeurons, conf.inputNeurons, nil)
	rnn.whh = mat.NewDense(conf.HiddenNeurons, conf.HiddenNeurons, nil)
	rnn.why = mat.NewDense(conf.outputNeurons, conf.HiddenNeurons, nil)
	rnn.bh = make([]float64, conf.HiddenNeurons)
	rnn.by = make([]float64, conf.outputNeurons)
	wHiddenRaw := rnn.wxh.RawMatrix().Data
	wHiddenHiddenRaw := rnn.whh.RawMatrix().Data
	wOutRaw := rnn.why.RawMatrix().Data

	for _, param := range [][]float64{
		wHiddenRaw,
		wHiddenHiddenRaw,
		wOutRaw,
	} {
		for i := range param {
			randSource := rand.NewSource(time.Now().UnixNano())
			randGen := rand.New(randSource)
			param[i] = randGen.NormFloat64() * conf.RandomFactor
		}
	}

	// initialise the hidden vector to zero
	rnn.hprev = make([]float64, conf.HiddenNeurons)

	return &rnn
}

// RNNs have a deceptively simple API:
// They accept an input vector x and give you an output vector y.
// However, crucially this output vector’s contents are influenced
// not only by the input you just fed in,
// but also on the entire history of inputs you’ve fed in in the past.
// Written as a class, the RNN’s API consists of a single step function:
func (r *RNN) step(x, hprev []float64) (y, h []float64) {
	h = tanh(
		add(
			dot(r.wxh, x),
			dot(r.whh, hprev),
			r.bh,
		))
	y = add(
		dot(r.why, h),
		r.by)
	return
}

// forwardPass takes a matrix of inputs and returns
// the corresponding outputs matrix
// and a matrix of the hidden states that will be used
// for the backpropagation
func (r *RNN) forwardPass(xs [][]float64, hprev []float64) (ys, hs [][]float64) {
	inputSize := len(xs)
	hp := make([]float64, len(hprev))
	copy(hp, hprev)
	// un-normalized log probabilities for next chars
	ys = make([][]float64, inputSize)
	hs = make([][]float64, inputSize)
	for t := 0; t < inputSize; t++ {
		// Initialization of the arrays
		ys[t] = make([]float64, r.config.outputNeurons)
		hs[t] = make([]float64, r.config.HiddenNeurons)
		ys[t], hs[t] = r.step(xs[t], hp)
		hp = hs[t]
	}
	return
}

// Do a backpropagation of the RNNs and returns the derivates
// xs is the input matrix
// ts is the target matrices
// ps is the normalized log probability
// hs is a matrix of hidden vector
func (r *RNN) backPropagation(xs, ps, hs, ts [][]float64) (dwxh, dwhh, dwhy *mat.Dense, dbh, dby []float64) {
	inputSize := len(xs)
	dwxh = mat.NewDense(r.config.HiddenNeurons, r.config.inputNeurons, nil)
	dwhh = mat.NewDense(r.config.HiddenNeurons, r.config.HiddenNeurons, nil)
	dwhy = mat.NewDense(r.config.outputNeurons, r.config.HiddenNeurons, nil)
	dhnext := make([]float64, r.config.outputNeurons)
	dbh = make([]float64, r.config.HiddenNeurons)
	dby = make([]float64, r.config.outputNeurons)
	dhraw := make([]float64, len(hs[0]))

	for t := inputSize - 1; t >= 0; t-- {
		dy := make([]float64, r.config.outputNeurons)
		for i := range ps[t] {
			dy[i] = ps[t][i] - ts[t][i]
		}
		dwhy.Add(dwhy,
			dotVec(dy, hs[t]),
		)
		dby = add(dby, dy)

		dh := add(
			dot(r.why.T(), dy),
			dhnext,
		)

		for i := range hs[t] {
			dhraw[i] = (1 - hs[t][i]*hs[t][i]) * dh[i]
		}

		dbh = add(dbh, dhraw)
		dwxh.Add(dwxh, dotVec(dhraw, xs[t]))
		if t == 0 {
			dwhh.Add(dwhh, dotVec(dhraw, hs[inputSize-1]))
		} else {
			dwhh.Add(dwhh, dotVec(dhraw, hs[t-1]))
		}
		dhnext = dot(r.whh.T(), dhraw)
	}

	return
}

// TrainingSet represents an input matrix and the expected
// result when passed through a rnn
type TrainingSet struct {
	Inputs  [][]float64
	Targets [][]float64
}

// Train the network.
// The train mechanisme is launched in a seperate go-routine
// it is waiting for an input to be sent in the feeding channel
func (r *RNN) Train() (feed chan TrainingSet, info chan float64) {
	feed = make(chan TrainingSet, 1)
	info = make(chan float64, 1)

	adagrad := newAdagrad(r.config)
	go func(feed chan TrainingSet, info chan float64) {
		// When we have new data
		for tset := range feed {
			// Forward pass
			xs := make([][]float64, len(tset.Inputs))
			copy(xs, tset.Inputs)
			ts := make([][]float64, len(tset.Targets))
			copy(ts, tset.Targets)
			r.Lock()
			ys, hs := r.forwardPass(xs, r.hprev)
			// Save the last state for future training
			copy(r.hprev, hs[len(hs)-1])
			//r.hprev = hs[len(hs)-1]
			r.Unlock()
			ps := normalizeByRow(ys)
			// Loss evaluation
			loss := float64(0)
			for t := 0; t < len(ps); t++ {
				l := float64(0)
				for i := 0; i < len(ps[t]); i++ {
					l += ps[t][i] * ts[t][i]
				}
				loss -= math.Log(l)
			}
			// Send info on a non blocking channel
			select {
			case info <- loss:
			default:
			}

			// Backpass
			r.Lock()
			dwxh, dwhh, dwhy, dbh, dby := r.backPropagation(xs, ps, hs, ts)
			r.Unlock()
			// Clip to mitigate exploding gradients
			for _, param := range [][]float64{
				dwxh.RawMatrix().Data,
				dwhh.RawMatrix().Data,
				dwhy.RawMatrix().Data,
				dby,
				dbh,
			} {
				func(param []float64) {
					for i := range param {
						if param[i] > 5 {
							param[i] = 5
						}
						if param[i] < -5 {
							param[i] = -5
						}
					}
				}(param)
			}
			// Adaptation
			r.Lock()
			adagrad.apply(r, dwxh, dwhh, dwhy, dbh, dby)
			r.Unlock()
		}
	}(feed, info)
	return feed, info
}

// Sample the rnn
// the choose function returns the index choosen, this allows to put
// some randomness if needed
// TODO: this function needs to be rewritten
func (r *RNN) Sample(seed, n int, choose func([]float64) int) []int {
	res := make([]int, n)
	h := make([]float64, len(r.hprev))
	//copy(h, r.hprev)
	y := make([]float64, r.config.outputNeurons)
	for i := 0; i < n; i++ {
		x := make([]float64, r.config.inputNeurons)
		if i == 0 {
			x[seed] = 1
		} else {
			x[res[i-1]] = 1
		}

		yr, hr := r.step(x, h)
		copy(y, yr)
		copy(h, hr)
		expY := exp(y)
		p := div(expY, sum(expY))
		res[i] = choose(p)
	}

	return res
}
<|endoftext|>


<commit_before>package dht

import (
	"bytes"
	"encoding/json"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	peer "github.com/jbenet/go-ipfs/peer"
	kb "github.com/jbenet/go-ipfs/routing/kbucket"
	u "github.com/jbenet/go-ipfs/util"
)

// This file implements the Routing interface for the IpfsDHT struct.

// Basic Put/Get

// PutValue adds value corresponding to given Key.
// This is the top level "Store" operation of the DHT
func (dht *IpfsDHT) PutValue(ctx context.Context, key u.Key, value []byte) error {
	peers := []*peer.Peer{}

	// get the peers we need to announce to
	for _, route := range dht.routingTables {
		npeers := route.NearestPeers(kb.ConvertKey(key), KValue)
		peers = append(peers, npeers...)
	}

	query := newQuery(key, func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {
		u.DOut("[%s] PutValue qry part %v\n", dht.self.ID.Pretty(), p.ID.Pretty())
		err := dht.putValueToNetwork(ctx, p, string(key), value)
		if err != nil {
			return nil, err
		}
		return &dhtQueryResult{success: true}, nil
	})

	_, err := query.Run(ctx, peers)
	u.DOut("[%s] PutValue %v %v\n", dht.self.ID.Pretty(), key, value)
	return err
}

// GetValue searches for the value corresponding to given Key.
// If the search does not succeed, a multiaddr string of a closer peer is
// returned along with util.ErrSearchIncomplete
func (dht *IpfsDHT) GetValue(ctx context.Context, key u.Key) ([]byte, error) {
	ll := startNewRPC("GET")
	defer ll.EndAndPrint()

	// If we have it local, dont bother doing an RPC!
	// NOTE: this might not be what we want to do...
	val, err := dht.getLocal(key)
	if err == nil {
		ll.Success = true
		return val, nil
	}

	// get closest peers in the routing tables
	routeLevel := 0
	closest := dht.routingTables[routeLevel].NearestPeers(kb.ConvertKey(key), PoolSize)
	if closest == nil || len(closest) == 0 {
		return nil, nil
	}

	// setup the Query
	query := newQuery(key, func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {

		val, peers, err := dht.getValueOrPeers(ctx, p, key, routeLevel)
		if err != nil {
			return nil, err
		}

		res := &dhtQueryResult{value: val, closerPeers: peers}
		if val != nil {
			res.success = true
		}

		return res, nil
	})

	// run it!
	result, err := query.Run(ctx, closest)
	if err != nil {
		return nil, err
	}

	u.DOut("[%s] GetValue %v %v\n", dht.self.ID.Pretty(), key, result.value)
	if result.value == nil {
		return nil, u.ErrNotFound
	}

	return result.value, nil
}

// Value provider layer of indirection.
// This is what DSHTs (Coral and MainlineDHT) do to store large values in a DHT.

// Provide makes this node announce that it can provide a value for the given key
func (dht *IpfsDHT) Provide(ctx context.Context, key u.Key) error {

	dht.providers.AddProvider(key, dht.self)
	peers := dht.routingTables[0].NearestPeers(kb.ConvertKey(key), PoolSize)
	if len(peers) == 0 {
		return nil
	}

	//TODO FIX: this doesn't work! it needs to be sent to the actual nearest peers.
	// `peers` are the closest peers we have, not the ones that should get the value.
	for _, p := range peers {
		err := dht.putProvider(ctx, p, string(key))
		if err != nil {
			return err
		}
	}
	return nil
}

// NB: not actually async. Used to keep the interface consistent while the
// actual async method, FindProvidersAsync2 is under construction
func (dht *IpfsDHT) FindProvidersAsync(ctx context.Context, key u.Key, count int) <-chan *peer.Peer {
	ch := make(chan *peer.Peer)
	providers, err := dht.FindProviders(ctx, key)
	if err != nil {
		close(ch)
		return ch
	}
	go func() {
		defer close(ch)
		for _, p := range providers {
			ch <- p
		}
	}()
	return ch
}

// FIXME: there's a bug here!
func (dht *IpfsDHT) FindProvidersAsync2(ctx context.Context, key u.Key, count int) <-chan *peer.Peer {
	peerOut := make(chan *peer.Peer, count)
	go func() {
		ps := newPeerSet()
		provs := dht.providers.GetProviders(key)
		for _, p := range provs {
			count--
			// NOTE: assuming that this list of peers is unique
			ps.Add(p)
			peerOut <- p
			if count <= 0 {
				return
			}
		}

		peers := dht.routingTables[0].NearestPeers(kb.ConvertKey(key), AlphaValue)
		for _, pp := range peers {
			ppp := pp
			go func() {
				pmes, err := dht.findProvidersSingle(ctx, ppp, key, 0)
				if err != nil {
					u.PErr("%v\n", err)
					return
				}
				dht.addPeerListAsync(key, pmes.GetProviderPeers(), ps, count, peerOut)
			}()
		}

	}()
	return peerOut
}

//TODO: this function could also be done asynchronously
func (dht *IpfsDHT) addPeerListAsync(k u.Key, peers []*Message_Peer, ps *peerSet, count int, out chan *peer.Peer) {
	for _, pbp := range peers {

		// construct new peer
		p, err := dht.ensureConnectedToPeer(pbp)
		if err != nil {
			continue
		}

		dht.providers.AddProvider(k, p)
		if ps.AddIfSmallerThan(p, count) {
			out <- p
		} else if ps.Size() >= count {
			return
		}
	}
}

// FindProviders searches for peers who can provide the value for given key.
func (dht *IpfsDHT) FindProviders(ctx context.Context, key u.Key) ([]*peer.Peer, error) {
	ll := startNewRPC("FindProviders")
	ll.EndAndPrint()

	// get closest peer
	u.DOut("Find providers for: '%s'\n", key)
	p := dht.routingTables[0].NearestPeer(kb.ConvertKey(key))
	if p == nil {
		return nil, nil
	}

	for level := 0; level < len(dht.routingTables); {

		// attempt retrieving providers
		pmes, err := dht.findProvidersSingle(ctx, p, key, level)
		if err != nil {
			return nil, err
		}

		// handle providers
		provs := pmes.GetProviderPeers()
		if provs != nil {
			u.DOut("Got providers back from findProviders call!\n")
			return dht.addProviders(key, provs), nil
		}

		u.DOut("Didnt get providers, just closer peers.\n")
		closer := pmes.GetCloserPeers()
		if len(closer) == 0 {
			level++
			continue
		}

		np, err := dht.peerFromInfo(closer[0])
		if err != nil {
			u.DOut("no peerFromInfo")
			level++
			continue
		}
		p = np
	}
	return nil, u.ErrNotFound
}

// Find specific Peer

// FindPeer searches for a peer with given ID.
func (dht *IpfsDHT) FindPeer(ctx context.Context, id peer.ID) (*peer.Peer, error) {

	// Check if were already connected to them
	p, _ := dht.Find(id)
	if p != nil {
		return p, nil
	}

	// @whyrusleeping why is this here? doesn't the dht.Find above cover it?
	routeLevel := 0
	p = dht.routingTables[routeLevel].NearestPeer(kb.ConvertPeerID(id))
	if p == nil {
		return nil, nil
	}
	if p.ID.Equal(id) {
		return p, nil
	}

	for routeLevel < len(dht.routingTables) {
		pmes, err := dht.findPeerSingle(ctx, p, id, routeLevel)
		plist := pmes.GetCloserPeers()
		if plist == nil || len(plist) == 0 {
			routeLevel++
			continue
		}
		found := plist[0]

		nxtPeer, err := dht.ensureConnectedToPeer(found)
		if err != nil {
			routeLevel++
			continue
		}

		if nxtPeer.ID.Equal(id) {
			return nxtPeer, nil
		}

		p = nxtPeer
	}
	return nil, u.ErrNotFound
}

func (dht *IpfsDHT) findPeerMultiple(ctx context.Context, id peer.ID) (*peer.Peer, error) {

	// Check if were already connected to them
	p, _ := dht.Find(id)
	if p != nil {
		return p, nil
	}

	// get the peers we need to announce to
	routeLevel := 0
	peers := dht.routingTables[routeLevel].NearestPeers(kb.ConvertPeerID(id), AlphaValue)
	if len(peers) == 0 {
		return nil, nil
	}

	// setup query function
	query := newQuery(u.Key(id), func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {
		pmes, err := dht.findPeerSingle(ctx, p, id, routeLevel)
		if err != nil {
			u.DErr("getPeer error: %v\n", err)
			return nil, err
		}

		plist := pmes.GetCloserPeers()
		if len(plist) == 0 {
			routeLevel++
		}

		nxtprs := make([]*peer.Peer, len(plist))
		for i, fp := range plist {
			nxtp, err := dht.peerFromInfo(fp)
			if err != nil {
				u.DErr("findPeer error: %v\n", err)
				continue
			}

			if nxtp.ID.Equal(id) {
				return &dhtQueryResult{peer: nxtp, success: true}, nil
			}

			nxtprs[i] = nxtp
		}

		return &dhtQueryResult{closerPeers: nxtprs}, nil
	})

	result, err := query.Run(ctx, peers)
	if err != nil {
		return nil, err
	}

	if result.peer == nil {
		return nil, u.ErrNotFound
	}
	return result.peer, nil
}

// Ping a peer, log the time it took
func (dht *IpfsDHT) Ping(ctx context.Context, p *peer.Peer) error {
	// Thoughts: maybe this should accept an ID and do a peer lookup?
	u.DOut("[%s] ping %s start\n", dht.self.ID.Pretty(), p.ID.Pretty())

	pmes := newMessage(Message_PING, "", 0)
	_, err := dht.sendRequest(ctx, p, pmes)
	u.DOut("[%s] ping %s end (err = %s)\n", dht.self.ID.Pretty(), p.ID.Pretty(), err)
	return err
}

func (dht *IpfsDHT) getDiagnostic(ctx context.Context) ([]*diagInfo, error) {

	u.DOut("Begin Diagnostic")
	peers := dht.routingTables[0].NearestPeers(kb.ConvertPeerID(dht.self.ID), 10)
	var out []*diagInfo

	query := newQuery(dht.self.Key(), func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {
		pmes := newMessage(Message_DIAGNOSTIC, "", 0)
		rpmes, err := dht.sendRequest(ctx, p, pmes)
		if err != nil {
			return nil, err
		}

		dec := json.NewDecoder(bytes.NewBuffer(rpmes.GetValue()))
		for {
			di := new(diagInfo)
			err := dec.Decode(di)
			if err != nil {
				break
			}

			out = append(out, di)
		}
		return &dhtQueryResult{success: true}, nil
	})

	_, err := query.Run(ctx, peers)
	return out, err
}
<commit_msg>fixes to make interface more usable<commit_after>package dht

import (
	"bytes"
	"encoding/json"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	peer "github.com/jbenet/go-ipfs/peer"
	kb "github.com/jbenet/go-ipfs/routing/kbucket"
	u "github.com/jbenet/go-ipfs/util"
)

// This file implements the Routing interface for the IpfsDHT struct.

// Basic Put/Get

// PutValue adds value corresponding to given Key.
// This is the top level "Store" operation of the DHT
func (dht *IpfsDHT) PutValue(ctx context.Context, key u.Key, value []byte) error {
	peers := []*peer.Peer{}

	// get the peers we need to announce to
	for _, route := range dht.routingTables {
		npeers := route.NearestPeers(kb.ConvertKey(key), KValue)
		peers = append(peers, npeers...)
	}

	query := newQuery(key, func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {
		u.DOut("[%s] PutValue qry part %v\n", dht.self.ID.Pretty(), p.ID.Pretty())
		err := dht.putValueToNetwork(ctx, p, string(key), value)
		if err != nil {
			return nil, err
		}
		return &dhtQueryResult{success: true}, nil
	})

	_, err := query.Run(ctx, peers)
	u.DOut("[%s] PutValue %v %v\n", dht.self.ID.Pretty(), key, value)
	return err
}

// GetValue searches for the value corresponding to given Key.
// If the search does not succeed, a multiaddr string of a closer peer is
// returned along with util.ErrSearchIncomplete
func (dht *IpfsDHT) GetValue(ctx context.Context, key u.Key) ([]byte, error) {
	ll := startNewRPC("GET")
	defer ll.EndAndPrint()

	// If we have it local, dont bother doing an RPC!
	// NOTE: this might not be what we want to do...
	val, err := dht.getLocal(key)
	if err == nil {
		ll.Success = true
		return val, nil
	}

	// get closest peers in the routing tables
	routeLevel := 0
	closest := dht.routingTables[routeLevel].NearestPeers(kb.ConvertKey(key), PoolSize)
	if closest == nil || len(closest) == 0 {
		return nil, nil
	}

	// setup the Query
	query := newQuery(key, func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {

		val, peers, err := dht.getValueOrPeers(ctx, p, key, routeLevel)
		if err != nil {
			return nil, err
		}

		res := &dhtQueryResult{value: val, closerPeers: peers}
		if val != nil {
			res.success = true
		}

		return res, nil
	})

	// run it!
	result, err := query.Run(ctx, closest)
	if err != nil {
		return nil, err
	}

	u.DOut("[%s] GetValue %v %v\n", dht.self.ID.Pretty(), key, result.value)
	if result.value == nil {
		return nil, u.ErrNotFound
	}

	return result.value, nil
}

// Value provider layer of indirection.
// This is what DSHTs (Coral and MainlineDHT) do to store large values in a DHT.

// Provide makes this node announce that it can provide a value for the given key
func (dht *IpfsDHT) Provide(ctx context.Context, key u.Key) error {

	dht.providers.AddProvider(key, dht.self)
	peers := dht.routingTables[0].NearestPeers(kb.ConvertKey(key), PoolSize)
	if len(peers) == 0 {
		// Early out for no targets
		return nil
	}

	//TODO FIX: this doesn't work! it needs to be sent to the actual nearest peers.
	// `peers` are the closest peers we have, not the ones that should get the value.
	for _, p := range peers {
		err := dht.putProvider(ctx, p, string(key))
		if err != nil {
			return err
		}
	}
	return nil
}

// NB: not actually async. Used to keep the interface consistent while the
// actual async method, FindProvidersAsync2 is under construction
func (dht *IpfsDHT) FindProvidersAsync(ctx context.Context, key u.Key, count int) <-chan *peer.Peer {
	ch := make(chan *peer.Peer)
	providers, err := dht.FindProviders(ctx, key)
	if err != nil {
		close(ch)
		return ch
	}
	go func() {
		defer close(ch)
		for _, p := range providers {
			ch <- p
		}
	}()
	return ch
}

// FIXME: there's a bug here!
func (dht *IpfsDHT) FindProvidersAsync2(ctx context.Context, key u.Key, count int) <-chan *peer.Peer {
	peerOut := make(chan *peer.Peer, count)
	go func() {
		ps := newPeerSet()
		provs := dht.providers.GetProviders(key)
		for _, p := range provs {
			count--
			// NOTE: assuming that this list of peers is unique
			ps.Add(p)
			peerOut <- p
			if count <= 0 {
				return
			}
		}

		peers := dht.routingTables[0].NearestPeers(kb.ConvertKey(key), AlphaValue)
		for _, pp := range peers {
			ppp := pp
			go func() {
				pmes, err := dht.findProvidersSingle(ctx, ppp, key, 0)
				if err != nil {
					u.PErr("%v\n", err)
					return
				}
				dht.addPeerListAsync(key, pmes.GetProviderPeers(), ps, count, peerOut)
			}()
		}

	}()
	return peerOut
}

//TODO: this function could also be done asynchronously
func (dht *IpfsDHT) addPeerListAsync(k u.Key, peers []*Message_Peer, ps *peerSet, count int, out chan *peer.Peer) {
	for _, pbp := range peers {

		// construct new peer
		p, err := dht.ensureConnectedToPeer(pbp)
		if err != nil {
			continue
		}

		dht.providers.AddProvider(k, p)
		if ps.AddIfSmallerThan(p, count) {
			out <- p
		} else if ps.Size() >= count {
			return
		}
	}
}

// FindProviders searches for peers who can provide the value for given key.
func (dht *IpfsDHT) FindProviders(ctx context.Context, key u.Key) ([]*peer.Peer, error) {
	ll := startNewRPC("FindProviders")
	ll.EndAndPrint()

	// get closest peer
	u.DOut("Find providers for: '%s'\n", key)
	p := dht.routingTables[0].NearestPeer(kb.ConvertKey(key))
	if p == nil {
		return nil, nil
	}

	for level := 0; level < len(dht.routingTables); {

		// attempt retrieving providers
		pmes, err := dht.findProvidersSingle(ctx, p, key, level)
		if err != nil {
			return nil, err
		}

		// handle providers
		provs := pmes.GetProviderPeers()
		if provs != nil {
			u.DOut("Got providers back from findProviders call!\n")
			return dht.addProviders(key, provs), nil
		}

		u.DOut("Didnt get providers, just closer peers.\n")
		closer := pmes.GetCloserPeers()
		if len(closer) == 0 {
			level++
			continue
		}

		np, err := dht.peerFromInfo(closer[0])
		if err != nil {
			u.DOut("no peerFromInfo")
			level++
			continue
		}
		p = np
	}
	return nil, u.ErrNotFound
}

// Find specific Peer

// FindPeer searches for a peer with given ID.
func (dht *IpfsDHT) FindPeer(ctx context.Context, id peer.ID) (*peer.Peer, error) {

	// Check if were already connected to them
	p, _ := dht.Find(id)
	if p != nil {
		return p, nil
	}

	// @whyrusleeping why is this here? doesn't the dht.Find above cover it?
	routeLevel := 0
	p = dht.routingTables[routeLevel].NearestPeer(kb.ConvertPeerID(id))
	if p == nil {
		return nil, nil
	}
	if p.ID.Equal(id) {
		return p, nil
	}

	for routeLevel < len(dht.routingTables) {
		pmes, err := dht.findPeerSingle(ctx, p, id, routeLevel)
		plist := pmes.GetCloserPeers()
		if plist == nil || len(plist) == 0 {
			routeLevel++
			continue
		}
		found := plist[0]

		nxtPeer, err := dht.ensureConnectedToPeer(found)
		if err != nil {
			routeLevel++
			continue
		}

		if nxtPeer.ID.Equal(id) {
			return nxtPeer, nil
		}

		p = nxtPeer
	}
	return nil, u.ErrNotFound
}

func (dht *IpfsDHT) findPeerMultiple(ctx context.Context, id peer.ID) (*peer.Peer, error) {

	// Check if were already connected to them
	p, _ := dht.Find(id)
	if p != nil {
		return p, nil
	}

	// get the peers we need to announce to
	routeLevel := 0
	peers := dht.routingTables[routeLevel].NearestPeers(kb.ConvertPeerID(id), AlphaValue)
	if len(peers) == 0 {
		return nil, nil
	}

	// setup query function
	query := newQuery(u.Key(id), func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {
		pmes, err := dht.findPeerSingle(ctx, p, id, routeLevel)
		if err != nil {
			u.DErr("getPeer error: %v\n", err)
			return nil, err
		}

		plist := pmes.GetCloserPeers()
		if len(plist) == 0 {
			routeLevel++
		}

		nxtprs := make([]*peer.Peer, len(plist))
		for i, fp := range plist {
			nxtp, err := dht.peerFromInfo(fp)
			if err != nil {
				u.DErr("findPeer error: %v\n", err)
				continue
			}

			if nxtp.ID.Equal(id) {
				return &dhtQueryResult{peer: nxtp, success: true}, nil
			}

			nxtprs[i] = nxtp
		}

		return &dhtQueryResult{closerPeers: nxtprs}, nil
	})

	result, err := query.Run(ctx, peers)
	if err != nil {
		return nil, err
	}

	if result.peer == nil {
		return nil, u.ErrNotFound
	}
	return result.peer, nil
}

// Ping a peer, log the time it took
func (dht *IpfsDHT) Ping(ctx context.Context, p *peer.Peer) error {
	// Thoughts: maybe this should accept an ID and do a peer lookup?
	u.DOut("[%s] ping %s start\n", dht.self.ID.Pretty(), p.ID.Pretty())

	pmes := newMessage(Message_PING, "", 0)
	_, err := dht.sendRequest(ctx, p, pmes)
	u.DOut("[%s] ping %s end (err = %s)\n", dht.self.ID.Pretty(), p.ID.Pretty(), err)
	return err
}

func (dht *IpfsDHT) getDiagnostic(ctx context.Context) ([]*diagInfo, error) {

	u.DOut("Begin Diagnostic")
	peers := dht.routingTables[0].NearestPeers(kb.ConvertPeerID(dht.self.ID), 10)
	var out []*diagInfo

	query := newQuery(dht.self.Key(), func(ctx context.Context, p *peer.Peer) (*dhtQueryResult, error) {
		pmes := newMessage(Message_DIAGNOSTIC, "", 0)
		rpmes, err := dht.sendRequest(ctx, p, pmes)
		if err != nil {
			return nil, err
		}

		dec := json.NewDecoder(bytes.NewBuffer(rpmes.GetValue()))
		for {
			di := new(diagInfo)
			err := dec.Decode(di)
			if err != nil {
				break
			}

			out = append(out, di)
		}
		return &dhtQueryResult{success: true}, nil
	})

	_, err := query.Run(ctx, peers)
	return out, err
}
<|endoftext|>


<commit_before>package pop

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"text/tabwriter"
	"time"

	"github.com/pkg/errors"
)

var mrx = regexp.MustCompile(`(\d+)_([^\.]+)(\.[a-z]+)?\.(up|down)\.(sql|fizz)`)

// NewMigrator returns a new "blank" migrator. It is recommended
// to use something like MigrationBox or FileMigrator. A "blank"
// Migrator should only be used as the basis for a new type of
// migration system.
func NewMigrator(c *Connection) Migrator {
	return Migrator{
		Connection: c,
		Migrations: map[string]Migrations{
			"up":   Migrations{},
			"down": Migrations{},
		},
	}
}

// Migrator forms the basis of all migrations systems.
// It does the actual heavy lifting of running migrations.
// When building a new migration system, you should embed this
// type into your migrator.
type Migrator struct {
	Connection *Connection
	SchemaPath string
	Migrations map[string]Migrations
}

// Up runs pending "up" migrations and applies them to the database.
func (m Migrator) Up() error {
	c := m.Connection
	return m.exec(func() error {
		mtn := c.MigrationTableName()
		mfs := m.Migrations["up"]
		sort.Sort(mfs)
		for _, mi := range mfs {
			if mi.DBType != "all" && mi.DBType != c.Dialect.Name() {
				// Skip migration for non-matching dialect
				continue
			}
			exists, err := c.Where("version = ?", mi.Version).Exists(mtn)
			if err != nil {
				return errors.Wrapf(err, "problem checking for migration version %s", mi.Version)
			}
			if exists {
				continue
			}
			err = c.Transaction(func(tx *Connection) error {
				err := mi.Run(tx)
				if err != nil {
					return err
				}
				_, err = tx.Store.Exec(fmt.Sprintf("insert into %s (version) values ('%s')", mtn, mi.Version))
				return errors.Wrapf(err, "problem inserting migration version %s", mi.Version)
			})
			if err != nil {
				return errors.WithStack(err)
			}
			fmt.Printf("> %s\n", mi.Name)
		}
		return nil
	})
}

// Down runs pending "down" migrations and rolls back the
// database by the specified number of steps.
func (m Migrator) Down(step int) error {
	c := m.Connection
	return m.exec(func() error {
		mtn := c.MigrationTableName()
		count, err := c.Count(mtn)
		if err != nil {
			return errors.Wrap(err, "migration down: unable count existing migration")
		}
		mfs := m.Migrations["down"]
		sort.Sort(sort.Reverse(mfs))
		// skip all runned migration
		if len(mfs) > count {
			mfs = mfs[len(mfs)-count:]
		}
		// run only required steps
		if step > 0 && len(mfs) >= step {
			mfs = mfs[:step]
		}
		for _, mi := range mfs {
			exists, err := c.Where("version = ?", mi.Version).Exists(mtn)
			if err != nil || !exists {
				return errors.Wrapf(err, "problem checking for migration version %s", mi.Version)
			}
			err = c.Transaction(func(tx *Connection) error {
				err := mi.Run(tx)
				if err != nil {
					return err
				}
				err = tx.RawQuery("delete from %s where version = ?", mtn, mi.Version).Exec()
				return errors.Wrapf(err, "problem deleting migration version %s", mi.Version)
			})
			if err != nil {
				return err
			}

			fmt.Printf("< %s\n", mi.Name)
		}
		return nil
	})
}

// Reset the database by runing the down migrations followed by the up migrations.
func (m Migrator) Reset() error {
	err := m.Down(-1)
	if err != nil {
		return errors.WithStack(err)
	}
	return m.Up()
}

// CreateSchemaMigrations sets up a table to track migrations. This is an idempotent
// operation.
func (m Migrator) CreateSchemaMigrations() error {
	c := m.Connection
	mtn := c.MigrationTableName()
	err := c.Open()
	if err != nil {
		return errors.Wrap(err, "could not open connection")
	}
	_, err = c.Store.Exec(fmt.Sprintf("select * from %s", mtn))
	if err == nil {
		return nil
	}

	return c.Transaction(func(tx *Connection) error {
		schemaMigrations := newSchemaMigrations(mtn)
		smSQL, err := c.Dialect.FizzTranslator().CreateTable(schemaMigrations)
		if err != nil {
			return errors.Wrap(err, "could not build SQL for schema migration table")
		}
		err = tx.RawQuery(smSQL).Exec()
		if err != nil {
			return errors.WithStack(errors.Wrap(err, smSQL))
		}
		return nil
	})
}

// Status prints out the status of applied/pending migrations.
func (m Migrator) Status() error {
	err := m.CreateSchemaMigrations()
	if err != nil {
		return errors.WithStack(err)
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', tabwriter.TabIndent)
	fmt.Fprintln(w, "Version\tName\tStatus\t")
	for _, mf := range m.Migrations["up"] {
		exists, err := m.Connection.Where("version = ?", mf.Version).Exists(m.Connection.MigrationTableName())
		if err != nil {
			return errors.Wrapf(err, "problem with migration")
		}
		state := "Pending"
		if exists {
			state = "Applied"
		}
		fmt.Fprintf(w, "%s\t%s\t%s\t\n", mf.Version, mf.Name, state)
	}
	return w.Flush()
}

// DumpMigrationSchema will generate a file of the current database schema
// based on the value of Migrator.SchemaPath
func (m Migrator) DumpMigrationSchema() error {
	if m.SchemaPath == "" {
		return nil
	}
	c := m.Connection
	f, err := os.Create(filepath.Join(m.SchemaPath, "schema.sql"))
	if err != nil {
		return errors.WithStack(err)
	}
	err = c.Dialect.DumpSchema(f)
	if err != nil {

		return errors.WithStack(err)
	}
	return nil
}

func (m Migrator) exec(fn func() error) error {
	now := time.Now()
	defer m.DumpMigrationSchema()
	defer printTimer(now)

	err := m.CreateSchemaMigrations()
	if err != nil {
		return errors.Wrap(err, "Migrator: problem creating schema migrations")
	}
	return fn()
}

func printTimer(timerStart time.Time) {
	diff := time.Now().Sub(timerStart).Seconds()
	if diff > 60 {
		fmt.Printf("\n%.4f minutes\n", diff/60)
	} else {
		fmt.Printf("\n%.4f seconds\n", diff)
	}
}
<commit_msg>Fix a regression with migrations (#69)<commit_after>package pop

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"text/tabwriter"
	"time"

	"github.com/pkg/errors"
)

var mrx = regexp.MustCompile(`(\d+)_([^\.]+)(\.[a-z]+)?\.(up|down)\.(sql|fizz)`)

// NewMigrator returns a new "blank" migrator. It is recommended
// to use something like MigrationBox or FileMigrator. A "blank"
// Migrator should only be used as the basis for a new type of
// migration system.
func NewMigrator(c *Connection) Migrator {
	return Migrator{
		Connection: c,
		Migrations: map[string]Migrations{
			"up":   Migrations{},
			"down": Migrations{},
		},
	}
}

// Migrator forms the basis of all migrations systems.
// It does the actual heavy lifting of running migrations.
// When building a new migration system, you should embed this
// type into your migrator.
type Migrator struct {
	Connection *Connection
	SchemaPath string
	Migrations map[string]Migrations
}

// Up runs pending "up" migrations and applies them to the database.
func (m Migrator) Up() error {
	c := m.Connection
	return m.exec(func() error {
		mtn := c.MigrationTableName()
		mfs := m.Migrations["up"]
		sort.Sort(mfs)
		for _, mi := range mfs {
			if mi.DBType != "all" && mi.DBType != c.Dialect.Name() {
				// Skip migration for non-matching dialect
				continue
			}
			exists, err := c.Where("version = ?", mi.Version).Exists(mtn)
			if err != nil {
				return errors.Wrapf(err, "problem checking for migration version %s", mi.Version)
			}
			if exists {
				continue
			}
			err = c.Transaction(func(tx *Connection) error {
				err := mi.Run(tx)
				if err != nil {
					return err
				}
				_, err = tx.Store.Exec(fmt.Sprintf("insert into %s (version) values ('%s')", mtn, mi.Version))
				return errors.Wrapf(err, "problem inserting migration version %s", mi.Version)
			})
			if err != nil {
				return errors.WithStack(err)
			}
			fmt.Printf("> %s\n", mi.Name)
		}
		return nil
	})
}

// Down runs pending "down" migrations and rolls back the
// database by the specified number of steps.
func (m Migrator) Down(step int) error {
	c := m.Connection
	return m.exec(func() error {
		mtn := c.MigrationTableName()
		count, err := c.Count(mtn)
		if err != nil {
			return errors.Wrap(err, "migration down: unable count existing migration")
		}
		mfs := m.Migrations["down"]
		sort.Sort(sort.Reverse(mfs))
		// skip all runned migration
		if len(mfs) > count {
			mfs = mfs[len(mfs)-count:]
		}
		// run only required steps
		if step > 0 && len(mfs) >= step {
			mfs = mfs[:step]
		}
		for _, mi := range mfs {
			exists, err := c.Where("version = ?", mi.Version).Exists(mtn)
			if err != nil || !exists {
				return errors.Wrapf(err, "problem checking for migration version %s", mi.Version)
			}
			err = c.Transaction(func(tx *Connection) error {
				err := mi.Run(tx)
				if err != nil {
					return err
				}
				err = tx.RawQuery(fmt.Sprintf("delete from %s where version = ?", mtn), mi.Version).Exec()
				return errors.Wrapf(err, "problem deleting migration version %s", mi.Version)
			})
			if err != nil {
				return err
			}

			fmt.Printf("< %s\n", mi.Name)
		}
		return nil
	})
}

// Reset the database by runing the down migrations followed by the up migrations.
func (m Migrator) Reset() error {
	err := m.Down(-1)
	if err != nil {
		return errors.WithStack(err)
	}
	return m.Up()
}

// CreateSchemaMigrations sets up a table to track migrations. This is an idempotent
// operation.
func (m Migrator) CreateSchemaMigrations() error {
	c := m.Connection
	mtn := c.MigrationTableName()
	err := c.Open()
	if err != nil {
		return errors.Wrap(err, "could not open connection")
	}
	_, err = c.Store.Exec(fmt.Sprintf("select * from %s", mtn))
	if err == nil {
		return nil
	}

	return c.Transaction(func(tx *Connection) error {
		schemaMigrations := newSchemaMigrations(mtn)
		smSQL, err := c.Dialect.FizzTranslator().CreateTable(schemaMigrations)
		if err != nil {
			return errors.Wrap(err, "could not build SQL for schema migration table")
		}
		err = tx.RawQuery(smSQL).Exec()
		if err != nil {
			return errors.WithStack(errors.Wrap(err, smSQL))
		}
		return nil
	})
}

// Status prints out the status of applied/pending migrations.
func (m Migrator) Status() error {
	err := m.CreateSchemaMigrations()
	if err != nil {
		return errors.WithStack(err)
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', tabwriter.TabIndent)
	fmt.Fprintln(w, "Version\tName\tStatus\t")
	for _, mf := range m.Migrations["up"] {
		exists, err := m.Connection.Where("version = ?", mf.Version).Exists(m.Connection.MigrationTableName())
		if err != nil {
			return errors.Wrapf(err, "problem with migration")
		}
		state := "Pending"
		if exists {
			state = "Applied"
		}
		fmt.Fprintf(w, "%s\t%s\t%s\t\n", mf.Version, mf.Name, state)
	}
	return w.Flush()
}

// DumpMigrationSchema will generate a file of the current database schema
// based on the value of Migrator.SchemaPath
func (m Migrator) DumpMigrationSchema() error {
	if m.SchemaPath == "" {
		return nil
	}
	c := m.Connection
	f, err := os.Create(filepath.Join(m.SchemaPath, "schema.sql"))
	if err != nil {
		return errors.WithStack(err)
	}
	err = c.Dialect.DumpSchema(f)
	if err != nil {

		return errors.WithStack(err)
	}
	return nil
}

func (m Migrator) exec(fn func() error) error {
	now := time.Now()
	defer m.DumpMigrationSchema()
	defer printTimer(now)

	err := m.CreateSchemaMigrations()
	if err != nil {
		return errors.Wrap(err, "Migrator: problem creating schema migrations")
	}
	return fn()
}

func printTimer(timerStart time.Time) {
	diff := time.Now().Sub(timerStart).Seconds()
	if diff > 60 {
		fmt.Printf("\n%.4f minutes\n", diff/60)
	} else {
		fmt.Printf("\n%.4f seconds\n", diff)
	}
}
<|endoftext|>


<commit_before>package masters

import (
	"context"
	"fmt"

	"github.com/giantswarm/microerror"
	"github.com/giantswarm/to"

	"github.com/giantswarm/azure-operator/v4/pkg/checksum"
	"github.com/giantswarm/azure-operator/v4/service/controller/blobclient"
	"github.com/giantswarm/azure-operator/v4/service/controller/internal/state"
	"github.com/giantswarm/azure-operator/v4/service/controller/key"
)

func (r *Resource) deploymentCompletedTransition(ctx context.Context, obj interface{}, currentState state.State) (state.State, error) {
	cr, err := key.ToCustomResource(obj)
	if err != nil {
		return Empty, microerror.Mask(err)
	}
	deploymentsClient, err := r.getDeploymentsClient(ctx)
	if err != nil {
		return Empty, microerror.Mask(err)
	}

	d, err := deploymentsClient.Get(ctx, key.ClusterID(cr), key.MastersVmssDeploymentName)
	if IsDeploymentNotFound(err) {
		r.logger.LogCtx(ctx, "level", "debug", "message", "deployment should be completed but is not found")
		r.logger.LogCtx(ctx, "level", "debug", "message", "going back to DeploymentUninitialized")
		return Empty, nil
	} else if err != nil {
		return Empty, microerror.Mask(err)
	}

	s := *d.Properties.ProvisioningState
	r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("deployment is in state '%s'", s))

	if key.IsSucceededProvisioningState(s) {
		if d.Location == nil {
			d.Location = to.StringP("<unknown>")
		}

		computedDeployment, err := r.newDeployment(ctx, cr, nil, *d.Location)
		if blobclient.IsBlobNotFound(err) {
			r.logger.LogCtx(ctx, "level", "debug", "message", "ignition blob not found")
			return currentState, nil
		} else if err != nil {
			return "", microerror.Mask(err)
		} else {
			desiredDeploymentTemplateChk, err := checksum.GetDeploymentTemplateChecksum(computedDeployment)
			if err != nil {
				return "", microerror.Mask(err)
			}

			desiredDeploymentParametersChk, err := checksum.GetDeploymentParametersChecksum(computedDeployment)
			if err != nil {
				return "", microerror.Mask(err)
			}

			currentDeploymentTemplateChk, err := r.getResourceStatus(cr, DeploymentTemplateChecksum)
			if err != nil {
				return "", microerror.Mask(err)
			}

			currentDeploymentParametersChk, err := r.getResourceStatus(cr, DeploymentParametersChecksum)
			if err != nil {
				return "", microerror.Mask(err)
			}

			if currentDeploymentTemplateChk != desiredDeploymentTemplateChk || currentDeploymentParametersChk != desiredDeploymentParametersChk {
				r.logger.LogCtx(ctx, "level", "debug", "message", "template or parameters changed")
				// As current and desired state differs, start process from the beginning.
				return Empty, nil
			}

			r.logger.LogCtx(ctx, "level", "debug", "message", "template and parameters unchanged")

			return currentState, nil
		}
	} else if key.IsFinalProvisioningState(s) {
		// Deployment has failed. Restart from beginning.
		return Empty, nil
	}

	r.logger.LogCtx(ctx, "level", "warning", "message", "instances reconciliation process reached unexpected state")

	// Normally the process should never get here. In case this happens, start
	// from the beginning.
	return Empty, nil
}
<commit_msg>Fix checkcum calculation for GWC region (#834)<commit_after>package masters

import (
	"context"
	"fmt"

	"github.com/giantswarm/microerror"

	"github.com/giantswarm/azure-operator/v4/pkg/checksum"
	"github.com/giantswarm/azure-operator/v4/service/controller/blobclient"
	"github.com/giantswarm/azure-operator/v4/service/controller/internal/state"
	"github.com/giantswarm/azure-operator/v4/service/controller/key"
)

func (r *Resource) deploymentCompletedTransition(ctx context.Context, obj interface{}, currentState state.State) (state.State, error) {
	cr, err := key.ToCustomResource(obj)
	if err != nil {
		return Empty, microerror.Mask(err)
	}
	deploymentsClient, err := r.getDeploymentsClient(ctx)
	if err != nil {
		return Empty, microerror.Mask(err)
	}
	groupsClient, err := r.getGroupsClient(ctx)
	if err != nil {
		return currentState, microerror.Mask(err)
	}

	d, err := deploymentsClient.Get(ctx, key.ClusterID(cr), key.MastersVmssDeploymentName)
	if IsDeploymentNotFound(err) {
		r.logger.LogCtx(ctx, "level", "debug", "message", "deployment should be completed but is not found")
		r.logger.LogCtx(ctx, "level", "debug", "message", "going back to DeploymentUninitialized")
		return Empty, nil
	} else if err != nil {
		return Empty, microerror.Mask(err)
	}

	s := *d.Properties.ProvisioningState
	r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("deployment is in state '%s'", s))

	group, err := groupsClient.Get(ctx, key.ClusterID(cr))
	if err != nil {
		return currentState, microerror.Mask(err)
	}

	if key.IsSucceededProvisioningState(s) {
		computedDeployment, err := r.newDeployment(ctx, cr, nil, *group.Location)
		if blobclient.IsBlobNotFound(err) {
			r.logger.LogCtx(ctx, "level", "debug", "message", "ignition blob not found")
			return currentState, nil
		} else if err != nil {
			return "", microerror.Mask(err)
		} else {
			desiredDeploymentTemplateChk, err := checksum.GetDeploymentTemplateChecksum(computedDeployment)
			if err != nil {
				return "", microerror.Mask(err)
			}

			desiredDeploymentParametersChk, err := checksum.GetDeploymentParametersChecksum(computedDeployment)
			if err != nil {
				return "", microerror.Mask(err)
			}

			currentDeploymentTemplateChk, err := r.getResourceStatus(cr, DeploymentTemplateChecksum)
			if err != nil {
				return "", microerror.Mask(err)
			}

			currentDeploymentParametersChk, err := r.getResourceStatus(cr, DeploymentParametersChecksum)
			if err != nil {
				return "", microerror.Mask(err)
			}

			if currentDeploymentTemplateChk != desiredDeploymentTemplateChk || currentDeploymentParametersChk != desiredDeploymentParametersChk {
				r.logger.LogCtx(ctx, "level", "debug", "message", "template or parameters changed")
				// As current and desired state differs, start process from the beginning.
				return Empty, nil
			}

			r.logger.LogCtx(ctx, "level", "debug", "message", "template and parameters unchanged")

			return currentState, nil
		}
	} else if key.IsFinalProvisioningState(s) {
		// Deployment has failed. Restart from beginning.
		return Empty, nil
	}

	r.logger.LogCtx(ctx, "level", "warning", "message", "instances reconciliation process reached unexpected state")

	// Normally the process should never get here. In case this happens, start
	// from the beginning.
	return Empty, nil
}
<|endoftext|>


<commit_before>package shared

import (
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/bytefmt"
	"code.cloudfoundry.org/cli/actor/v7action"
	"code.cloudfoundry.org/cli/api/cloudcontroller/ccv3/constant"
	"code.cloudfoundry.org/cli/command"
	log "github.com/sirupsen/logrus"
)

type AppSummaryDisplayer struct {
	UI command.UI
}

func NewAppSummaryDisplayer(ui command.UI) *AppSummaryDisplayer {
	return &AppSummaryDisplayer{
		UI: ui,
	}
}

func (display AppSummaryDisplayer) AppDisplay(summary v7action.DetailedApplicationSummary, displayStartCommand bool) {
	var isoRow []string
	if name, exists := summary.GetIsolationSegmentName(); exists {
		isoRow = append(isoRow, display.UI.TranslateText("isolation segment:"), name)
	}

	var lifecycleInfo []string
	if summary.LifecycleType == constant.AppLifecycleTypeDocker {
		lifecycleInfo = []string{display.UI.TranslateText("docker image:"), summary.CurrentDroplet.Image}
	} else {
		lifecycleInfo = []string{display.UI.TranslateText("buildpacks:"), display.buildpackNames(summary.CurrentDroplet.Buildpacks)}
	}

	keyValueTable := [][]string{
		{display.UI.TranslateText("name:"), summary.Application.Name},
		{display.UI.TranslateText("requested state:"), strings.ToLower(string(summary.State))},
		isoRow,
		{display.UI.TranslateText("routes:"), routeSummary(summary.Routes)},
		{display.UI.TranslateText("last uploaded:"), display.getCreatedTime(summary)},
		{display.UI.TranslateText("stack:"), summary.CurrentDroplet.Stack},
		lifecycleInfo,
	}

	display.UI.DisplayKeyValueTable("", keyValueTable, 3)

	display.displayProcessTable(summary, displayStartCommand)
}

func routeSummary(rs []v7action.Route) string {
	formattedRoutes := []string{}
	for _, route := range rs {
		formattedRoutes = append(formattedRoutes, route.URL)
	}
	return strings.Join(formattedRoutes, ", ")
}

func (display AppSummaryDisplayer) displayAppInstancesTable(processSummary v7action.ProcessSummary) {
	table := [][]string{
		{
			"",
			display.UI.TranslateText("state"),
			display.UI.TranslateText("since"),
			display.UI.TranslateText("cpu"),
			display.UI.TranslateText("memory"),
			display.UI.TranslateText("disk"),
			display.UI.TranslateText("details"),
		},
	}

	for _, instance := range processSummary.InstanceDetails {
		table = append(table, []string{
			fmt.Sprintf("#%d", instance.Index),
			display.UI.TranslateText(strings.ToLower(string(instance.State))),
			display.appInstanceDate(instance.StartTime()),
			fmt.Sprintf("%.1f%%", instance.CPU*100),
			display.UI.TranslateText("{{.MemUsage}} of {{.MemQuota}}", map[string]interface{}{
				"MemUsage": bytefmt.ByteSize(instance.MemoryUsage),
				"MemQuota": bytefmt.ByteSize(instance.MemoryQuota),
			}),
			display.UI.TranslateText("{{.DiskUsage}} of {{.DiskQuota}}", map[string]interface{}{
				"DiskUsage": bytefmt.ByteSize(instance.DiskUsage),
				"DiskQuota": bytefmt.ByteSize(instance.DiskQuota),
			}),
			instance.Details,
		})
	}

	display.UI.DisplayInstancesTableForApp(table)
}

func (display AppSummaryDisplayer) displayProcessTable(summary v7action.DetailedApplicationSummary, displayStartCommand bool) {
	for _, process := range summary.ProcessSummaries {
		display.UI.DisplayNewline()

		var startCommandRow []string
		if displayStartCommand && len(process.Command.Value) > 0 {
			startCommandRow = append(startCommandRow, display.UI.TranslateText("start command:"), process.Command.Value)
		}

		var processSidecars []string
		for _, sidecar := range process.Sidecars {
			processSidecars = append(processSidecars, sidecar.Name)
		}

		keyValueTable := [][]string{
			{display.UI.TranslateText("type:"), process.Type},
			{display.UI.TranslateText("sidecars:"), strings.Join(processSidecars, ", ")},
			{display.UI.TranslateText("instances:"), fmt.Sprintf("%d/%d", process.HealthyInstanceCount(), process.TotalInstanceCount())},
			{display.UI.TranslateText("memory usage:"), fmt.Sprintf("%dM", process.MemoryInMB.Value)},
			startCommandRow,
		}

		display.UI.DisplayKeyValueTable("", keyValueTable, 3)

		if len(process.InstanceDetails) == 0 {
			display.UI.DisplayNewline()
			display.UI.DisplayText("There are no running instances of this process.")
			continue
		}
		display.displayAppInstancesTable(process)
	}
}

func (display AppSummaryDisplayer) getCreatedTime(summary v7action.DetailedApplicationSummary) string {
	if summary.CurrentDroplet.CreatedAt != "" {
		timestamp, err := time.Parse(time.RFC3339, summary.CurrentDroplet.CreatedAt)
		if err != nil {
			log.WithField("createdAt", summary.CurrentDroplet.CreatedAt).Errorln("error parsing created at:", err)
		}

		return display.UI.UserFriendlyDate(timestamp)
	}

	return ""
}

func (AppSummaryDisplayer) buildpackNames(buildpacks []v7action.DropletBuildpack) string {
	var names []string
	for _, buildpack := range buildpacks {
		if buildpack.DetectOutput != "" {
			names = append(names, buildpack.DetectOutput)
		} else {
			names = append(names, buildpack.Name)
		}
	}

	return strings.Join(names, ", ")
}

func (AppSummaryDisplayer) appInstanceDate(input time.Time) string {
	return input.UTC().Format(time.RFC3339)
}
<commit_msg>Remove newline before 'no running instances' in app_summary_displayer<commit_after>package shared

import (
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/bytefmt"
	"code.cloudfoundry.org/cli/actor/v7action"
	"code.cloudfoundry.org/cli/api/cloudcontroller/ccv3/constant"
	"code.cloudfoundry.org/cli/command"
	log "github.com/sirupsen/logrus"
)

type AppSummaryDisplayer struct {
	UI command.UI
}

func NewAppSummaryDisplayer(ui command.UI) *AppSummaryDisplayer {
	return &AppSummaryDisplayer{
		UI: ui,
	}
}

func (display AppSummaryDisplayer) AppDisplay(summary v7action.DetailedApplicationSummary, displayStartCommand bool) {
	var isoRow []string
	if name, exists := summary.GetIsolationSegmentName(); exists {
		isoRow = append(isoRow, display.UI.TranslateText("isolation segment:"), name)
	}

	var lifecycleInfo []string
	if summary.LifecycleType == constant.AppLifecycleTypeDocker {
		lifecycleInfo = []string{display.UI.TranslateText("docker image:"), summary.CurrentDroplet.Image}
	} else {
		lifecycleInfo = []string{display.UI.TranslateText("buildpacks:"), display.buildpackNames(summary.CurrentDroplet.Buildpacks)}
	}

	keyValueTable := [][]string{
		{display.UI.TranslateText("name:"), summary.Application.Name},
		{display.UI.TranslateText("requested state:"), strings.ToLower(string(summary.State))},
		isoRow,
		{display.UI.TranslateText("routes:"), routeSummary(summary.Routes)},
		{display.UI.TranslateText("last uploaded:"), display.getCreatedTime(summary)},
		{display.UI.TranslateText("stack:"), summary.CurrentDroplet.Stack},
		lifecycleInfo,
	}

	display.UI.DisplayKeyValueTable("", keyValueTable, 3)

	display.displayProcessTable(summary, displayStartCommand)
}

func routeSummary(rs []v7action.Route) string {
	formattedRoutes := []string{}
	for _, route := range rs {
		formattedRoutes = append(formattedRoutes, route.URL)
	}
	return strings.Join(formattedRoutes, ", ")
}

func (display AppSummaryDisplayer) displayAppInstancesTable(processSummary v7action.ProcessSummary) {
	table := [][]string{
		{
			"",
			display.UI.TranslateText("state"),
			display.UI.TranslateText("since"),
			display.UI.TranslateText("cpu"),
			display.UI.TranslateText("memory"),
			display.UI.TranslateText("disk"),
			display.UI.TranslateText("details"),
		},
	}

	for _, instance := range processSummary.InstanceDetails {
		table = append(table, []string{
			fmt.Sprintf("#%d", instance.Index),
			display.UI.TranslateText(strings.ToLower(string(instance.State))),
			display.appInstanceDate(instance.StartTime()),
			fmt.Sprintf("%.1f%%", instance.CPU*100),
			display.UI.TranslateText("{{.MemUsage}} of {{.MemQuota}}", map[string]interface{}{
				"MemUsage": bytefmt.ByteSize(instance.MemoryUsage),
				"MemQuota": bytefmt.ByteSize(instance.MemoryQuota),
			}),
			display.UI.TranslateText("{{.DiskUsage}} of {{.DiskQuota}}", map[string]interface{}{
				"DiskUsage": bytefmt.ByteSize(instance.DiskUsage),
				"DiskQuota": bytefmt.ByteSize(instance.DiskQuota),
			}),
			instance.Details,
		})
	}

	display.UI.DisplayInstancesTableForApp(table)
}

func (display AppSummaryDisplayer) displayProcessTable(summary v7action.DetailedApplicationSummary, displayStartCommand bool) {
	for _, process := range summary.ProcessSummaries {
		display.UI.DisplayNewline()

		var startCommandRow []string
		if displayStartCommand && len(process.Command.Value) > 0 {
			startCommandRow = append(startCommandRow, display.UI.TranslateText("start command:"), process.Command.Value)
		}

		var processSidecars []string
		for _, sidecar := range process.Sidecars {
			processSidecars = append(processSidecars, sidecar.Name)
		}

		keyValueTable := [][]string{
			{display.UI.TranslateText("type:"), process.Type},
			{display.UI.TranslateText("sidecars:"), strings.Join(processSidecars, ", ")},
			{display.UI.TranslateText("instances:"), fmt.Sprintf("%d/%d", process.HealthyInstanceCount(), process.TotalInstanceCount())},
			{display.UI.TranslateText("memory usage:"), fmt.Sprintf("%dM", process.MemoryInMB.Value)},
			startCommandRow,
		}

		display.UI.DisplayKeyValueTable("", keyValueTable, 3)

		if len(process.InstanceDetails) == 0 {
			display.UI.DisplayText("There are no running instances of this process.")
			continue
		}
		display.displayAppInstancesTable(process)
	}
}

func (display AppSummaryDisplayer) getCreatedTime(summary v7action.DetailedApplicationSummary) string {
	if summary.CurrentDroplet.CreatedAt != "" {
		timestamp, err := time.Parse(time.RFC3339, summary.CurrentDroplet.CreatedAt)
		if err != nil {
			log.WithField("createdAt", summary.CurrentDroplet.CreatedAt).Errorln("error parsing created at:", err)
		}

		return display.UI.UserFriendlyDate(timestamp)
	}

	return ""
}

func (AppSummaryDisplayer) buildpackNames(buildpacks []v7action.DropletBuildpack) string {
	var names []string
	for _, buildpack := range buildpacks {
		if buildpack.DetectOutput != "" {
			names = append(names, buildpack.DetectOutput)
		} else {
			names = append(names, buildpack.Name)
		}
	}

	return strings.Join(names, ", ")
}

func (AppSummaryDisplayer) appInstanceDate(input time.Time) string {
	return input.UTC().Format(time.RFC3339)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// kube2sky is a bridge between Kubernetes and SkyDNS.  It watches the
// Kubernetes master for changes in Services and manifests them into etcd for
// SkyDNS to serve as DNS records.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	klabels "github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	kwatch "github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
	etcd "github.com/coreos/go-etcd/etcd"
	skymsg "github.com/skynetservices/skydns/msg"
)

var (
	domain                = flag.String("domain", "kubernetes.local", "domain under which to create names")
	verbose               = flag.Bool("verbose", false, "log extra information")
	etcd_mutation_timeout = flag.Duration("etcd_mutation_timeout", 10*time.Second, "crash after retrying etcd mutation for a specified duration")
)

func removeDNS(record string, etcdClient *etcd.Client) error {
	log.Printf("Removing %s from DNS", record)
	_, err := etcdClient.Delete(skymsg.Path(record), true)
	return err
}

func addDNS(record string, service *kapi.Service, etcdClient *etcd.Client) error {
	svc := skymsg.Service{
		Host:     service.Spec.PortalIP,
		Port:     service.Spec.Port,
		Priority: 10,
		Weight:   10,
		Ttl:      30,
	}
	b, err := json.Marshal(svc)
	if err != nil {
		return err
	}
	// Set with no TTL, and hope that kubernetes events are accurate.

	log.Printf("Setting dns record: %v -> %s:%d\n", record, service.Spec.PortalIP, service.Spec.Port)
	_, err = etcdClient.Set(skymsg.Path(record), string(b), uint64(0))
	return err
}

// Implements retry logic for arbitrary mutator. Crashes after retrying for
// etcd_mutation_timeout.
func mutateEtcdOrDie(mutator func() error) {
	timeout := time.After(*etcd_mutation_timeout)
	for {
		select {
		case <-timeout:
			log.Fatalf("Failed to mutate etcd for %v using mutator: %v", *etcd_mutation_timeout, mutator)
		default:
			if err := mutator(); err != nil {
				delay := 50 * time.Millisecond
				log.Printf("Failed to mutate etcd using mutator: %v due to: %v. Will retry in: %v", mutator, err, delay)
				time.Sleep(delay)
			} else {
				return
			}
		}
	}
}

func newEtcdClient() (client *etcd.Client) {
	// TODO: take a flag for etcd server(s).
	client = etcd.NewClient([]string{"http://127.0.0.1:4001"})
	if client == nil {
		return nil
	}
	client.SyncCluster()

	return client
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	config := &kclient.Config{}

	masterHost := os.Getenv("KUBERNETES_RO_SERVICE_HOST")
	if masterHost == "" {
		log.Fatalf("KUBERNETES_RO_SERVICE_HOST is not defined")
	}
	masterPort := os.Getenv("KUBERNETES_RO_SERVICE_PORT")
	if masterPort == "" {
		log.Fatalf("KUBERNETES_RO_SERVICE_PORT is not defined")
	}
	config.Host = fmt.Sprintf("http://%s:%s", masterHost, masterPort)
	log.Printf("Using %s for kubernetes master", config.Host)

	config.Version = "v1beta1"
	log.Printf("Using kubernetes API %s", config.Version)

	return kclient.New(config)
}

func buildNameString(service, namespace, domain string) string {
	return fmt.Sprintf("%s.%s.%s.", service, namespace, domain)
}

func watchOnce(etcdClient *etcd.Client, kubeClient *kclient.Client) {
	// Start the goroutine to produce update events.
	updates := make(chan serviceUpdate)
	startWatching(kubeClient.Services(kapi.NamespaceAll), updates)

	// This loop will break if the channel closes, which is how the
	// goroutine signals an error.
	for ev := range updates {
		if *verbose {
			log.Printf("Received update event: %#v", ev)
		}
		switch ev.Op {
		case SetServices, AddService:
			for i := range ev.Services {
				s := &ev.Services[i]
				name := buildNameString(s.Name, s.Namespace, *domain)
				mutateEtcdOrDie(func() error { return addDNS(name, s, etcdClient) })
			}
		case RemoveService:
			for i := range ev.Services {
				s := &ev.Services[i]
				name := buildNameString(s.Name, s.Namespace, *domain)
				mutateEtcdOrDie(func() error { return removeDNS(name, etcdClient) })
			}
		}
	}
	//TODO: fully resync periodically.
}

func main() {
	flag.Parse()

	etcdClient := newEtcdClient()
	if etcdClient == nil {
		log.Fatal("Failed to create etcd client")
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		log.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	// In case of error, the watch will be aborted.  At that point we just
	// retry.
	for {
		watchOnce(etcdClient, kubeClient)
	}
}

//FIXME: make the below part of the k8s client lib?

// servicesWatcher is capable of listing and watching for changes to services
// across ALL namespaces
type servicesWatcher interface {
	List(label klabels.Selector) (*kapi.ServiceList, error)
	Watch(label, field klabels.Selector, resourceVersion string) (kwatch.Interface, error)
}

type operation int

// These are the available operation types.
const (
	SetServices operation = iota
	AddService
	RemoveService
)

// serviceUpdate describes an operation of services, sent on the channel.
//
// You can add or remove a single service by sending an array of size one with
// Op == AddService|RemoveService.  For setting the state of the system to a given state, just
// set Services as desired and Op to SetServices, which will reset the system
// state to that specified in this operation for this source channel. To remove
// all services, set Services to empty array and Op to SetServices
type serviceUpdate struct {
	Services []kapi.Service
	Op       operation
}

// startWatching launches a goroutine that watches for changes to services.
func startWatching(watcher servicesWatcher, updates chan<- serviceUpdate) {
	serviceVersion := ""
	go watchLoop(watcher, updates, &serviceVersion)
}

// watchLoop loops forever looking for changes to services.  If an error occurs
// it will close the channel and return.
func watchLoop(svcWatcher servicesWatcher, updates chan<- serviceUpdate, resourceVersion *string) {
	defer close(updates)

	if len(*resourceVersion) == 0 {
		services, err := svcWatcher.List(klabels.Everything())
		if err != nil {
			log.Printf("Failed to load services: %v", err)
			return
		}
		*resourceVersion = services.ResourceVersion
		updates <- serviceUpdate{Op: SetServices, Services: services.Items}
	}

	watcher, err := svcWatcher.Watch(klabels.Everything(), klabels.Everything(), *resourceVersion)
	if err != nil {
		log.Printf("Failed to watch for service changes: %v", err)
		return
	}
	defer watcher.Stop()

	ch := watcher.ResultChan()
	for {
		select {
		case event, ok := <-ch:
			if !ok {
				log.Printf("watchLoop channel closed")
				return
			}

			if event.Type == kwatch.Error {
				if status, ok := event.Object.(*kapi.Status); ok {
					log.Printf("Error during watch: %#v", status)
					return
				}
				log.Fatalf("Received unexpected error: %#v", event.Object)
			}

			if service, ok := event.Object.(*kapi.Service); ok {
				sendUpdate(updates, event, service, resourceVersion)
				continue
			}
		}
	}
}

func sendUpdate(updates chan<- serviceUpdate, event kwatch.Event, service *kapi.Service, resourceVersion *string) {
	*resourceVersion = service.ResourceVersion

	switch event.Type {
	case kwatch.Added, kwatch.Modified:
		updates <- serviceUpdate{Op: AddService, Services: []kapi.Service{*service}}
	case kwatch.Deleted:
		updates <- serviceUpdate{Op: RemoveService, Services: []kapi.Service{*service}}
	default:
		log.Fatalf("Unknown event.Type: %v", event.Type)
	}
}
<commit_msg>Add etcd server check at startup to kube2sky (with retry loop).<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// kube2sky is a bridge between Kubernetes and SkyDNS.  It watches the
// Kubernetes master for changes in Services and manifests them into etcd for
// SkyDNS to serve as DNS records.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	klabels "github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	tools "github.com/GoogleCloudPlatform/kubernetes/pkg/tools"
	kwatch "github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
	etcd "github.com/coreos/go-etcd/etcd"
	skymsg "github.com/skynetservices/skydns/msg"
)

var (
	domain                = flag.String("domain", "kubernetes.local", "domain under which to create names")
	etcd_mutation_timeout = flag.Duration("etcd_mutation_timeout", 10*time.Second, "crash after retrying etcd mutation for a specified duration")
	etcd_server           = flag.String("etcd-server", "http://127.0.0.1:4001", "URL to etcd server")
	verbose               = flag.Bool("verbose", false, "log extra information")
)

func removeDNS(record string, etcdClient *etcd.Client) error {
	log.Printf("Removing %s from DNS", record)
	_, err := etcdClient.Delete(skymsg.Path(record), true)
	return err
}

func addDNS(record string, service *kapi.Service, etcdClient *etcd.Client) error {
	svc := skymsg.Service{
		Host:     service.Spec.PortalIP,
		Port:     service.Spec.Port,
		Priority: 10,
		Weight:   10,
		Ttl:      30,
	}
	b, err := json.Marshal(svc)
	if err != nil {
		return err
	}
	// Set with no TTL, and hope that kubernetes events are accurate.

	log.Printf("Setting dns record: %v -> %s:%d\n", record, service.Spec.PortalIP, service.Spec.Port)
	_, err = etcdClient.Set(skymsg.Path(record), string(b), uint64(0))
	return err
}

// Implements retry logic for arbitrary mutator. Crashes after retrying for
// etcd_mutation_timeout.
func mutateEtcdOrDie(mutator func() error) {
	timeout := time.After(*etcd_mutation_timeout)
	for {
		select {
		case <-timeout:
			log.Fatalf("Failed to mutate etcd for %v using mutator: %v", *etcd_mutation_timeout, mutator)
		default:
			if err := mutator(); err != nil {
				delay := 50 * time.Millisecond
				log.Printf("Failed to mutate etcd using mutator: %v due to: %v. Will retry in: %v", mutator, err, delay)
				time.Sleep(delay)
			} else {
				return
			}
		}
	}
}

func newEtcdClient() (client *etcd.Client) {
	maxConnectRetries := 12
	for maxConnectRetries > 0 {
		if _, _, err := tools.GetEtcdVersion(*etcd_server); err != nil {
			log.Fatalf("Failed to connect to etcd server: %v, error: %v", *etcd_server, err)
			if maxConnectRetries > 0 {
				log.Println("Retrying request after 5 second sleep.")
				time.Sleep(5 * time.Second)
				maxConnectRetries--
			} else {
				return nil
			}
		} else {
			log.Printf("Etcd server found: %v", *etcd_server)
			break
		}
	}
	client = etcd.NewClient([]string{*etcd_server})
	if client == nil {
		return nil
	}
	client.SyncCluster()

	return client
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	config := &kclient.Config{}

	masterHost := os.Getenv("KUBERNETES_RO_SERVICE_HOST")
	if masterHost == "" {
		log.Fatalf("KUBERNETES_RO_SERVICE_HOST is not defined")
	}
	masterPort := os.Getenv("KUBERNETES_RO_SERVICE_PORT")
	if masterPort == "" {
		log.Fatalf("KUBERNETES_RO_SERVICE_PORT is not defined")
	}
	config.Host = fmt.Sprintf("http://%s:%s", masterHost, masterPort)
	log.Printf("Using %s for kubernetes master", config.Host)

	config.Version = "v1beta1"
	log.Printf("Using kubernetes API %s", config.Version)

	return kclient.New(config)
}

func buildNameString(service, namespace, domain string) string {
	return fmt.Sprintf("%s.%s.%s.", service, namespace, domain)
}

func watchOnce(etcdClient *etcd.Client, kubeClient *kclient.Client) {
	// Start the goroutine to produce update events.
	updates := make(chan serviceUpdate)
	startWatching(kubeClient.Services(kapi.NamespaceAll), updates)

	// This loop will break if the channel closes, which is how the
	// goroutine signals an error.
	for ev := range updates {
		if *verbose {
			log.Printf("Received update event: %#v", ev)
		}
		switch ev.Op {
		case SetServices, AddService:
			for i := range ev.Services {
				s := &ev.Services[i]
				name := buildNameString(s.Name, s.Namespace, *domain)
				mutateEtcdOrDie(func() error { return addDNS(name, s, etcdClient) })
			}
		case RemoveService:
			for i := range ev.Services {
				s := &ev.Services[i]
				name := buildNameString(s.Name, s.Namespace, *domain)
				mutateEtcdOrDie(func() error { return removeDNS(name, etcdClient) })
			}
		}
	}
	//TODO: fully resync periodically.
}

func main() {
	flag.Parse()

	etcdClient := newEtcdClient()
	if etcdClient == nil {
		log.Fatal("Failed to create etcd client")
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		log.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	// In case of error, the watch will be aborted.  At that point we just
	// retry.
	for {
		watchOnce(etcdClient, kubeClient)
	}
}

//FIXME: make the below part of the k8s client lib?

// servicesWatcher is capable of listing and watching for changes to services
// across ALL namespaces
type servicesWatcher interface {
	List(label klabels.Selector) (*kapi.ServiceList, error)
	Watch(label, field klabels.Selector, resourceVersion string) (kwatch.Interface, error)
}

type operation int

// These are the available operation types.
const (
	SetServices operation = iota
	AddService
	RemoveService
)

// serviceUpdate describes an operation of services, sent on the channel.
//
// You can add or remove a single service by sending an array of size one with
// Op == AddService|RemoveService.  For setting the state of the system to a given state, just
// set Services as desired and Op to SetServices, which will reset the system
// state to that specified in this operation for this source channel. To remove
// all services, set Services to empty array and Op to SetServices
type serviceUpdate struct {
	Services []kapi.Service
	Op       operation
}

// startWatching launches a goroutine that watches for changes to services.
func startWatching(watcher servicesWatcher, updates chan<- serviceUpdate) {
	serviceVersion := ""
	go watchLoop(watcher, updates, &serviceVersion)
}

// watchLoop loops forever looking for changes to services.  If an error occurs
// it will close the channel and return.
func watchLoop(svcWatcher servicesWatcher, updates chan<- serviceUpdate, resourceVersion *string) {
	defer close(updates)

	if len(*resourceVersion) == 0 {
		services, err := svcWatcher.List(klabels.Everything())
		if err != nil {
			log.Printf("Failed to load services: %v", err)
			return
		}
		*resourceVersion = services.ResourceVersion
		updates <- serviceUpdate{Op: SetServices, Services: services.Items}
	}

	watcher, err := svcWatcher.Watch(klabels.Everything(), klabels.Everything(), *resourceVersion)
	if err != nil {
		log.Printf("Failed to watch for service changes: %v", err)
		return
	}
	defer watcher.Stop()

	ch := watcher.ResultChan()
	for {
		select {
		case event, ok := <-ch:
			if !ok {
				log.Printf("watchLoop channel closed")
				return
			}

			if event.Type == kwatch.Error {
				if status, ok := event.Object.(*kapi.Status); ok {
					log.Printf("Error during watch: %#v", status)
					return
				}
				log.Fatalf("Received unexpected error: %#v", event.Object)
			}

			if service, ok := event.Object.(*kapi.Service); ok {
				sendUpdate(updates, event, service, resourceVersion)
				continue
			}
		}
	}
}

func sendUpdate(updates chan<- serviceUpdate, event kwatch.Event, service *kapi.Service, resourceVersion *string) {
	*resourceVersion = service.ResourceVersion

	switch event.Type {
	case kwatch.Added, kwatch.Modified:
		updates <- serviceUpdate{Op: AddService, Services: []kapi.Service{*service}}
	case kwatch.Deleted:
		updates <- serviceUpdate{Op: RemoveService, Services: []kapi.Service{*service}}
	default:
		log.Fatalf("Unknown event.Type: %v", event.Type)
	}
}
<|endoftext|>