<commit_before>/*
A simple Feed Forward Neural Network can be constructed and trained as follows:

	// set the random seed to 0
	rand.Seed(0)

	// create the XOR representation patter to train the network
	patterns := [][][]float64{
	  {{0, 0}, {0}},
	  {{0, 1}, {1}},
	  {{1, 0}, {1}},
	  {{1, 1}, {0}},
	}

	// instantiate the Feed Forward
	ff := &nn.FeedForward{}

	// initialize the Neural Network;
	// the networks structure will contain:
	// 2 inputs, 2 hidden nodes and 1 output.
	ff.Init(2, 2, 1)

	// train the network using the XOR patterns
	// the training will run for 1000 epochs
	// the learning rate is set to 0.6 and the momentum factor to 0.4
	// use true in the last parameter to receive reports about the learning error
	ff.Train(patterns, 1000, 0.6, 0.4, true)

After running this code the network will be trained and ready to be used.

The network can be tested running using the `Test` method, for instance:

	ff.Test(patterns)

The test operation will print in the console something like:

	[0 0] -> [0.057503945708445]  :  [0]
	[0 1] -> [0.930100635071210]  :  [1]
	[1 0] -> [0.927809966227284]  :  [1]
	[1 1] -> [0.097408795324620]  :  [0]

Where the first values are the inputs, the values after the arrow `->` are the output values from the network and the values after `:` are the expected outputs.

The method `Update` can be used to predict the output given an input, for example:

	inputs := []float64{1, 1}
	ff.Update(inputs)

the output will be a vector with values ranging from `0` to `1`.
*/

package nn

import (
	"fmt"
	"log"
	"math"
)

type FeedForward struct {
	// Number of input, hidden and output nodes
	NInputs, NHiddens, NOutputs int
	// Whether it is regression or not
	Regression bool
	// Activations for nodes
	InputActivations, HiddenActivations, OutputActivations []float64
	// ElmanRNN contexts
	Contexts [][]float64
	// Weights
	InputWeights, OutputWeights [][]float64
	// Last change in weights for momentum
	InputChanges, OutputChanges [][]float64
}

// Initialize the neural network
func (nn *FeedForward) Init(inputs, hiddens, outputs int) {
	nn.NInputs = inputs + 1   // +1 for bias
	nn.NHiddens = hiddens + 1 // +1 for bias
	nn.NOutputs = outputs

	nn.InputActivations = vector(nn.NInputs, 1.0)
	nn.HiddenActivations = vector(nn.NHiddens, 1.0)
	nn.OutputActivations = vector(nn.NOutputs, 1.0)

	nn.InputWeights = matrix(nn.NInputs, nn.NHiddens)
	nn.OutputWeights = matrix(nn.NHiddens, nn.NOutputs)

	for i := 0; i < nn.NInputs; i++ {
		for j := 0; j < nn.NHiddens; j++ {
			nn.InputWeights[i][j] = random(-1, 1)
		}
	}

	for i := 0; i < nn.NHiddens; i++ {
		for j := 0; j < nn.NOutputs; j++ {
			nn.OutputWeights[i][j] = random(-1, 1)
		}
	}

	nn.InputChanges = matrix(nn.NInputs, nn.NHiddens)
	nn.OutputChanges = matrix(nn.NHiddens, nn.NOutputs)
}

func (nn *FeedForward) SetContexts(nContexts int, initValues [][]float64) {
	if initValues == nil {
		initValues = make([][]float64, nContexts)

		for i := 0; i < nContexts; i++ {
			initValues[i] = vector(nn.NHiddens, 0.5)
		}
	}

	nn.Contexts = initValues
}

func (nn *FeedForward) Update(inputs []float64) []float64 {
	if len(inputs) != nn.NInputs-1 {
		log.Fatal("Error: wrong number of inputs")
	}

	for i := 0; i < nn.NInputs-1; i++ {
		nn.InputActivations[i] = inputs[i]
	}

	for i := 0; i < nn.NHiddens-1; i++ {
		var sum float64 = 0.0

		for j := 0; j < nn.NInputs; j++ {
			sum += nn.InputActivations[j] * nn.InputWeights[j][i]
		}

		// compute contexts sum
		for k := 0; k < len(nn.Contexts); k++ {
			for j := 0; j < nn.NHiddens-1; j++ {
				sum += nn.Contexts[k][j]
			}
		}

		nn.HiddenActivations[i] = sigmoid(sum)
	}

	// update the contexts
	if len(nn.Contexts) > 0 {
		for i := len(nn.Contexts) - 1; i > 0; i-- {
			nn.Contexts[i] = nn.Contexts[i-1]
		}
		nn.Contexts[0] = nn.HiddenActivations
	}

	for i := 0; i < nn.NOutputs; i++ {
		var sum float64 = 0.0
		for j := 0; j < nn.NHiddens; j++ {
			sum += nn.HiddenActivations[j] * nn.OutputWeights[j][i]
		}

		nn.OutputActivations[i] = sigmoid(sum)
	}

	return nn.OutputActivations
}

func (nn *FeedForward) BackPropagate(targets []float64, lRate, mFactor float64) float64 {
	if len(targets) != nn.NOutputs {
		log.Fatal("Error: wrong number of target values")
	}

	outputDeltas := vector(nn.NOutputs, 0.0)
	for i := 0; i < nn.NOutputs; i++ {
		outputDeltas[i] = dsigmoid(nn.OutputActivations[i]) * (targets[i] - nn.OutputActivations[i])
	}

	hiddenDeltas := vector(nn.NHiddens, 0.0)
	for i := 0; i < nn.NHiddens; i++ {
		var e float64 = 0.0

		for j := 0; j < nn.NOutputs; j++ {
			e += outputDeltas[j] * nn.OutputWeights[i][j]
		}

		hiddenDeltas[i] = dsigmoid(nn.HiddenActivations[i]) * e
	}

	for i := 0; i < nn.NHiddens; i++ {
		for j := 0; j < nn.NOutputs; j++ {
			change := outputDeltas[j] * nn.HiddenActivations[i]
			nn.OutputWeights[i][j] = nn.OutputWeights[i][j] + lRate*change + mFactor*nn.OutputChanges[i][j]
			nn.OutputChanges[i][j] = change
		}
	}

	for i := 0; i < nn.NInputs; i++ {
		for j := 0; j < nn.NHiddens; j++ {
			change := hiddenDeltas[j] * nn.InputActivations[i]
			nn.InputWeights[i][j] = nn.InputWeights[i][j] + lRate*change + mFactor*nn.InputChanges[i][j]
			nn.InputChanges[i][j] = change
		}
	}

	var e float64 = 0.0

	for i := 0; i < len(targets); i++ {
		e += 0.5 * math.Pow(targets[i]-nn.OutputActivations[i], 2)
	}

	return e
}

func (nn *FeedForward) Train(patterns [][][]float64, iterations int, lRate, mFactor float64, debug bool) []float64 {
	errors := make([]float64, iterations)

	for i := 0; i < iterations; i++ {
		var e float64 = 0.0
		for _, p := range patterns {
			nn.Update(p[0])

			tmp := nn.BackPropagate(p[1], lRate, mFactor)
			e += tmp
		}

		errors[i] = e

		if debug && i%1000 == 0 {
			fmt.Println(i, e)
		}
	}

	return errors
}

func (nn *FeedForward) Test(patterns [][][]float64) {
	for _, p := range patterns {
		fmt.Println(p[0], "->", nn.Update(p[0]), " : ", p[1])
	}
}
<commit_msg>fix doc<commit_after>/*
A simple Feed Forward Neural Network can be constructed and trained as follows:

	// set the random seed to 0
	rand.Seed(0)

	// create the XOR representation patter to train the network
	patterns := [][][]float64{
	  {{0, 0}, {0}},
	  {{0, 1}, {1}},
	  {{1, 0}, {1}},
	  {{1, 1}, {0}},
	}

	// instantiate the Feed Forward
	ff := &nn.FeedForward{}

	// initialize the Neural Network;
	// the networks structure will contain:
	// 2 inputs, 2 hidden nodes and 1 output.
	ff.Init(2, 2, 1)

	// train the network using the XOR patterns
	// the training will run for 1000 epochs
	// the learning rate is set to 0.6 and the momentum factor to 0.4
	// use true in the last parameter to receive reports about the learning error
	ff.Train(patterns, 1000, 0.6, 0.4, true)

After running this code the network will be trained and ready to be used.

The network can be tested running using the `Test` method, for instance:

	ff.Test(patterns)

The test operation will print in the console something like:

	[0 0] -> [0.057503945708445]  :  [0]
	[0 1] -> [0.930100635071210]  :  [1]
	[1 0] -> [0.927809966227284]  :  [1]
	[1 1] -> [0.097408795324620]  :  [0]

Where the first values are the inputs, the values after the arrow `->` are the output values from the network and the values after `:` are the expected outputs.

The method `Update` can be used to predict the output given an input, for example:

	inputs := []float64{1, 1}
	ff.Update(inputs)

the output will be a vector with values ranging from `0` to `1`.
*/
package nn

import (
	"fmt"
	"log"
	"math"
)

type FeedForward struct {
	// Number of input, hidden and output nodes
	NInputs, NHiddens, NOutputs int
	// Whether it is regression or not
	Regression bool
	// Activations for nodes
	InputActivations, HiddenActivations, OutputActivations []float64
	// ElmanRNN contexts
	Contexts [][]float64
	// Weights
	InputWeights, OutputWeights [][]float64
	// Last change in weights for momentum
	InputChanges, OutputChanges [][]float64
}

// Initialize the neural network
func (nn *FeedForward) Init(inputs, hiddens, outputs int) {
	nn.NInputs = inputs + 1   // +1 for bias
	nn.NHiddens = hiddens + 1 // +1 for bias
	nn.NOutputs = outputs

	nn.InputActivations = vector(nn.NInputs, 1.0)
	nn.HiddenActivations = vector(nn.NHiddens, 1.0)
	nn.OutputActivations = vector(nn.NOutputs, 1.0)

	nn.InputWeights = matrix(nn.NInputs, nn.NHiddens)
	nn.OutputWeights = matrix(nn.NHiddens, nn.NOutputs)

	for i := 0; i < nn.NInputs; i++ {
		for j := 0; j < nn.NHiddens; j++ {
			nn.InputWeights[i][j] = random(-1, 1)
		}
	}

	for i := 0; i < nn.NHiddens; i++ {
		for j := 0; j < nn.NOutputs; j++ {
			nn.OutputWeights[i][j] = random(-1, 1)
		}
	}

	nn.InputChanges = matrix(nn.NInputs, nn.NHiddens)
	nn.OutputChanges = matrix(nn.NHiddens, nn.NOutputs)
}

func (nn *FeedForward) SetContexts(nContexts int, initValues [][]float64) {
	if initValues == nil {
		initValues = make([][]float64, nContexts)

		for i := 0; i < nContexts; i++ {
			initValues[i] = vector(nn.NHiddens, 0.5)
		}
	}

	nn.Contexts = initValues
}

func (nn *FeedForward) Update(inputs []float64) []float64 {
	if len(inputs) != nn.NInputs-1 {
		log.Fatal("Error: wrong number of inputs")
	}

	for i := 0; i < nn.NInputs-1; i++ {
		nn.InputActivations[i] = inputs[i]
	}

	for i := 0; i < nn.NHiddens-1; i++ {
		var sum float64 = 0.0

		for j := 0; j < nn.NInputs; j++ {
			sum += nn.InputActivations[j] * nn.InputWeights[j][i]
		}

		// compute contexts sum
		for k := 0; k < len(nn.Contexts); k++ {
			for j := 0; j < nn.NHiddens-1; j++ {
				sum += nn.Contexts[k][j]
			}
		}

		nn.HiddenActivations[i] = sigmoid(sum)
	}

	// update the contexts
	if len(nn.Contexts) > 0 {
		for i := len(nn.Contexts) - 1; i > 0; i-- {
			nn.Contexts[i] = nn.Contexts[i-1]
		}
		nn.Contexts[0] = nn.HiddenActivations
	}

	for i := 0; i < nn.NOutputs; i++ {
		var sum float64 = 0.0
		for j := 0; j < nn.NHiddens; j++ {
			sum += nn.HiddenActivations[j] * nn.OutputWeights[j][i]
		}

		nn.OutputActivations[i] = sigmoid(sum)
	}

	return nn.OutputActivations
}

func (nn *FeedForward) BackPropagate(targets []float64, lRate, mFactor float64) float64 {
	if len(targets) != nn.NOutputs {
		log.Fatal("Error: wrong number of target values")
	}

	outputDeltas := vector(nn.NOutputs, 0.0)
	for i := 0; i < nn.NOutputs; i++ {
		outputDeltas[i] = dsigmoid(nn.OutputActivations[i]) * (targets[i] - nn.OutputActivations[i])
	}

	hiddenDeltas := vector(nn.NHiddens, 0.0)
	for i := 0; i < nn.NHiddens; i++ {
		var e float64 = 0.0

		for j := 0; j < nn.NOutputs; j++ {
			e += outputDeltas[j] * nn.OutputWeights[i][j]
		}

		hiddenDeltas[i] = dsigmoid(nn.HiddenActivations[i]) * e
	}

	for i := 0; i < nn.NHiddens; i++ {
		for j := 0; j < nn.NOutputs; j++ {
			change := outputDeltas[j] * nn.HiddenActivations[i]
			nn.OutputWeights[i][j] = nn.OutputWeights[i][j] + lRate*change + mFactor*nn.OutputChanges[i][j]
			nn.OutputChanges[i][j] = change
		}
	}

	for i := 0; i < nn.NInputs; i++ {
		for j := 0; j < nn.NHiddens; j++ {
			change := hiddenDeltas[j] * nn.InputActivations[i]
			nn.InputWeights[i][j] = nn.InputWeights[i][j] + lRate*change + mFactor*nn.InputChanges[i][j]
			nn.InputChanges[i][j] = change
		}
	}

	var e float64 = 0.0

	for i := 0; i < len(targets); i++ {
		e += 0.5 * math.Pow(targets[i]-nn.OutputActivations[i], 2)
	}

	return e
}

func (nn *FeedForward) Train(patterns [][][]float64, iterations int, lRate, mFactor float64, debug bool) []float64 {
	errors := make([]float64, iterations)

	for i := 0; i < iterations; i++ {
		var e float64 = 0.0
		for _, p := range patterns {
			nn.Update(p[0])

			tmp := nn.BackPropagate(p[1], lRate, mFactor)
			e += tmp
		}

		errors[i] = e

		if debug && i%1000 == 0 {
			fmt.Println(i, e)
		}
	}

	return errors
}

func (nn *FeedForward) Test(patterns [][][]float64) {
	for _, p := range patterns {
		fmt.Println(p[0], "->", nn.Update(p[0]), " : ", p[1])
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/amahi/spdy"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"time"
)

func handle(err error) {
	if err != nil {
		panic(err)
	}
}

const HOST_PORT = "localhost:1444"

func main() {

	root := flag.String("r", "./testdata", "root of the directory to serve")
	spdy_debug := flag.Bool("s", false, "enable SPDY debug output")
	flag.Parse()

	if *spdy_debug {
		// enable spdy debug messages
		spdy.EnableDebug()
	}

	for {
		const SLEEP_RETRY = 5
		var conn *tls.Conn
		var err error
		for i := 0; i < 10; i++ {
			// connect to P.
			conn, err = tls.Dial("tcp", HOST_PORT, &tls.Config{InsecureSkipVerify: true})
			if err != nil {
				time.Sleep(100 * time.Millisecond)
			} else {
				break
			}
		}
		if conn == nil {
			log.Println("Failed to connect. Waiting", SLEEP_RETRY, "seconds.")
			time.Sleep(SLEEP_RETRY * time.Second)
			continue
		}

		// build the request
		buf := new(bytes.Buffer)
		_, err = buf.WriteString("Hello from C")
		handle(err)
		req, err := http.NewRequest("PUT", "https://"+HOST_PORT, buf)
		handle(err)

		// make the client connection
		client := httputil.NewClientConn(conn, nil)
		res, err := client.Do(req)
		if err != nil {
			log.Println("Error: Failed to make connection to P:", err)
			continue
		}
		buf.Reset()
		_, err = io.Copy(buf, res.Body)
		handle(err)
		fmt.Printf("%q from P: %q.\n", res.Status, buf.String())

		c, _ := client.Hijack()
		conn = c.(*tls.Conn)
		server := new(http.Server)
		server.Handler = http.FileServer(http.Dir(*root))
		session := spdy.NewServerSession(conn, server)
		session.Serve()
	}
}
<commit_msg>POST support properly tested<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/amahi/spdy"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"time"
)

func handle(err error) {
	if err != nil {
		panic(err)
	}
}

type handler struct {
        data []byte
        rt string
}
func (h *handler) ServeHTTP(rw http.ResponseWriter,rq *http.Request) {
        if rq.Body!=nil {
                h.data = make([]byte, int(rq.ContentLength))
                _,err := rq.Body.(io.Reader).Read(h.data)
                if err != nil {
                        fmt.Println(err)
                }
                filename := "/tmp/postdat"
                f, err := os.Create(filename)
                if err != nil {
                        fmt.Println(err)
                }
                n, err := f.Write(h.data)
                if err != nil {
                        fmt.Println(n, err)
                }
                f.Close()
        }
        fileserver := http.FileServer(http.Dir(h.rt))
        fileserver.ServeHTTP(rw, rq)
}

const HOST_PORT = "localhost:1444"

func main() {

	root := flag.String("r", "./testdata", "root of the directory to serve")
	spdy_debug := flag.Bool("s", false, "enable SPDY debug output")
	flag.Parse()

	if *spdy_debug {
		// enable spdy debug messages
		spdy.EnableDebug()
	}

	for {
		const SLEEP_RETRY = 5
		var conn *tls.Conn
		var err error
		for i := 0; i < 10; i++ {
			// connect to P.
			conn, err = tls.Dial("tcp", HOST_PORT, &tls.Config{InsecureSkipVerify: true})
			if err != nil {
				time.Sleep(100 * time.Millisecond)
			} else {
				break
			}
		}
		if conn == nil {
			log.Println("Failed to connect. Waiting", SLEEP_RETRY, "seconds.")
			time.Sleep(SLEEP_RETRY * time.Second)
			continue
		}

		// build the request
		buf := new(bytes.Buffer)
		_, err = buf.WriteString("Hello from C")
		handle(err)
		req, err := http.NewRequest("PUT", "https://"+HOST_PORT, buf)
		handle(err)

		// make the client connection
		client := httputil.NewClientConn(conn, nil)
		res, err := client.Do(req)
		if err != nil {
			log.Println("Error: Failed to make connection to P:", err)
			continue
		}
		buf.Reset()
		_, err = io.Copy(buf, res.Body)
		handle(err)
		fmt.Printf("%q from P: %q.\n", res.Status, buf.String())

		c, _ := client.Hijack()
		conn = c.(*tls.Conn)
		server := new(http.Server)
		server.Handler = &handler{data:nil,rt:*root}
		//http.FileServer(http.Dir(*root))
		session := spdy.NewServerSession(conn, server)
		session.Serve()
	}
}
<|endoftext|>


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

package session

import (
	"encoding/json"
	"github.com/realglobe-Inc/go-lib/erro"
)

// OpenID Connect Core 1.0 Section 5.5 を参照。

// 認証リクエストの claims.id_token や claims.userinfo パラメータの要素。
type ClaimEntry struct {
	// essential
	ess bool
	// value
	val interface{}
	// values
	vals []interface{}

	// 言語タグ。
	lang string
}

func (this *ClaimEntry) Essential() bool {
	return this.ess
}

func (this *ClaimEntry) Value() interface{} {
	return this.val
}

func (this *ClaimEntry) Values() []interface{} {
	return this.vals
}

func (this *ClaimEntry) Language() string {
	return this.lang
}

func (this *ClaimEntry) setLanguage(lang string) {
	this.lang = lang
}

//  {
//      "essential": <必須か>,
//      "value": <指定値>,
//      "values": [
//          <候補値>,
//          ....
//      ]
//  }
func (this *ClaimEntry) MarshalJSON() (data []byte, err error) {
	return json.Marshal(map[string]interface{}{
		"essential": this.ess,
		"value":     this.val,
		"values":    this.vals,
	})
}

func (this *ClaimEntry) UnmarshalJSON(data []byte) error {
	var buff struct {
		Ess  bool          `json:"essential"`
		Val  interface{}   `json:"value"`
		Vals []interface{} `json:"values"`
	}
	if err := json.Unmarshal(data, &buff); err != nil {
		return erro.Wrap(err)
	}
	this.ess = buff.Ess
	this.val = buff.Val
	this.vals = buff.Vals
	return nil
}
<commit_msg>テスト用に要求クレームを作成する関数を追加<commit_after>// Copyright 2015 realglobe, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package session

import (
	"encoding/json"
	"github.com/realglobe-Inc/go-lib/erro"
)

// OpenID Connect Core 1.0 Section 5.5 を参照。

// 認証リクエストの claims.id_token や claims.userinfo パラメータの要素。
type ClaimEntry struct {
	// essential
	ess bool
	// value
	val interface{}
	// values
	vals []interface{}

	// 言語タグ。
	lang string
}

// 主にテスト用。
func NewClaimEntry(ess bool, val interface{}, vals []interface{}, lang string) *ClaimEntry {
	return &ClaimEntry{
		ess:  ess,
		val:  val,
		vals: vals,
		lang: lang,
	}
}

func (this *ClaimEntry) Essential() bool {
	return this.ess
}

func (this *ClaimEntry) Value() interface{} {
	return this.val
}

func (this *ClaimEntry) Values() []interface{} {
	return this.vals
}

func (this *ClaimEntry) Language() string {
	return this.lang
}

func (this *ClaimEntry) setLanguage(lang string) {
	this.lang = lang
}

//  {
//      "essential": <必須か>,
//      "value": <指定値>,
//      "values": [
//          <候補値>,
//          ....
//      ]
//  }
func (this *ClaimEntry) MarshalJSON() (data []byte, err error) {
	return json.Marshal(map[string]interface{}{
		"essential": this.ess,
		"value":     this.val,
		"values":    this.vals,
	})
}

func (this *ClaimEntry) UnmarshalJSON(data []byte) error {
	var buff struct {
		Ess  bool          `json:"essential"`
		Val  interface{}   `json:"value"`
		Vals []interface{} `json:"values"`
	}
	if err := json.Unmarshal(data, &buff); err != nil {
		return erro.Wrap(err)
	}
	this.ess = buff.Ess
	this.val = buff.Val
	this.vals = buff.Vals
	return nil
}
<|endoftext|>


<commit_before>package xfer // import "github.com/docker/docker/distribution/xfer"

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/docker/image"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/progress"
	digest "github.com/opencontainers/go-digest"
)

const maxDownloadConcurrency = 3

type mockLayer struct {
	layerData bytes.Buffer
	diffID    layer.DiffID
	chainID   layer.ChainID
	parent    layer.Layer
	os        string
}

func (ml *mockLayer) TarStream() (io.ReadCloser, error) {
	return ioutil.NopCloser(bytes.NewBuffer(ml.layerData.Bytes())), nil
}

func (ml *mockLayer) TarStreamFrom(layer.ChainID) (io.ReadCloser, error) {
	return nil, fmt.Errorf("not implemented")
}

func (ml *mockLayer) ChainID() layer.ChainID {
	return ml.chainID
}

func (ml *mockLayer) DiffID() layer.DiffID {
	return ml.diffID
}

func (ml *mockLayer) Parent() layer.Layer {
	return ml.parent
}

func (ml *mockLayer) Size() (size int64, err error) {
	return 0, nil
}

func (ml *mockLayer) DiffSize() (size int64, err error) {
	return 0, nil
}

func (ml *mockLayer) Metadata() (map[string]string, error) {
	return make(map[string]string), nil
}

type mockLayerStore struct {
	layers map[layer.ChainID]*mockLayer
}

func createChainIDFromParent(parent layer.ChainID, dgsts ...layer.DiffID) layer.ChainID {
	if len(dgsts) == 0 {
		return parent
	}
	if parent == "" {
		return createChainIDFromParent(layer.ChainID(dgsts[0]), dgsts[1:]...)
	}
	// H = "H(n-1) SHA256(n)"
	dgst := digest.FromBytes([]byte(string(parent) + " " + string(dgsts[0])))
	return createChainIDFromParent(layer.ChainID(dgst), dgsts[1:]...)
}

func (ls *mockLayerStore) Map() map[layer.ChainID]layer.Layer {
	layers := map[layer.ChainID]layer.Layer{}

	for k, v := range ls.layers {
		layers[k] = v
	}

	return layers
}

func (ls *mockLayerStore) Register(reader io.Reader, parentID layer.ChainID) (layer.Layer, error) {
	return ls.RegisterWithDescriptor(reader, parentID, distribution.Descriptor{})
}

func (ls *mockLayerStore) RegisterWithDescriptor(reader io.Reader, parentID layer.ChainID, _ distribution.Descriptor) (layer.Layer, error) {
	var (
		parent layer.Layer
		err    error
	)

	if parentID != "" {
		parent, err = ls.Get(parentID)
		if err != nil {
			return nil, err
		}
	}

	l := &mockLayer{parent: parent}
	_, err = l.layerData.ReadFrom(reader)
	if err != nil {
		return nil, err
	}
	l.diffID = layer.DiffID(digest.FromBytes(l.layerData.Bytes()))
	l.chainID = createChainIDFromParent(parentID, l.diffID)

	ls.layers[l.chainID] = l
	return l, nil
}

func (ls *mockLayerStore) Get(chainID layer.ChainID) (layer.Layer, error) {
	l, ok := ls.layers[chainID]
	if !ok {
		return nil, layer.ErrLayerDoesNotExist
	}
	return l, nil
}

func (ls *mockLayerStore) Release(l layer.Layer) ([]layer.Metadata, error) {
	return []layer.Metadata{}, nil
}
func (ls *mockLayerStore) CreateRWLayer(string, layer.ChainID, *layer.CreateRWLayerOpts) (layer.RWLayer, error) {
	return nil, errors.New("not implemented")
}

func (ls *mockLayerStore) GetRWLayer(string) (layer.RWLayer, error) {
	return nil, errors.New("not implemented")
}

func (ls *mockLayerStore) ReleaseRWLayer(layer.RWLayer) ([]layer.Metadata, error) {
	return nil, errors.New("not implemented")
}
func (ls *mockLayerStore) GetMountID(string) (string, error) {
	return "", errors.New("not implemented")
}

func (ls *mockLayerStore) Cleanup() error {
	return nil
}

func (ls *mockLayerStore) DriverStatus() [][2]string {
	return [][2]string{}
}

func (ls *mockLayerStore) DriverName() string {
	return "mock"
}

type mockDownloadDescriptor struct {
	currentDownloads *int32
	id               string
	diffID           layer.DiffID
	registeredDiffID layer.DiffID
	expectedDiffID   layer.DiffID
	simulateRetries  int
}

// Key returns the key used to deduplicate downloads.
func (d *mockDownloadDescriptor) Key() string {
	return d.id
}

// ID returns the ID for display purposes.
func (d *mockDownloadDescriptor) ID() string {
	return d.id
}

// DiffID should return the DiffID for this layer, or an error
// if it is unknown (for example, if it has not been downloaded
// before).
func (d *mockDownloadDescriptor) DiffID() (layer.DiffID, error) {
	if d.diffID != "" {
		return d.diffID, nil
	}
	return "", errors.New("no diffID available")
}

func (d *mockDownloadDescriptor) Registered(diffID layer.DiffID) {
	d.registeredDiffID = diffID
}

func (d *mockDownloadDescriptor) mockTarStream() io.ReadCloser {
	// The mock implementation returns the ID repeated 5 times as a tar
	// stream instead of actual tar data. The data is ignored except for
	// computing IDs.
	return ioutil.NopCloser(bytes.NewBuffer([]byte(d.id + d.id + d.id + d.id + d.id)))
}

// Download is called to perform the download.
func (d *mockDownloadDescriptor) Download(ctx context.Context, progressOutput progress.Output) (io.ReadCloser, int64, error) {
	if d.currentDownloads != nil {
		defer atomic.AddInt32(d.currentDownloads, -1)

		if atomic.AddInt32(d.currentDownloads, 1) > maxDownloadConcurrency {
			return nil, 0, errors.New("concurrency limit exceeded")
		}
	}

	// Sleep a bit to simulate a time-consuming download.
	for i := int64(0); i <= 10; i++ {
		select {
		case <-ctx.Done():
			return nil, 0, ctx.Err()
		case <-time.After(10 * time.Millisecond):
			progressOutput.WriteProgress(progress.Progress{ID: d.ID(), Action: "Downloading", Current: i, Total: 10})
		}
	}

	if d.simulateRetries != 0 {
		d.simulateRetries--
		return nil, 0, errors.New("simulating retry")
	}

	return d.mockTarStream(), 0, nil
}

func (d *mockDownloadDescriptor) Close() {
}

func downloadDescriptors(currentDownloads *int32) []DownloadDescriptor {
	return []DownloadDescriptor{
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id1",
			expectedDiffID:   layer.DiffID("sha256:68e2c75dc5c78ea9240689c60d7599766c213ae210434c53af18470ae8c53ec1"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id2",
			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id3",
			expectedDiffID:   layer.DiffID("sha256:58745a8bbd669c25213e9de578c4da5c8ee1c836b3581432c2b50e38a6753300"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id2",
			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id4",
			expectedDiffID:   layer.DiffID("sha256:0dfb5b9577716cc173e95af7c10289322c29a6453a1718addc00c0c5b1330936"),
			simulateRetries:  1,
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id5",
			expectedDiffID:   layer.DiffID("sha256:0a5f25fa1acbc647f6112a6276735d0fa01e4ee2aa7ec33015e337350e1ea23d"),
		},
	}
}

func TestSuccessfulDownload(t *testing.T) {
	// TODO Windows: Fix this unit text
	if runtime.GOOS == "windows" {
		t.Skip("Needs fixing on Windows")
	}

	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
	lsMap := make(map[string]layer.Store)
	lsMap[runtime.GOOS] = layerStore
	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })

	progressChan := make(chan progress.Progress)
	progressDone := make(chan struct{})
	receivedProgress := make(map[string]progress.Progress)

	go func() {
		for p := range progressChan {
			receivedProgress[p.ID] = p
		}
		close(progressDone)
	}()

	var currentDownloads int32
	descriptors := downloadDescriptors(&currentDownloads)

	firstDescriptor := descriptors[0].(*mockDownloadDescriptor)

	// Pre-register the first layer to simulate an already-existing layer
	l, err := layerStore.Register(firstDescriptor.mockTarStream(), "")
	if err != nil {
		t.Fatal(err)
	}
	firstDescriptor.diffID = l.DiffID()

	rootFS, releaseFunc, err := ldm.Download(context.Background(), *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
	if err != nil {
		t.Fatalf("download error: %v", err)
	}

	releaseFunc()

	close(progressChan)
	<-progressDone

	if len(rootFS.DiffIDs) != len(descriptors) {
		t.Fatal("got wrong number of diffIDs in rootfs")
	}

	for i, d := range descriptors {
		descriptor := d.(*mockDownloadDescriptor)

		if descriptor.diffID != "" {
			if receivedProgress[d.ID()].Action != "Already exists" {
				t.Fatalf("did not get 'Already exists' message for %v", d.ID())
			}
		} else if receivedProgress[d.ID()].Action != "Pull complete" {
			t.Fatalf("did not get 'Pull complete' message for %v", d.ID())
		}

		if rootFS.DiffIDs[i] != descriptor.expectedDiffID {
			t.Fatalf("rootFS item %d has the wrong diffID (expected: %v got: %v)", i, descriptor.expectedDiffID, rootFS.DiffIDs[i])
		}

		if descriptor.diffID == "" && descriptor.registeredDiffID != rootFS.DiffIDs[i] {
			t.Fatal("diffID mismatch between rootFS and Registered callback")
		}
	}
}

func TestCancelledDownload(t *testing.T) {
	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
	lsMap := make(map[string]layer.Store)
	lsMap[runtime.GOOS] = layerStore
	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })
	progressChan := make(chan progress.Progress)
	progressDone := make(chan struct{})

	go func() {
		for range progressChan {
		}
		close(progressDone)
	}()

	ctx, cancel := context.WithCancel(context.Background())

	go func() {
		<-time.After(time.Millisecond)
		cancel()
	}()

	descriptors := downloadDescriptors(nil)
	_, _, err := ldm.Download(ctx, *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
	if err != context.Canceled {
		t.Fatal("expected download to be cancelled")
	}

	close(progressChan)
	<-progressDone
}
<commit_msg>distribution/xfer/download_test: rm unused field<commit_after>package xfer // import "github.com/docker/docker/distribution/xfer"

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/docker/image"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/progress"
	digest "github.com/opencontainers/go-digest"
)

const maxDownloadConcurrency = 3

type mockLayer struct {
	layerData bytes.Buffer
	diffID    layer.DiffID
	chainID   layer.ChainID
	parent    layer.Layer
}

func (ml *mockLayer) TarStream() (io.ReadCloser, error) {
	return ioutil.NopCloser(bytes.NewBuffer(ml.layerData.Bytes())), nil
}

func (ml *mockLayer) TarStreamFrom(layer.ChainID) (io.ReadCloser, error) {
	return nil, fmt.Errorf("not implemented")
}

func (ml *mockLayer) ChainID() layer.ChainID {
	return ml.chainID
}

func (ml *mockLayer) DiffID() layer.DiffID {
	return ml.diffID
}

func (ml *mockLayer) Parent() layer.Layer {
	return ml.parent
}

func (ml *mockLayer) Size() (size int64, err error) {
	return 0, nil
}

func (ml *mockLayer) DiffSize() (size int64, err error) {
	return 0, nil
}

func (ml *mockLayer) Metadata() (map[string]string, error) {
	return make(map[string]string), nil
}

type mockLayerStore struct {
	layers map[layer.ChainID]*mockLayer
}

func createChainIDFromParent(parent layer.ChainID, dgsts ...layer.DiffID) layer.ChainID {
	if len(dgsts) == 0 {
		return parent
	}
	if parent == "" {
		return createChainIDFromParent(layer.ChainID(dgsts[0]), dgsts[1:]...)
	}
	// H = "H(n-1) SHA256(n)"
	dgst := digest.FromBytes([]byte(string(parent) + " " + string(dgsts[0])))
	return createChainIDFromParent(layer.ChainID(dgst), dgsts[1:]...)
}

func (ls *mockLayerStore) Map() map[layer.ChainID]layer.Layer {
	layers := map[layer.ChainID]layer.Layer{}

	for k, v := range ls.layers {
		layers[k] = v
	}

	return layers
}

func (ls *mockLayerStore) Register(reader io.Reader, parentID layer.ChainID) (layer.Layer, error) {
	return ls.RegisterWithDescriptor(reader, parentID, distribution.Descriptor{})
}

func (ls *mockLayerStore) RegisterWithDescriptor(reader io.Reader, parentID layer.ChainID, _ distribution.Descriptor) (layer.Layer, error) {
	var (
		parent layer.Layer
		err    error
	)

	if parentID != "" {
		parent, err = ls.Get(parentID)
		if err != nil {
			return nil, err
		}
	}

	l := &mockLayer{parent: parent}
	_, err = l.layerData.ReadFrom(reader)
	if err != nil {
		return nil, err
	}
	l.diffID = layer.DiffID(digest.FromBytes(l.layerData.Bytes()))
	l.chainID = createChainIDFromParent(parentID, l.diffID)

	ls.layers[l.chainID] = l
	return l, nil
}

func (ls *mockLayerStore) Get(chainID layer.ChainID) (layer.Layer, error) {
	l, ok := ls.layers[chainID]
	if !ok {
		return nil, layer.ErrLayerDoesNotExist
	}
	return l, nil
}

func (ls *mockLayerStore) Release(l layer.Layer) ([]layer.Metadata, error) {
	return []layer.Metadata{}, nil
}
func (ls *mockLayerStore) CreateRWLayer(string, layer.ChainID, *layer.CreateRWLayerOpts) (layer.RWLayer, error) {
	return nil, errors.New("not implemented")
}

func (ls *mockLayerStore) GetRWLayer(string) (layer.RWLayer, error) {
	return nil, errors.New("not implemented")
}

func (ls *mockLayerStore) ReleaseRWLayer(layer.RWLayer) ([]layer.Metadata, error) {
	return nil, errors.New("not implemented")
}
func (ls *mockLayerStore) GetMountID(string) (string, error) {
	return "", errors.New("not implemented")
}

func (ls *mockLayerStore) Cleanup() error {
	return nil
}

func (ls *mockLayerStore) DriverStatus() [][2]string {
	return [][2]string{}
}

func (ls *mockLayerStore) DriverName() string {
	return "mock"
}

type mockDownloadDescriptor struct {
	currentDownloads *int32
	id               string
	diffID           layer.DiffID
	registeredDiffID layer.DiffID
	expectedDiffID   layer.DiffID
	simulateRetries  int
}

// Key returns the key used to deduplicate downloads.
func (d *mockDownloadDescriptor) Key() string {
	return d.id
}

// ID returns the ID for display purposes.
func (d *mockDownloadDescriptor) ID() string {
	return d.id
}

// DiffID should return the DiffID for this layer, or an error
// if it is unknown (for example, if it has not been downloaded
// before).
func (d *mockDownloadDescriptor) DiffID() (layer.DiffID, error) {
	if d.diffID != "" {
		return d.diffID, nil
	}
	return "", errors.New("no diffID available")
}

func (d *mockDownloadDescriptor) Registered(diffID layer.DiffID) {
	d.registeredDiffID = diffID
}

func (d *mockDownloadDescriptor) mockTarStream() io.ReadCloser {
	// The mock implementation returns the ID repeated 5 times as a tar
	// stream instead of actual tar data. The data is ignored except for
	// computing IDs.
	return ioutil.NopCloser(bytes.NewBuffer([]byte(d.id + d.id + d.id + d.id + d.id)))
}

// Download is called to perform the download.
func (d *mockDownloadDescriptor) Download(ctx context.Context, progressOutput progress.Output) (io.ReadCloser, int64, error) {
	if d.currentDownloads != nil {
		defer atomic.AddInt32(d.currentDownloads, -1)

		if atomic.AddInt32(d.currentDownloads, 1) > maxDownloadConcurrency {
			return nil, 0, errors.New("concurrency limit exceeded")
		}
	}

	// Sleep a bit to simulate a time-consuming download.
	for i := int64(0); i <= 10; i++ {
		select {
		case <-ctx.Done():
			return nil, 0, ctx.Err()
		case <-time.After(10 * time.Millisecond):
			progressOutput.WriteProgress(progress.Progress{ID: d.ID(), Action: "Downloading", Current: i, Total: 10})
		}
	}

	if d.simulateRetries != 0 {
		d.simulateRetries--
		return nil, 0, errors.New("simulating retry")
	}

	return d.mockTarStream(), 0, nil
}

func (d *mockDownloadDescriptor) Close() {
}

func downloadDescriptors(currentDownloads *int32) []DownloadDescriptor {
	return []DownloadDescriptor{
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id1",
			expectedDiffID:   layer.DiffID("sha256:68e2c75dc5c78ea9240689c60d7599766c213ae210434c53af18470ae8c53ec1"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id2",
			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id3",
			expectedDiffID:   layer.DiffID("sha256:58745a8bbd669c25213e9de578c4da5c8ee1c836b3581432c2b50e38a6753300"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id2",
			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id4",
			expectedDiffID:   layer.DiffID("sha256:0dfb5b9577716cc173e95af7c10289322c29a6453a1718addc00c0c5b1330936"),
			simulateRetries:  1,
		},
		&mockDownloadDescriptor{
			currentDownloads: currentDownloads,
			id:               "id5",
			expectedDiffID:   layer.DiffID("sha256:0a5f25fa1acbc647f6112a6276735d0fa01e4ee2aa7ec33015e337350e1ea23d"),
		},
	}
}

func TestSuccessfulDownload(t *testing.T) {
	// TODO Windows: Fix this unit text
	if runtime.GOOS == "windows" {
		t.Skip("Needs fixing on Windows")
	}

	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
	lsMap := make(map[string]layer.Store)
	lsMap[runtime.GOOS] = layerStore
	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })

	progressChan := make(chan progress.Progress)
	progressDone := make(chan struct{})
	receivedProgress := make(map[string]progress.Progress)

	go func() {
		for p := range progressChan {
			receivedProgress[p.ID] = p
		}
		close(progressDone)
	}()

	var currentDownloads int32
	descriptors := downloadDescriptors(&currentDownloads)

	firstDescriptor := descriptors[0].(*mockDownloadDescriptor)

	// Pre-register the first layer to simulate an already-existing layer
	l, err := layerStore.Register(firstDescriptor.mockTarStream(), "")
	if err != nil {
		t.Fatal(err)
	}
	firstDescriptor.diffID = l.DiffID()

	rootFS, releaseFunc, err := ldm.Download(context.Background(), *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
	if err != nil {
		t.Fatalf("download error: %v", err)
	}

	releaseFunc()

	close(progressChan)
	<-progressDone

	if len(rootFS.DiffIDs) != len(descriptors) {
		t.Fatal("got wrong number of diffIDs in rootfs")
	}

	for i, d := range descriptors {
		descriptor := d.(*mockDownloadDescriptor)

		if descriptor.diffID != "" {
			if receivedProgress[d.ID()].Action != "Already exists" {
				t.Fatalf("did not get 'Already exists' message for %v", d.ID())
			}
		} else if receivedProgress[d.ID()].Action != "Pull complete" {
			t.Fatalf("did not get 'Pull complete' message for %v", d.ID())
		}

		if rootFS.DiffIDs[i] != descriptor.expectedDiffID {
			t.Fatalf("rootFS item %d has the wrong diffID (expected: %v got: %v)", i, descriptor.expectedDiffID, rootFS.DiffIDs[i])
		}

		if descriptor.diffID == "" && descriptor.registeredDiffID != rootFS.DiffIDs[i] {
			t.Fatal("diffID mismatch between rootFS and Registered callback")
		}
	}
}

func TestCancelledDownload(t *testing.T) {
	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
	lsMap := make(map[string]layer.Store)
	lsMap[runtime.GOOS] = layerStore
	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })
	progressChan := make(chan progress.Progress)
	progressDone := make(chan struct{})

	go func() {
		for range progressChan {
		}
		close(progressDone)
	}()

	ctx, cancel := context.WithCancel(context.Background())

	go func() {
		<-time.After(time.Millisecond)
		cancel()
	}()

	descriptors := downloadDescriptors(nil)
	_, _, err := ldm.Download(ctx, *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
	if err != context.Canceled {
		t.Fatal("expected download to be cancelled")
	}

	close(progressChan)
	<-progressDone
}
<|endoftext|>


<commit_before>package tests

import (
	"database/sql"
	"fmt"
	"github.com/deis/deis/tests/dockercli"
	"github.com/deis/deis/tests/mock"
	"github.com/deis/deis/tests/utils"
	"github.com/lib/pq"
	"testing"
	"time"
)

func OpenDeisDatabase(t *testing.T, host string, port string) *sql.DB {
	db, err := sql.Open("postgres", "postgres://deis:changeme123@"+host+":"+port+"/deis?sslmode=disable")
	if err != nil {
		t.Fatal(err)
	}
	WaitForDatabase(t, db)
	return db
}

func WaitForDatabase(t *testing.T, db *sql.DB) {
	fmt.Printf("--- Waiting for pg to be ready")
	for {
		_, err := db.Query("select 1")
		if err, ok := err.(*pq.Error); ok {
			if err.Code.Name() == "cannot_connect_now" {
				fmt.Printf(".")
				time.Sleep(500 * time.Millisecond)
				continue
			}
			t.Fatal(err)
		}
		fmt.Printf("\n")
		break
	}
}

func TryTableSelect(t *testing.T, db *sql.DB, tableName string, expectFailure bool) {
	_, err := db.Query("select * from " + tableName)

	if expectFailure {
		if err == nil {
			t.Fatal("The table should not exist")
		}
	} else {
		if err != nil {
			t.Fatal(err)
		}
	}
}

func execSql(t *testing.T, db *sql.DB, q string) {
	_, err := db.Query(q)
	if err != nil {
		t.Fatal(err)
	}
}

func TestDatabaseRecovery(t *testing.T) {
	var err error
	tag, etcdPort := utils.BuildTag(), utils.RandomPort()
	cli, stdout, _ := dockercli.NewClient()
	imageName := utils.ImagePrefix() + "database" + ":" + tag

	// start etcd container
	etcdName := "deis-etcd-" + tag
	dockercli.RunTestEtcd(t, etcdName, etcdPort)
	defer cli.CmdRm("-f", etcdName)

	// run mock ceph containers
	cephName := "deis-ceph-" + tag
	mock.RunMockCeph(t, cephName, cli, etcdPort)
	defer cli.CmdRm("-f", cephName)

	// create volumes
	databaseVolumeA := "deis-database-data-a-" + tag
	databaseVolumeB := "deis-database-data-b-" + tag
	defer cli.CmdRm("-f", databaseVolumeA)
	defer cli.CmdRm("-f", databaseVolumeB)
	go func() {
		fmt.Printf("--- Creating Volume A\n")
		_ = cli.CmdRm("-f", "-v", databaseVolumeA)
		dockercli.CreateVolume(t, cli, databaseVolumeA, "/var/lib/postgresql")

		fmt.Printf("--- Creating Volume B\n")

		_ = cli.CmdRm("-f", databaseVolumeB)
		dockercli.CreateVolume(t, cli, databaseVolumeB, "/var/lib/postgresql")
	}()
	dockercli.WaitForLine(t, stdout, databaseVolumeB, true)

	// setup database container start/stop routines
	host, port := utils.HostAddress(), utils.RandomPort()
	fmt.Printf("--- Run deis/database:%s at %s:%s\n", tag, host, port)
	name := "deis-database-" + tag
	defer cli.CmdRm("-f", name)
	startDatabase := func(volumeName string) {
		_ = cli.CmdRm("-f", name)
		err = dockercli.RunContainer(cli,
			"--name", name,
			"--volumes-from", volumeName,
			"--rm",
			"-p", port+":5432",
			"-e", "EXTERNAL_PORT="+port,
			"-e", "HOST="+host,
			"-e", "ETCD_PORT="+etcdPort,
			"-e", "ETCD_TTL=2",
			"-e", "BACKUP_FREQUENCY=0",
			"-e", "BACKUPS_TO_RETAIN=100",
			imageName)
	}

	stopDatabase := func() {
		fmt.Print("--- Stopping data-database... ")
		if err = stdout.Close(); err != nil {
			t.Fatal("Failed to closeStdout")
		}
		_ = cli.CmdStop(name)
		fmt.Println("Done")
	}

	//ACTION

	//STEP 1: start db with volume A and wait for init to complete
	fmt.Print("--- Starting database with Volume A... ")
	go startDatabase(databaseVolumeA)
	dockercli.WaitForLine(t, stdout, "database: postgres is running...", true)
	fmt.Println("Done")

	db := OpenDeisDatabase(t, host, port)
	TryTableSelect(t, db, "api_foo", true)

	stopDatabase()

	//STEP 2a: start db with volume B, wait for init and create the table
	cli, stdout, _ = dockercli.NewClient()
	fmt.Printf("--- Starting database with Volume B... ")
	go startDatabase(databaseVolumeB)
	dockercli.WaitForLine(t, stdout, "database: postgres is running...", true)
	fmt.Println("Done")

	db = OpenDeisDatabase(t, host, port)
	TryTableSelect(t, db, "api_foo", true)

	fmt.Println("--- Creating the table")
	execSql(t, db, "create table api_foo(t text)")

	//STEP 2b: make sure we observed full backup cycle after forced checkpoint
	fmt.Print("--- Waiting for the change to be backed up... ")
	dockercli.WaitForLine(t, stdout, "database: performing a backup...", true)
	dockercli.WaitForLine(t, stdout, "database: backup has been completed.", true)
	fmt.Println("Done")

	stopDatabase()

	//STEP 3: start db with volume A again and assert table existence
	cli, stdout, _ = dockercli.NewClient()
	fmt.Printf("--- Starting database with Volume A again... ")
	go startDatabase(databaseVolumeA)
	dockercli.WaitForLine(t, stdout, "database: postgres is running...", true)
	fmt.Println("Done")

	db = OpenDeisDatabase(t, host, port)
	TryTableSelect(t, db, "api_foo", false)

}
<commit_msg>ref(tests+database): set connection_timeout, increase polling interval, use db.Ping()<commit_after>package tests

import (
	"database/sql"
	"fmt"
	"github.com/deis/deis/tests/dockercli"
	"github.com/deis/deis/tests/mock"
	"github.com/deis/deis/tests/utils"
	"github.com/lib/pq"
	"testing"
	"time"
)

func OpenDeisDatabase(t *testing.T, host string, port string) *sql.DB {
	db, err := sql.Open("postgres", "postgres://deis:changeme123@"+host+":"+port+"/deis?sslmode=disable&connect_timeout=4")
	if err != nil {
		t.Fatal(err)
	}
	WaitForDatabase(t, db)
	return db
}

func WaitForDatabase(t *testing.T, db *sql.DB) {
	fmt.Println("--- Waiting for pg to be ready")
	for {
		err := db.Ping()
		if err, ok := err.(*pq.Error); ok {
			if err.Code.Name() == "cannot_connect_now" {
				fmt.Println(err.Message)
				time.Sleep(1000 * time.Millisecond)
				continue
			}
			t.Fatal(err)
		}
		fmt.Println("Ready")
		break
	}
}

func TryTableSelect(t *testing.T, db *sql.DB, tableName string, expectFailure bool) {
	_, err := db.Query("select * from " + tableName)

	if expectFailure {
		if err == nil {
			t.Fatal("The table should not exist")
		}
	} else {
		if err != nil {
			t.Fatal(err)
		}
	}
}

func execSql(t *testing.T, db *sql.DB, q string) {
	_, err := db.Query(q)
	if err != nil {
		t.Fatal(err)
	}
}

func TestDatabaseRecovery(t *testing.T) {
	var err error
	tag, etcdPort := utils.BuildTag(), utils.RandomPort()
	cli, stdout, _ := dockercli.NewClient()
	imageName := utils.ImagePrefix() + "database" + ":" + tag

	// start etcd container
	etcdName := "deis-etcd-" + tag
	dockercli.RunTestEtcd(t, etcdName, etcdPort)
	defer cli.CmdRm("-f", etcdName)

	// run mock ceph containers
	cephName := "deis-ceph-" + tag
	mock.RunMockCeph(t, cephName, cli, etcdPort)
	defer cli.CmdRm("-f", cephName)

	// create volumes
	databaseVolumeA := "deis-database-data-a-" + tag
	databaseVolumeB := "deis-database-data-b-" + tag
	defer cli.CmdRm("-f", databaseVolumeA)
	defer cli.CmdRm("-f", databaseVolumeB)
	go func() {
		fmt.Printf("--- Creating Volume A\n")
		_ = cli.CmdRm("-f", "-v", databaseVolumeA)
		dockercli.CreateVolume(t, cli, databaseVolumeA, "/var/lib/postgresql")

		fmt.Printf("--- Creating Volume B\n")

		_ = cli.CmdRm("-f", databaseVolumeB)
		dockercli.CreateVolume(t, cli, databaseVolumeB, "/var/lib/postgresql")
	}()
	dockercli.WaitForLine(t, stdout, databaseVolumeB, true)

	// setup database container start/stop routines
	host, port := utils.HostAddress(), utils.RandomPort()
	fmt.Printf("--- Run deis/database:%s at %s:%s\n", tag, host, port)
	name := "deis-database-" + tag
	defer cli.CmdRm("-f", name)
	startDatabase := func(volumeName string) {
		_ = cli.CmdRm("-f", name)
		err = dockercli.RunContainer(cli,
			"--name", name,
			"--volumes-from", volumeName,
			"--rm",
			"-p", port+":5432",
			"-e", "EXTERNAL_PORT="+port,
			"-e", "HOST="+host,
			"-e", "ETCD_PORT="+etcdPort,
			"-e", "ETCD_TTL=2",
			"-e", "BACKUP_FREQUENCY=0",
			"-e", "BACKUPS_TO_RETAIN=100",
			imageName)
	}

	stopDatabase := func() {
		fmt.Print("--- Stopping data-database... ")
		if err = stdout.Close(); err != nil {
			t.Fatal("Failed to closeStdout")
		}
		_ = cli.CmdStop(name)
		fmt.Println("Done")
	}

	//ACTION

	//STEP 1: start db with volume A and wait for init to complete
	fmt.Print("--- Starting database with Volume A... ")
	go startDatabase(databaseVolumeA)
	dockercli.WaitForLine(t, stdout, "database: postgres is running...", true)
	fmt.Println("Done")

	db := OpenDeisDatabase(t, host, port)
	TryTableSelect(t, db, "api_foo", true)

	stopDatabase()

	//STEP 2a: start db with volume B, wait for init and create the table
	cli, stdout, _ = dockercli.NewClient()
	fmt.Printf("--- Starting database with Volume B... ")
	go startDatabase(databaseVolumeB)
	dockercli.WaitForLine(t, stdout, "database: postgres is running...", true)
	fmt.Println("Done")

	db = OpenDeisDatabase(t, host, port)
	TryTableSelect(t, db, "api_foo", true)

	fmt.Println("--- Creating the table")
	execSql(t, db, "create table api_foo(t text)")

	//STEP 2b: make sure we observed full backup cycle after forced checkpoint
	fmt.Print("--- Waiting for the change to be backed up... ")
	dockercli.WaitForLine(t, stdout, "database: performing a backup...", true)
	dockercli.WaitForLine(t, stdout, "database: backup has been completed.", true)
	fmt.Println("Done")

	stopDatabase()

	//STEP 3: start db with volume A again and assert table existence
	cli, stdout, _ = dockercli.NewClient()
	fmt.Printf("--- Starting database with Volume A again... ")
	go startDatabase(databaseVolumeA)
	dockercli.WaitForLine(t, stdout, "database: postgres is running...", true)
	fmt.Println("Done")

	db = OpenDeisDatabase(t, host, port)
	TryTableSelect(t, db, "api_foo", false)

}
<|endoftext|>


<commit_before>package note

import (
	"reflect"
	"testing"
)

func TestUnMarshallTimeLog(t *testing.T) {

	cases := []struct {
		Note string
		Want CommitNote
	}{
		{
			`
[ver:1,total:1425]
environment/drone/run-tests.sh:725,1460066400:705,1460070000:20,m
environment/drone/run-tests-cron.sh:700,1460066400:540,1460070000:160,m
`,
			CommitNote{
				Files: []FileDetail{
					FileDetail{
						SourceFile: "environment/drone/run-tests.sh",
						TimeSpent:  725,
						Timeline:   map[int64]int{int64(1460066400): 705, int64(1460070000): 20},
						Status:     "m"},
					FileDetail{
						SourceFile: "environment/drone/run-tests-cron.sh",
						TimeSpent:  700,
						Timeline:   map[int64]int{int64(1460066400): 540, int64(1460070000): 160},
						Status:     "m"},
				},
			},
		},
		{
			`

[ver:1,total:1425]
environment/drone/run-tests.sh:725,1460066400:705,1460070000:20,m
environment/drone/run-tests-cron.sh:700,1460066400:540,1460070000:160,m

[ver:1,total:1425]
environment/drone/test.go:60,1460070000:60,r

`,
			CommitNote{
				Files: []FileDetail{
					FileDetail{
						SourceFile: "environment/drone/run-tests.sh",
						TimeSpent:  725,
						Timeline:   map[int64]int{int64(1460066400): 705, int64(1460070000): 20},
						Status:     "m"},
					FileDetail{
						SourceFile: "environment/drone/run-tests-cron.sh",
						TimeSpent:  700,
						Timeline:   map[int64]int{int64(1460066400): 540, int64(1460070000): 160},
						Status:     "m"},
					FileDetail{
						SourceFile: "environment/drone/test.go",
						TimeSpent:  60,
						Timeline:   map[int64]int{int64(1460070000): 60},
						Status:     "r"},
				},
			},
		},
	}

	for _, tc := range cases {
		got, err := UnMarshal(tc.Note)
		if err != nil {
			t.Errorf("unMarshalTimelog(%s), want error nil got error %s", tc.Note, err)
		}
		if !reflect.DeepEqual(tc.Want, got) {
			t.Errorf("unMarshalTimelog(%s), want:\n%+v\n got:\n%+v\n", tc.Note, tc.Want, got)
		}
	}

}
<commit_msg>Fix note test data<commit_after>package note

import (
	"reflect"
	"testing"
)

func TestUnMarshallTimeLog(t *testing.T) {

	cases := []struct {
		Note string
		Want CommitNote
	}{
		{
			`
[ver:1,total:1425]
environment/drone/run-tests.sh:725,1460066400:705,1460070000:20,m
environment/drone/run-tests-cron.sh:700,1460066400:540,1460070000:160,m
`,
			CommitNote{
				Files: []FileDetail{
					FileDetail{
						SourceFile: "environment/drone/run-tests.sh",
						TimeSpent:  725,
						Timeline:   map[int64]int{int64(1460066400): 705, int64(1460070000): 20},
						Status:     "m"},
					FileDetail{
						SourceFile: "environment/drone/run-tests-cron.sh",
						TimeSpent:  700,
						Timeline:   map[int64]int{int64(1460066400): 540, int64(1460070000): 160},
						Status:     "m"},
				},
			},
		},
		{
			`

[ver:1,total:1425]
environment/drone/run-tests.sh:725,1460066400:705,1460070000:20,m
environment/drone/run-tests-cron.sh:700,1460066400:540,1460070000:160,m

[ver:1,total:60]
environment/drone/test.go:60,1460070000:60,r

`,
			CommitNote{
				Files: []FileDetail{
					FileDetail{
						SourceFile: "environment/drone/run-tests.sh",
						TimeSpent:  725,
						Timeline:   map[int64]int{int64(1460066400): 705, int64(1460070000): 20},
						Status:     "m"},
					FileDetail{
						SourceFile: "environment/drone/run-tests-cron.sh",
						TimeSpent:  700,
						Timeline:   map[int64]int{int64(1460066400): 540, int64(1460070000): 160},
						Status:     "m"},
					FileDetail{
						SourceFile: "environment/drone/test.go",
						TimeSpent:  60,
						Timeline:   map[int64]int{int64(1460070000): 60},
						Status:     "r"},
				},
			},
		},
	}

	for _, tc := range cases {
		got, err := UnMarshal(tc.Note)
		if err != nil {
			t.Errorf("unMarshalTimelog(%s), want error nil got error %s", tc.Note, err)
		}
		if !reflect.DeepEqual(tc.Want, got) {
			t.Errorf("unMarshalTimelog(%s), want:\n%+v\n got:\n%+v\n", tc.Note, tc.Want, got)
		}
	}

}
<|endoftext|>


<commit_before>// Copyright 2012 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.
package fracserv

import (
	"encoding/gob"
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	bookmarkFn = flag.String("bookmarkFn", "/tmp/bookmark.db",
		"location to store bookmarked links")
)

func init() {
	b := NewBookmarks()
	http.HandleFunc("/bookmarks", func(w http.ResponseWriter, r *http.Request) {
		b.ListHandler(w, r)
	})
	http.HandleFunc("/bookmarks/add", func(w http.ResponseWriter, r *http.Request) {
		b.AddHandler(w, r)
	})
}

type Bookmark struct {
	Name  string    `json:"name"`
	Url   string    `json:"url"`
	Added time.Time `json:"added"`
}

type Bookmarks struct {
	Bookmarks []Bookmark
	addCh     chan Bookmark
	mu        sync.RWMutex
}

func NewBookmarks() *Bookmarks {
	b := &Bookmarks{
		Bookmarks: make([]Bookmark, 0),
		addCh:     make(chan Bookmark, 1),
	}
	go b.Loop()
	return b
}

func (b *Bookmarks) Add(bookmark Bookmark) {
	log.Print("Adding bookmark ", bookmark)
	b.addCh <- bookmark
}

func (b *Bookmarks) Load(fn string) error {
	log.Print("Loading bookmarks from ", fn)

	f, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer f.Close()

	gf := gob.NewDecoder(f)
	err = gf.Decode(&b.Bookmarks)
	return err
}

func (b *Bookmarks) Save(fn string) error {
	log.Print("Saving bookmarks to ", fn)
	f, err := os.OpenFile(fn, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	gf := gob.NewEncoder(f)
	err = gf.Encode(b.Bookmarks)
	return err
}

func (b *Bookmarks) Loop() {
	err := b.Load(*bookmarkFn)
	if err != nil {
		log.Print("Error loading bookmarks ", err)
	}
	for {
		select {
		case bookmark := <-b.addCh:
			b.mu.Lock()
			b.Bookmarks = append(b.Bookmarks, bookmark)
			// TODO This doesn't really scale
			err = b.Save(*bookmarkFn)
			if err != nil {
				log.Print("Error saving bookmarks ", err)
			}
			b.mu.Unlock()
		}
	}
}

func (b *Bookmarks) AddHandler(w http.ResponseWriter, req *http.Request) {
	q := req.URL.Query()
	name := q.Get("name")
	url := q.Get("url")
	if url == "" || name == "" {
		http.Error(w, "missing url or name: "+req.URL.String(),
			http.StatusInternalServerError)
		return
	}
	b.Add(Bookmark{
		Name:  q.Get("name"),
		Url:   q.Get("url"),
		Added: time.Now(),
	})

	http.Redirect(w, req, "/", http.StatusMovedPermanently)
}

func (b *Bookmarks) ListHandler(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	b.mu.Lock()
	defer b.mu.Unlock()
	j, err := json.Marshal(b.Bookmarks)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	w.Write(j)
}
<commit_msg>Load bookmarks on first access to let flags be set.<commit_after>// Copyright 2012 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.
package fracserv

import (
	"encoding/gob"
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	bookmarkFn = flag.String("bookmarkFn", "/tmp/bookmark.db",
		"location to store bookmarked links")
	loaded = sync.Once{}
)

func init() {
	b := NewBookmarks()
	http.HandleFunc("/bookmarks", func(w http.ResponseWriter, r *http.Request) {
		b.ListHandler(w, r)
	})
	http.HandleFunc("/bookmarks/add", func(w http.ResponseWriter, r *http.Request) {
		b.AddHandler(w, r)
	})
}

type Bookmark struct {
	Name  string    `json:"name"`
	Url   string    `json:"url"`
	Added time.Time `json:"added"`
}

type Bookmarks struct {
	Bookmarks []Bookmark
	addCh     chan Bookmark
	mu        sync.RWMutex
}

func NewBookmarks() *Bookmarks {
	return &Bookmarks{
		Bookmarks: make([]Bookmark, 0),
		addCh:     make(chan Bookmark, 1),
	}
}

func (b *Bookmarks) Load(fn string) error {
	log.Print("Loading bookmarks from ", fn)

	f, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer f.Close()

	gf := gob.NewDecoder(f)
	err = gf.Decode(&b.Bookmarks)
	return err
}

func (b *Bookmarks) Save(fn string) error {
	log.Print("Saving bookmarks to ", fn)
	f, err := os.OpenFile(fn, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	gf := gob.NewEncoder(f)
	err = gf.Encode(b.Bookmarks)
	return err
}

func (b *Bookmarks) Add(bookmark Bookmark) {
	log.Print("Adding bookmark ", bookmark)
	b.mu.Lock()
	defer b.mu.Unlock()
	b.Bookmarks = append(b.Bookmarks, bookmark)
	// TODO This doesn't really scale
	err := b.Save(*bookmarkFn)
	if err != nil {
		log.Print("Error saving bookmarks ", err)
	}
}

func (b *Bookmarks) loadOnce() {
	err := b.Load(*bookmarkFn)
	if err != nil {
		log.Print("Error loading bookmarks ", err)
	}
}

func (b *Bookmarks) AddHandler(w http.ResponseWriter, req *http.Request) {
	loaded.Do(func() { b.loadOnce() })

	q := req.URL.Query()
	name := q.Get("name")
	url := q.Get("url")
	if url == "" || name == "" {
		http.Error(w, "missing url or name: "+req.URL.String(),
			http.StatusInternalServerError)
		return
	}
	b.Add(Bookmark{
		Name:  q.Get("name"),
		Url:   q.Get("url"),
		Added: time.Now(),
	})

	http.Redirect(w, req, "/", http.StatusMovedPermanently)
}

func (b *Bookmarks) ListHandler(w http.ResponseWriter, req *http.Request) {
	loaded.Do(func() { b.loadOnce() })

	w.Header().Set("Content-Type", "application/json")
	b.mu.Lock()
	defer b.mu.Unlock()
	j, err := json.Marshal(b.Bookmarks)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	w.Write(j)
}
<|endoftext|>


<commit_before>package utils

import (
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"time"
)

type JSONError struct {
	Code    int    `json:"code,omitempty"`
	Message string `json:"message,omitempty"`
}

func (e *JSONError) Error() string {
	return e.Message
}

type JSONProgress struct {
	Current int   `json:"current,omitempty"`
	Total   int   `json:"total,omitempty"`
	Start   int64 `json:"start,omitempty"`
}

func (p *JSONProgress) String() string {
	if p.Current == 0 && p.Total == 0 {
		return ""
	}
	current := HumanSize(int64(p.Current))
	if p.Total == 0 {
		return fmt.Sprintf("%8v/?", current)
	}
	total := HumanSize(int64(p.Total))
	percentage := int(float64(p.Current)/float64(p.Total)*100) / 2

	fromStart := time.Now().UTC().Sub(time.Unix(int64(p.Start), 0))
	perEntry := fromStart / time.Duration(p.Current)
	left := time.Duration(p.Total-p.Current) * perEntry
	left = (left / time.Second) * time.Second
	return fmt.Sprintf("[%s>%s] %8v/%v %s", strings.Repeat("=", percentage), strings.Repeat(" ", 50-percentage), current, total, left.String())
}

type JSONMessage struct {
	Status          string        `json:"status,omitempty"`
	Progress        *JSONProgress `json:"progressDetail,omitempty"`
	ProgressMessage string        `json:"progress,omitempty"` //deprecated
	ID              string        `json:"id,omitempty"`
	From            string        `json:"from,omitempty"`
	Time            int64         `json:"time,omitempty"`
	Error           *JSONError    `json:"errorDetail,omitempty"`
	ErrorMessage    string        `json:"error,omitempty"` //deprecated
}

func (jm *JSONMessage) Display(out io.Writer, isTerminal bool) error {
	if jm.Error != nil {
		if jm.Error.Code == 401 {
			return fmt.Errorf("Authentication is required.")
		}
		return jm.Error
	}
	endl := ""
	if isTerminal {
		// <ESC>[2K = erase entire current line
		fmt.Fprintf(out, "%c[2K\r", 27)
		endl = "\r\n"
	}
	if jm.Time != 0 {
		fmt.Fprintf(out, "[%s] ", time.Unix(jm.Time, 0))
	}
	if jm.ID != "" {
		fmt.Fprintf(out, "%s: ", jm.ID)
	}
	if jm.From != "" {
		fmt.Fprintf(out, "(from %s) ", jm.From)
	}
	if jm.Progress != nil {
		fmt.Fprintf(out, "%s %s%s", jm.Status, jm.Progress.String(), endl)
	} else if jm.ProgressMessage != "" { //deprecated
		fmt.Fprintf(out, "%s %s%s", jm.Status, jm.ProgressMessage, endl)
	} else {
		fmt.Fprintf(out, "%s%s", jm.Status, endl)
	}
	return nil
}

func DisplayJSONMessagesStream(in io.Reader, out io.Writer, isTerminal bool) error {
	dec := json.NewDecoder(in)
	ids := make(map[string]int)
	diff := 0
	for {
		jm := JSONMessage{}
		if err := dec.Decode(&jm); err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		if (jm.Progress != nil || jm.ProgressMessage != "") && jm.ID != "" {
			line, ok := ids[jm.ID]
			if !ok {
				line = len(ids)
				ids[jm.ID] = line
				fmt.Fprintf(out, "\n")
				diff = 0
			} else {
				diff = len(ids) - line
			}
			if isTerminal {
				// <ESC>[{diff}A = move cursor up diff rows
				fmt.Fprintf(out, "%c[%dA", 27, diff)
			}
		}
		err := jm.Display(out, isTerminal)
		if jm.ID != "" {
			if isTerminal {
				// <ESC>[{diff}B = move cursor down diff rows
				fmt.Fprintf(out, "%c[%dB", 27, diff)
			}
		}
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>small reformatting jsonmessage<commit_after>package utils

import (
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"time"
)

type JSONError struct {
	Code    int    `json:"code,omitempty"`
	Message string `json:"message,omitempty"`
}

func (e *JSONError) Error() string {
	return e.Message
}

type JSONProgress struct {
	Current int   `json:"current,omitempty"`
	Total   int   `json:"total,omitempty"`
	Start   int64 `json:"start,omitempty"`
}

func (p *JSONProgress) String() string {
	if p.Current == 0 && p.Total == 0 {
		return ""
	}
	current := HumanSize(int64(p.Current))
	if p.Total == 0 {
		return fmt.Sprintf("%8v/?", current)
	}
	total := HumanSize(int64(p.Total))
	percentage := int(float64(p.Current)/float64(p.Total)*100) / 2

	fromStart := time.Now().UTC().Sub(time.Unix(int64(p.Start), 0))
	perEntry := fromStart / time.Duration(p.Current)
	left := time.Duration(p.Total-p.Current) * perEntry
	left = (left / time.Second) * time.Second
	return fmt.Sprintf("[%s>%s] %8v/%v %s", strings.Repeat("=", percentage), strings.Repeat(" ", 50-percentage), current, total, left.String())
}

type JSONMessage struct {
	Status          string        `json:"status,omitempty"`
	Progress        *JSONProgress `json:"progressDetail,omitempty"`
	ProgressMessage string        `json:"progress,omitempty"` //deprecated
	ID              string        `json:"id,omitempty"`
	From            string        `json:"from,omitempty"`
	Time            int64         `json:"time,omitempty"`
	Error           *JSONError    `json:"errorDetail,omitempty"`
	ErrorMessage    string        `json:"error,omitempty"` //deprecated
}

func (jm *JSONMessage) Display(out io.Writer, isTerminal bool) error {
	if jm.Error != nil {
		if jm.Error.Code == 401 {
			return fmt.Errorf("Authentication is required.")
		}
		return jm.Error
	}
	var endl string
	if isTerminal {
		// <ESC>[2K = erase entire current line
		fmt.Fprintf(out, "%c[2K\r", 27)
		endl = "\r\n"
	}
	if jm.Time != 0 {
		fmt.Fprintf(out, "[%s] ", time.Unix(jm.Time, 0))
	}
	if jm.ID != "" {
		fmt.Fprintf(out, "%s: ", jm.ID)
	}
	if jm.From != "" {
		fmt.Fprintf(out, "(from %s) ", jm.From)
	}
	if jm.Progress != nil {
		fmt.Fprintf(out, "%s %s%s", jm.Status, jm.Progress.String(), endl)
	} else if jm.ProgressMessage != "" { //deprecated
		fmt.Fprintf(out, "%s %s%s", jm.Status, jm.ProgressMessage, endl)
	} else {
		fmt.Fprintf(out, "%s%s", jm.Status, endl)
	}
	return nil
}

func DisplayJSONMessagesStream(in io.Reader, out io.Writer, isTerminal bool) error {
	var (
		dec  = json.NewDecoder(in)
		ids  = make(map[string]int)
		diff = 0
	)
	for {
		var jm JSONMessage
		if err := dec.Decode(&jm); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if (jm.Progress != nil || jm.ProgressMessage != "") && jm.ID != "" {
			line, ok := ids[jm.ID]
			if !ok {
				line = len(ids)
				ids[jm.ID] = line
				fmt.Fprintf(out, "\n")
				diff = 0
			} else {
				diff = len(ids) - line
			}
			if isTerminal {
				// <ESC>[{diff}A = move cursor up diff rows
				fmt.Fprintf(out, "%c[%dA", 27, diff)
			}
		}
		err := jm.Display(out, isTerminal)
		if jm.ID != "" {
			if isTerminal {
				// <ESC>[{diff}B = move cursor down diff rows
				fmt.Fprintf(out, "%c[%dB", 27, diff)
			}
		}
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
    "gopkg.in/hypersleep/easyssh.v0"
    "strconv"
    "sync"
    "github.com/mgutz/ansi"
)

/**
################################################################################
							Miscallenous-Section
################################################################################
*/

/**
*	Prints the current Version of the goo application
*/
func printVersion(){
	fmt.Println("0.0.1")
}

/**
*	Prints the help dialog
*/
func printHelp(){
	fmt.Println("usage: goo [options...] <planet>... -c=\"<command>\"")
	fmt.Println("Options:")
	fmt.Println("-s=\"<path/to/script>\", --script=\"<path/to/script>\"  Execute script and return result")
	fmt.Println("-p, --pretty     Pretty print output as a table")
	fmt.Println("-t, --type       Show type of planet")
	fmt.Println("-h, --help       This help text")
	fmt.Println("-v, --version    Show version number")
	fmt.Println("-d, --debug	  Show extended debug informations")
}

/**
*	Formats and prints the given output and error.
*/
func throwErrOut(out []byte, err error){
	fmt.Print(fmt.Sprint(err) + " output is: " + string(out) + "called from ErrOut. ")
	os.Stderr.WriteString(fmt.Sprint(err) + " output is: " + string(out) + "called from ErrOut. ")
	os.Exit(1)
}

/**
*	Formats and prints the given error.
*/
func throwErr(err error){
	fmt.Print(fmt.Sprint(err)  + " called from Err. ")
	os.Stderr.WriteString(fmt.Sprint(err) + "called from Err. ")
	os.Exit(1)
}

/**
################################################################################
								SSH-Section
################################################################################
*/

/**
*	Executes a command on a remote ssh server
*	@params:
*		connDet: connection details in following form: user@hostname
*		cmd: command to be executed
*/
func execCommand(connDet string, cmd string, wg *sync.WaitGroup , wait bool, prettyFlag bool){
	if(prettyFlag){
		fmt.Println(ansi.Color("Executing command ","241") + ansi.Color(cmd,"white+hu") + ansi.Color(" on ","241") + ansi.Color(connDet,"white+hu"))
	}else{
		fmt.Println("Executing command " + cmd + " on " + connDet)
	}
	fmt.Println("")
	user := getUser(connDet)
	hostname := getHost(connDet)

	ssh := &easyssh.MakeConfig{
		User:   user,
		Server: hostname,
		Key:  os.Getenv("ORBIT_KEY"),
		Port: "22",
	}
	// Call Run method with command you want to run on remote server.
	out, err := ssh.Run(cmd)
	// Handle errors
	if err != nil {
		throwErr(err)
	} else {
		fmt.Println(out)
	}
	if (wait){
		wg.Done()
	}
}

/**
*	Uploads a file to the remote server
*/
func uploadFile(connDet string, path string){
	user := getUser(connDet)
	hostname := getHost(connDet)

	ssh := &easyssh.MakeConfig{
		User:   user,
		Server: hostname,
		Key:  os.Getenv("ORBIT_KEY"),
		Port: "22",
	}

	// Call Scp method with file you want to upload to remote server.
	err := ssh.Scp(path)

	// Handle errors
	if err != nil {
		throwErr(err)
	}
}

/**
*	Uploads and executes a script on a given planet
*	@params:
*		connDet: 	Connection details to planet
*		scriptPath: Path to script
*/
func upAndExecScript(connDet string, scriptPath string, wg *sync.WaitGroup, prettyFlag bool){
	uploadFile(connDet,scriptPath)
	path := strings.Split(scriptPath,"/")
	execCommand(connDet,"chmod +x " + path[len(path)-1],wg,false,prettyFlag)
	execCommand(connDet,"./" + path[len(path)-1],wg,false,prettyFlag)
	wg.Done()
}

/**
################################################################################
						Information-Retrieval-Section
################################################################################
*/

/**
*	Returns the contents of args in following order:
*	prettyprint flag
*	script flag
*	script path
*	command
*	planets
*/
func procArgs(args []string) (bool, bool, string, string, []string, bool, bool){
	prettyFlag := false
	scriptFlag := false
	typeFlag := false
	debugFlag := false
	var scriptPath string = ""
	var command string = ""
	planets := make([]string,0)

	cleanArgs := args[1:]

	for _, argument := range cleanArgs {
		if(strings.HasPrefix(argument,"-h") || strings.HasPrefix(argument,"--help")){
			printHelp()
			os.Exit(0)
		}else if(strings.HasPrefix(argument,"-p") || strings.HasPrefix(argument,"--pretty")){
			prettyFlag = true
		}else if(strings.HasPrefix(argument,"-t") || strings.HasPrefix(argument,"--type")){
			typeFlag = true
		}else if(strings.HasPrefix(argument,"-d") || strings.HasPrefix(argument,"--debug")){
			debugFlag = true
		}else if(strings.HasPrefix(argument,"-v") || strings.HasPrefix(argument,"--version")){
			printVersion()
			os.Exit(0)
		}else if(strings.HasPrefix(argument,"-c") || strings.HasPrefix(argument,"--command")){
			// TODO what if theres a = in the command itself?
			command = strings.TrimSuffix(strings.TrimPrefix(strings.Split(argument,"=")[1],"\""),"\"")
		}else if(strings.HasPrefix(argument,"-s") || strings.HasPrefix(argument,"--script")){
			scriptFlag = true
			scriptPath = strings.Split(argument,"=")[1]
		}else{
			planets = append(planets,argument)
		}
	}
	if(len(args) <3){
		printHelp()
		os.Exit(0)
	}

	_ = prettyFlag

	return prettyFlag,scriptFlag,scriptPath,command,planets,debugFlag,typeFlag
}

/**
*	Splits the given connectiondetails and returns the hostname
*	@params:
*		connDet: Connection details in following form: user@hostname
*	@return: hostname
*/
func getHost(connDet string) string{
	toReturn := strings.Split(connDet,"@")
	return toReturn[1]
}

/**
*	Splits the given connectiondetails and returns the user
*	@params:
*		connDet: Connection details in following form: user@hostname
*	@return: user
*/
func getUser(connDet string) string{
	toReturn := strings.Split(connDet,"@")
	return toReturn[0]
}

/**
*	Returns the type of a given planet
*	@params:
*		id: The planets id
*	@return: The planets type
*/
func getType(id string) string{
	cmd 	 := exec.Command("ff","-t" ,id)
	out, err := cmd.CombinedOutput()
	if err != nil {
    	throwErrOut(out,err)
	}
	return strings.TrimSpace(string(out))
}

/**
*	Returns the connection details to a given planet
*	@params:
*		id: The planets id
*	@return: The connection details to the planet
*/
func getConnDet(id string) string{
	cmd 	 := exec.Command("ff",id)
	out, err := cmd.CombinedOutput()
	if err != nil {
    	throwErrOut(out,err)
	}
	return strings.TrimSpace(string(out))
}

/**
*					DEPRECATED
*
*	Extracts the desired argument from the arguments list.
*	@params:
*		args: Arguments to be searched in.
*		type: Type of desired Argument (command,id)
*		position: starting position of desired argument
*	@return: The desired arguments
*/
func getArg(args []string, argType string, position int) string{
	switch argType{
		case "command":
			var command string  = args[position]
			var cmdArgs []string
			if(len(args) > (position+1)){
				cmdArgs = args[(position+1):(len(args))]
				for _, argument := range cmdArgs {
					command += (" " + argument)
				}
			}
			return command
		default:
			return args[position]
	}

}

/**
################################################################################
								Main-Section
################################################################################
*/

/**
*	Main function
*/
func main() {

	var args []string = os.Args

	prettyFlag,scriptFlag,scriptPath,command,planets,debugFlag,typeFlag := procArgs(args)

	_ = prettyFlag
	if(debugFlag){
		fmt.Println(args)
		fmt.Println("prettyflag " + strconv.FormatBool(prettyFlag))
		fmt.Println("scriptflag " + strconv.FormatBool(scriptFlag))
		fmt.Println("scriptpath " + scriptPath)
		fmt.Println("command " + command)
		for _, planet := range planets {
			fmt.Println("planet " + planet)
		}
	}

	var wg sync.WaitGroup
	wg.Add(len(planets))
	for _, planet := range planets {
		if(prettyFlag){
			fmt.Println(ansi.Color("################################################################################","blue"))
		}
		if(typeFlag){
			fmt.Println("The type of " + planet + " is " + getType(planet))
		}
		switch getType(planet) {
			case "server":
				var connDet string = getConnDet(planet)
				if(scriptFlag){
					go upAndExecScript(connDet,scriptPath,&wg,prettyFlag)
				}else{
					go execCommand(connDet,command,&wg,true,prettyFlag)
				}
			case "db":
				fmt.Fprintf(os.Stderr, "This Type of Connection is not yet supported.")
				os.Exit(1)
			case "web":
				fmt.Fprintf(os.Stderr, "This Type of Connection is not supported.")
				os.Exit(1)
			default:
				wg.Done()
		}
	}
	wg.Wait()
}
<commit_msg>improved prettyprint<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
    "gopkg.in/hypersleep/easyssh.v0"
    "strconv"
    "sync"
    "github.com/mgutz/ansi"
)

/**
################################################################################
							Miscallenous-Section
################################################################################
*/

/**
*	Prints the current Version of the goo application
*/
func printVersion(){
	fmt.Println("0.0.1")
}

/**
*	Prints the help dialog
*/
func printHelp(){
	fmt.Println("usage: goo [options...] <planet>... -c=\"<command>\"")
	fmt.Println("Options:")
	fmt.Println("-s=\"<path/to/script>\", --script=\"<path/to/script>\"  Execute script and return result")
	fmt.Println("-p, --pretty     Pretty print output as a table")
	fmt.Println("-t, --type       Show type of planet")
	fmt.Println("-h, --help       This help text")
	fmt.Println("-v, --version    Show version number")
	fmt.Println("-d, --debug	  Show extended debug informations")
}

/**
*	Formats and prints the given output and error.
*/
func throwErrOut(out []byte, err error){
	fmt.Print(fmt.Sprint(err) + " output is: " + string(out) + "called from ErrOut. ")
	os.Stderr.WriteString(fmt.Sprint(err) + " output is: " + string(out) + "called from ErrOut. ")
	os.Exit(1)
}

/**
*	Formats and prints the given error.
*/
func throwErr(err error){
	fmt.Print(fmt.Sprint(err)  + " called from Err. ")
	os.Stderr.WriteString(fmt.Sprint(err) + "called from Err. ")
	os.Exit(1)
}

/**
################################################################################
								SSH-Section
################################################################################
*/

/**
*	Executes a command on a remote ssh server
*	@params:
*		connDet: connection details in following form: user@hostname
*		cmd: command to be executed
*/
func execCommand(connDet string, cmd string, wg *sync.WaitGroup , wait bool, prettyFlag bool){

	user := getUser(connDet)
	hostname := getHost(connDet)

	ssh := &easyssh.MakeConfig{
		User:   user,
		Server: hostname,
		Key:  os.Getenv("ORBIT_KEY"),
		Port: "22",
	}
	// Call Run method with command you want to run on remote server.
	out, err := ssh.Run(cmd)
	// Handle errors
	if err != nil {
		throwErr(err)
	} else {
		if(prettyFlag){
			fmt.Println(ansi.Color("################################################################################","blue"))
			fmt.Println(ansi.Color("Executing command ","241") + ansi.Color(cmd,"white+hu") + ansi.Color(" on ","241") + ansi.Color(connDet,"white+hu"))
		}else{
			fmt.Println("Executing command " + cmd + " on " + connDet)
		}
		fmt.Println("")
		fmt.Println(out)
	}
	if (wait){
		wg.Done()
	}
}

/**
*	Uploads a file to the remote server
*/
func uploadFile(connDet string, path string){
	user := getUser(connDet)
	hostname := getHost(connDet)

	ssh := &easyssh.MakeConfig{
		User:   user,
		Server: hostname,
		Key:  os.Getenv("ORBIT_KEY"),
		Port: "22",
	}

	// Call Scp method with file you want to upload to remote server.
	err := ssh.Scp(path)

	// Handle errors
	if err != nil {
		throwErr(err)
	}
}

/**
*	Uploads and executes a script on a given planet
*	@params:
*		connDet: 	Connection details to planet
*		scriptPath: Path to script
*/
func upAndExecScript(connDet string, scriptPath string, wg *sync.WaitGroup, prettyFlag bool){
	uploadFile(connDet,scriptPath)
	path := strings.Split(scriptPath,"/")
	execCommand(connDet,"chmod +x " + path[len(path)-1],wg,false,prettyFlag)
	execCommand(connDet,"./" + path[len(path)-1],wg,false,prettyFlag)
	wg.Done()
}

/**
################################################################################
						Information-Retrieval-Section
################################################################################
*/

/**
*	Returns the contents of args in following order:
*	prettyprint flag
*	script flag
*	script path
*	command
*	planets
*/
func procArgs(args []string) (bool, bool, string, string, []string, bool, bool){
	prettyFlag := false
	scriptFlag := false
	typeFlag := false
	debugFlag := false
	var scriptPath string = ""
	var command string = ""
	planets := make([]string,0)

	cleanArgs := args[1:]

	for _, argument := range cleanArgs {
		if(strings.HasPrefix(argument,"-h") || strings.HasPrefix(argument,"--help")){
			printHelp()
			os.Exit(0)
		}else if(strings.HasPrefix(argument,"-p") || strings.HasPrefix(argument,"--pretty")){
			prettyFlag = true
		}else if(strings.HasPrefix(argument,"-t") || strings.HasPrefix(argument,"--type")){
			typeFlag = true
		}else if(strings.HasPrefix(argument,"-d") || strings.HasPrefix(argument,"--debug")){
			debugFlag = true
		}else if(strings.HasPrefix(argument,"-v") || strings.HasPrefix(argument,"--version")){
			printVersion()
			os.Exit(0)
		}else if(strings.HasPrefix(argument,"-c") || strings.HasPrefix(argument,"--command")){
			// TODO what if theres a = in the command itself?
			command = strings.TrimSuffix(strings.TrimPrefix(strings.Split(argument,"=")[1],"\""),"\"")
		}else if(strings.HasPrefix(argument,"-s") || strings.HasPrefix(argument,"--script")){
			scriptFlag = true
			scriptPath = strings.Split(argument,"=")[1]
		}else{
			planets = append(planets,argument)
		}
	}
	if(len(args) <3){
		printHelp()
		os.Exit(0)
	}

	_ = prettyFlag

	return prettyFlag,scriptFlag,scriptPath,command,planets,debugFlag,typeFlag
}

/**
*	Splits the given connectiondetails and returns the hostname
*	@params:
*		connDet: Connection details in following form: user@hostname
*	@return: hostname
*/
func getHost(connDet string) string{
	toReturn := strings.Split(connDet,"@")
	return toReturn[1]
}

/**
*	Splits the given connectiondetails and returns the user
*	@params:
*		connDet: Connection details in following form: user@hostname
*	@return: user
*/
func getUser(connDet string) string{
	toReturn := strings.Split(connDet,"@")
	return toReturn[0]
}

/**
*	Returns the type of a given planet
*	@params:
*		id: The planets id
*	@return: The planets type
*/
func getType(id string) string{
	cmd 	 := exec.Command("ff","-t" ,id)
	out, err := cmd.CombinedOutput()
	if err != nil {
    	throwErrOut(out,err)
	}
	return strings.TrimSpace(string(out))
}

/**
*	Returns the connection details to a given planet
*	@params:
*		id: The planets id
*	@return: The connection details to the planet
*/
func getConnDet(id string) string{
	cmd 	 := exec.Command("ff",id)
	out, err := cmd.CombinedOutput()
	if err != nil {
    	throwErrOut(out,err)
	}
	return strings.TrimSpace(string(out))
}

/**
*					DEPRECATED
*
*	Extracts the desired argument from the arguments list.
*	@params:
*		args: Arguments to be searched in.
*		type: Type of desired Argument (command,id)
*		position: starting position of desired argument
*	@return: The desired arguments
*/
func getArg(args []string, argType string, position int) string{
	switch argType{
		case "command":
			var command string  = args[position]
			var cmdArgs []string
			if(len(args) > (position+1)){
				cmdArgs = args[(position+1):(len(args))]
				for _, argument := range cmdArgs {
					command += (" " + argument)
				}
			}
			return command
		default:
			return args[position]
	}

}

/**
################################################################################
								Main-Section
################################################################################
*/

/**
*	Main function
*/
func main() {

	var args []string = os.Args

	prettyFlag,scriptFlag,scriptPath,command,planets,debugFlag,typeFlag := procArgs(args)

	_ = prettyFlag
	if(debugFlag){
		fmt.Println(args)
		fmt.Println("prettyflag " + strconv.FormatBool(prettyFlag))
		fmt.Println("scriptflag " + strconv.FormatBool(scriptFlag))
		fmt.Println("scriptpath " + scriptPath)
		fmt.Println("command " + command)
		for _, planet := range planets {
			fmt.Println("planet " + planet)
		}
	}

	var wg sync.WaitGroup
	wg.Add(len(planets))
	for _, planet := range planets {
		if(typeFlag){
			fmt.Println("The type of " + planet + " is " + getType(planet))
		}
		switch getType(planet) {
			case "server":
				var connDet string = getConnDet(planet)
				if(scriptFlag){
					go upAndExecScript(connDet,scriptPath,&wg,prettyFlag)
				}else{
					go execCommand(connDet,command,&wg,true,prettyFlag)
				}
			case "db":
				fmt.Fprintf(os.Stderr, "This Type of Connection is not yet supported.")
				os.Exit(1)
			case "web":
				fmt.Fprintf(os.Stderr, "This Type of Connection is not supported.")
				os.Exit(1)
			default:
				wg.Done()
		}
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>package tchannel

// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

import (
	"bytes"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel/golang/typed"
)

func TestNoFragmentation(t *testing.T) {
	in, out := buildChannels(ChecksumTypeCrc32)

	arg1 := []byte("Hello")
	w := newBodyWriter(out)
	require.Nil(t, w.WriteArgument(BytesOutput(arg1), true))

	// Should be a single frame
	// fragment flags(1), checksum type (1), checksum(5), chunk size(2), chunk(5)
	expectedFrames := combineBuffers([][]byte{
		[]byte{0x00, byte(ChecksumTypeCrc32)},
		ChecksumTypeCrc32.New().Add([]byte("Hello")),
		[]byte{0x00, 0x05},
		[]byte("Hello")})
	assertFramesEqual(t, expectedFrames, out.sentFragments, "no fragmentation")

	r1 := newBodyReader(in, true)
	rarg1 := make([]byte, len(arg1))
	if _, err := r1.Read(rarg1); err != nil {
		require.Nil(t, err)
	}

	assert.Equal(t, arg1, rarg1)
	require.Nil(t, r1.endArgument())
}

func TestFragmentationRoundTrip(t *testing.T) {
	in, out := buildChannels(ChecksumTypeCrc32)

	// Write three arguments, each of which should span fragments
	arg1 := make([]byte, MaxFramePayloadSize*2+756)
	for i := range arg1 {
		arg1[i] = byte(i % 0x0F)
	}
	w := newBodyWriter(out)
	if _, err := w.Write(arg1); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg2 := make([]byte, MaxFramePayloadSize+229)
	for i := range arg2 {
		arg2[i] = byte(i%0x0F) + 0x10
	}
	if _, err := w.Write(arg2); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg3 := make([]byte, MaxFramePayloadSize+72)
	for i := range arg3 {
		arg3[i] = byte(i%0x0F) + 0x20
	}
	if _, err := w.Write(arg3); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(true))

	// Read the three arguments
	r1 := newBodyReader(in, false)

	rarg1 := make([]byte, len(arg1))
	if _, err := r1.Read(rarg1); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg1, rarg1)
	require.Nil(t, r1.endArgument())

	r2 := newBodyReader(in, false)
	rarg2 := make([]byte, len(arg2))
	if _, err := r2.Read(rarg2); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg2, rarg2)
	require.Nil(t, r2.endArgument())

	r3 := newBodyReader(in, true)
	rarg3 := make([]byte, len(arg3))
	if _, err := r3.Read(rarg3); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg3, rarg3)
	require.Nil(t, r3.endArgument())
}

func TestArgEndOnFragmentBoundary(t *testing.T) {
	// Each argument should line up exactly at the end of each fragment
	in, out := buildChannels(ChecksumTypeCrc32)

	// Calculate the number of bytes available in the fragment content,
	// which is the size of the full frame minus the header content for the
	// fragment.  Header content consists of 1 byte flag, 1 byte checksum
	// type, 4 byte checksum value, for a total of 6 bytes
	fragmentContentSize := int(MaxFramePayloadSize) - 6
	arg1 := make([]byte, fragmentContentSize-2) // reserve 2 bytes for the arg chunk size
	for i := range arg1 {
		arg1[i] = byte(i % 0x0F)
	}
	w := newBodyWriter(out)
	if _, err := w.Write(arg1); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg2 := make([]byte, len(arg1)-2) // additional 2 byte trailing size for arg1
	for i := range arg2 {
		arg2[i] = byte(i % 0x1F)
	}
	if _, err := w.Write(arg2); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg3 := make([]byte, len(arg2)) // additional 2 byte trailing size for arg2
	for i := range arg3 {
		arg3[i] = byte(i % 0x2F)
	}
	if _, err := w.Write(arg3); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(true))

	// We should have sent 4 fragments (one for arg1, one for zero arg1
	// size + arg2, one for zero arg2 size + arg3, one for zero arg3 size)
	sentFragments := out.sentFragments
	require.Equal(t, 4, len(sentFragments))
	lastFragment := sentFragments[len(sentFragments)-1]

	// 1 byte flags, 1 byte checksum type, 4 bytes checksum, 2 bytes size (0)
	require.Equal(t, 8, int(lastFragment.Header.PayloadSize()))
	r1 := newBodyReader(in, false)

	rarg1 := make([]byte, len(arg1))
	if _, err := r1.Read(rarg1); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg1, rarg1)
	require.Nil(t, r1.endArgument())

	r2 := newBodyReader(in, false)
	rarg2 := make([]byte, len(arg2))
	if _, err := r2.Read(rarg2); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg2, rarg2)
	require.Nil(t, r2.endArgument())

	r3 := newBodyReader(in, true)
	rarg3 := make([]byte, len(arg3))
	if _, err := r3.Read(rarg3); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg3, rarg3)
	require.Nil(t, r3.endArgument())
}

func TestEmptyFragments(t *testing.T) {
	in, out := buildChannels(ChecksumTypeCrc32)

	w := newBodyWriter(out)
	require.Nil(t, w.WriteArgument(BytesOutput(nil), false))
	require.Nil(t, w.WriteArgument(BytesOutput(nil), true))

	r1 := newBodyReader(in, false)
	var arg1 BytesInput
	require.Nil(t, r1.ReadArgument(&arg1, false))
	assert.Equal(t, 0, len(arg1))

	r2 := newBodyReader(in, true)
	var arg2 BytesInput
	require.Nil(t, r2.ReadArgument(&arg2, true))
	assert.Equal(t, 0, len(arg2))
}

func buildChannels(checksumType ChecksumType) (*inFragments, *outFragments) {
	ch := make(chan *Frame, 512)

	in := &inFragments{ch: ch}
	out := &outFragments{ch: ch, checksum: checksumType.New()}
	return in, out
}

type inFragments struct {
	checksum Checksum
	ch       <-chan *Frame
	current  *inFragment
}

type sampleMessage struct{}

func (m *sampleMessage) ID() uint32                       { return 0xDEADBEEF }
func (m *sampleMessage) messageType() messageType         { return messageTypeCallReq }
func (m *sampleMessage) read(r *typed.ReadBuffer) error   { return nil }
func (m *sampleMessage) write(w *typed.WriteBuffer) error { return nil }

func (in *inFragments) waitForFragment() (*inFragment, error) {
	if in.current == nil || !in.current.hasMoreChunks() {
		var err error
		f := <-in.ch
		if in.current, err = newInboundFragment(f, &sampleMessage{}, in.checksum); err != nil {
			return nil, err
		}

		in.checksum = in.current.checksum
	}

	return in.current, nil
}

type outFragments struct {
	fragmentSize  int
	checksum      Checksum
	ch            chan<- *Frame
	sentFragments []*Frame
}

func (out *outFragments) beginFragment() (*outFragment, error) {
	return newOutboundFragment(NewFrame(MaxFramePayloadSize), &sampleMessage{}, out.checksum)
}

func (out *outFragments) flushFragment(toSend *outFragment, last bool) error {
	f := toSend.finish(last)
	out.ch <- f
	out.sentFragments = append(out.sentFragments, f)
	return nil
}

func assertFramesEqual(t *testing.T, expected [][]byte, frames []*Frame, msg string) {
	require.Equal(t, len(expected), len(frames), fmt.Sprintf("incorrect number of frames for %s", msg))

	for i := range expected {
		assert.Equal(t, len(expected[i]), int(frames[i].Header.PayloadSize()),
			fmt.Sprintf("incorrect size for frame %d of %s", i, msg))
		assert.Equal(t, expected[i], frames[i].Payload[:frames[i].Header.PayloadSize()])
	}
}

func combineBuffers(elements ...[][]byte) [][]byte {
	var buffers [][]byte
	for i := range elements {
		buffers = append(buffers, bytes.Join(elements[i], []byte{}))
	}

	return buffers
}
<commit_msg>Start on fragmentation test<commit_after>package tchannel

import (
	"bytes"
	"sync"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel/golang/typed"
)

const (
	testFragmentHeaderSize = 1 /* flags */ + 1 /* checksum type */ + 4 /* CRC32 checksum */

	testFragmentPayloadSize = 10 // enough room for a small payload
	testFragmentSize        = testFragmentHeaderSize + testFragmentPayloadSize
)

func TestFragmentationEmptyArgs(t *testing.T) {
	sendCh := make(fragmentChannel, 10)
	recvCh := make(fragmentChannel, 10)

	w := newFragmentingWriter(sendCh, ChecksumTypeCrc32.New())
	r := newFragmentingReader(recvCh)

	var fragments [][]byte
	var args [][]byte
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for fragment := range sendCh {
			fragments = append(fragments, fragment)
			recvCh <- fragment
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		var arg BytesInput
		require.NoError(t, r.ReadArgument(&arg, false))
		args = append(args, arg)
		arg = nil
		require.NoError(t, r.ReadArgument(&arg, false))
		args = append(args, arg)
		arg = nil
		require.NoError(t, r.ReadArgument(&arg, true))
		args = append(args, arg)
		arg = nil
	}()

	require.NoError(t, w.WriteArgument(BytesOutput(nil), false))
	require.NoError(t, w.WriteArgument(BytesOutput(nil), false))
	require.NoError(t, w.WriteArgument(BytesOutput(nil), true))
	close(sendCh)

	wg.Wait()
	assert.Equal(t, [][]byte{[]byte{}, []byte{}, []byte{}}, args)

	// Make sure the fragments look as we expected
	expectedFragments := buffers([][]byte{{
		0x00, // flags
		byte(ChecksumTypeCrc32), 0x00, 0x00, 0x00, 0x00, // empty checksum
		0x00, 0x00, // arg 1 (length no body)
		0x00, 0x00, // arg 2 (length no body)
		0x00, 0x00}, // arg 3 (length no body)
	})
	assert.Equal(t, expectedFragments, fragments)
}

func TestSingleFragment(t *testing.T) {
}

func TestMultipleFragments(t *testing.T) {
}

func TestMiddleArgOnFragmentBoundary(t *testing.T) {
}

func TestLastArgOnFragmentBoundary(t *testing.T) {
}

type fragmentChannel chan []byte

func (ch fragmentChannel) newFragment(initial bool, checksum Checksum) (*writableFragment, error) {
	wbuf := typed.NewWriteBuffer(make([]byte, testFragmentSize))
	fragment := new(writableFragment)
	fragment.flagsRef = wbuf.DeferByte()
	wbuf.WriteByte(byte(checksum.TypeCode()))
	fragment.checksumRef = wbuf.DeferBytes(checksum.Size())
	fragment.checksum = checksum
	fragment.contents = wbuf
	return fragment, wbuf.Err()
}

func (ch fragmentChannel) flushFragment(fragment *writableFragment) error {
	var buf bytes.Buffer
	fragment.contents.FlushTo(&buf)
	ch <- buf.Bytes()
	return nil
}

func (ch fragmentChannel) recvNextFragment(initial bool) (*readableFragment, error) {
	rbuf := typed.NewReadBuffer(<-ch)
	fragment := new(readableFragment)
	fragment.flags = rbuf.ReadByte()
	fragment.checksumType = ChecksumType(rbuf.ReadByte())
	fragment.checksum = rbuf.ReadBytes(fragment.checksumType.ChecksumSize())
	fragment.contents = rbuf
	return fragment, rbuf.Err()
}

func buffers(elements ...[][]byte) [][]byte {
	var buffers [][]byte
	for i := range elements {
		buffers = append(buffers, bytes.Join(elements[i], []byte{}))
	}

	return buffers
}
<|endoftext|>


<commit_before>package jlogutil

import (
	"errors"

	"github.com/fastly/jlog-go"
)

// JReader embeds a reader with a subscriber name.
type JReader struct {
	jlog.Reader
	subscriber string
}

// Reader opens the jlog at the given path, attaches the given subscriber
// for reading and returns the opened jlog.
func Reader(path, subscriber string) (*JReader, error) {
	log, err := jlog.NewReader(path, nil)
	if err != nil {
		return nil, err
	}
	subs, err := log.ListSubscribers()
	if err != nil {
		return nil, err
	}

	jreader := &JReader{Reader: log}

	for _, sub := range subs {
		if sub == subscriber {
			return nil, errors.New("subscriber already exists")
		}
	}

	err = log.AddSubscriber(subscriber, jlog.END)
	if err != nil {
		return nil, err
	}
	err = log.Open(subscriber)
	if err != nil {
		jreader.Close()
		return nil, err
	}
	jreader.subscriber = subscriber
	return jreader, nil
}

// Removes the JReader's specific reader and closes the jlog.
func (j *JReader) Close() {
	j.RemoveSubscriber(j.subscriber)
	j.Close()
}
<commit_msg>jlogutil<commit_after>package jlogutil

import (
	"errors"

	"github.com/fastly/jlog-go"
)

// JReader embeds a reader with a subscriber name.
type JReader struct {
	jlog.Reader
	subscriber string
}

// NewReader opens the jlog at the given path, attaches the given subscriber
// for reading and returns the opened jlog.
func NewReader(path, subscriber string) (*JReader, error) {
	log, err := jlog.NewReader(path, nil)
	if err != nil {
		return nil, err
	}
	subs, err := log.ListSubscribers()
	if err != nil {
		return nil, err
	}

	jreader := &JReader{Reader: log}

	for _, sub := range subs {
		if sub == subscriber {
			return nil, errors.New("subscriber already exists")
		}
	}

	err = log.AddSubscriber(subscriber, jlog.END)
	if err != nil {
		return nil, err
	}
	err = log.Open(subscriber)
	if err != nil {
		jreader.Close()
		return nil, err
	}
	jreader.subscriber = subscriber
	return jreader, nil
}

// Removes the JReader's specific reader and closes the jlog.
func (j *JReader) Close() {
	j.RemoveSubscriber(j.subscriber)
	j.Close()
}
<|endoftext|>