<commit_before>package mithril

import "fmt"

var (
	// SelfClosingTagToken is a list of self closing elements tag
	SelfClosingTagToken = []string{"area", "base", "br", "col", "command", "embed", "hr", "img", "input", "keygen", "link", "meta", "param", "source", "track", "wbr", "!doctype"}
)

// VirtualElement is an object which can convert to real DOM element
type VirtualElement struct {
	Tag      string
	Attrs    []Attribute
	Children interface{}
}

// Attr returns or sets the attribute
func (el *VirtualElement) Attr(key string, val ...string) string {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			if len(val) > 0 {
				attr.Val(val[0])
			}
			return attr.Val()
		}
	}
	if len(val) > 0 {
		el.Attrs = append(el.Attrs, NewStringAttr(key, val[0]))
		return val[0]
	}
	return ""
}

// HasAttr returns true if VirtualElement has this attribute in attributes array
func (el *VirtualElement) HasAttr(key string) bool {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			return true
		}
	}
	return false
}

// IsSelfClosingTag returns true if Tag is a self closing tag element
func (el *VirtualElement) IsSelfClosingTag() bool {
	for _, val := range SelfClosingTagToken {
		if val == el.Tag {
			return true
		}
	}
	return false
}

// String returns the html string of the element
func (el *VirtualElement) String() string {
	res := make(chan string)
	go func() {
		html := "<"
		if el.Tag != "" {
			html += el.Tag
		} else {
			html += "div"
		}
		for _, attr := range el.Attrs {
			html += (" " + attr.Key())
			if len(attr.Val()) > 0 {
				html += ("=\"" + attr.Val() + "\"")
			}
		}
		if el.IsSelfClosingTag() {
			html += " />"
		} else {
			html += ">"
			switch obj := el.Children.(type) {
			case *VirtualElement:
				html += obj.String()
			case string:
				html += obj
			case int, int32, int64:
				html += fmt.Sprintf("%d", obj)
			case float32, float64:
				html += fmt.Sprintf("%f", obj)
			}
			html += "</"
			if el.Tag != "" {
				html += el.Tag
			} else {
				html += "div"
			}
			html += ">"
		}
		res <- html
	}()
	return <-res
}
<commit_msg>improve element render implementation<commit_after>package mithril

import "fmt"

var (
	// SelfClosingTagToken is a list of self closing elements tag
	SelfClosingTagToken = []string{"area", "base", "br", "col", "command", "embed", "hr", "img", "input", "keygen", "link", "meta", "param", "source", "track", "wbr", "!doctype"}
)

// VirtualElement is an object which can convert to real DOM element
type VirtualElement struct {
	Tag      string
	Attrs    []Attribute
	Children interface{}
}

// Attr returns or sets the attribute
func (el *VirtualElement) Attr(key string, val ...string) string {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			if len(val) > 0 {
				attr.Val(val[0])
			}
			return attr.Val()
		}
	}
	if len(val) > 0 {
		el.Attrs = append(el.Attrs, NewStringAttr(key, val[0]))
		return val[0]
	}
	return ""
}

// HasAttr returns true if VirtualElement has this attribute in attributes array
func (el *VirtualElement) HasAttr(key string) bool {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			return true
		}
	}
	return false
}

// IsSelfClosingTag returns true if Tag is a self closing tag element
func (el *VirtualElement) IsSelfClosingTag() bool {
	for _, val := range SelfClosingTagToken {
		if val == el.Tag {
			return true
		}
	}
	return false
}

// Render returns html string of anything
func (el *VirtualElement) Render(item interface{}) string {
	str := ""
	switch obj := item.(type) {
	case []interface{}:
		for _, item := range obj {
			str += el.Render(item)
		}
	case *VirtualElement:
		str += obj.String()
	case string:
		str += obj
	case int, int32, int64:
		str += fmt.Sprintf("%d", obj)
	case float32, float64:
		str += fmt.Sprintf("%f", obj)
	}
	return str
}

// String returns the html string of the element
func (el *VirtualElement) String() string {
	res := make(chan string)
	go func() {
		html := "<"
		if el.Tag != "" {
			html += el.Tag
		} else {
			html += "div"
		}
		for _, attr := range el.Attrs {
			html += (" " + attr.Key())
			if len(attr.Val()) > 0 {
				html += ("=\"" + attr.Val() + "\"")
			}
		}
		if el.IsSelfClosingTag() {
			html += " />"
		} else {
			html += ">"
			html += el.Render(el.Children)
			html += "</"
			if el.Tag != "" {
				html += el.Tag
			} else {
				html += "div"
			}
			html += ">"
		}
		res <- html
	}()
	return <-res
}
<|endoftext|>


<commit_before>// Play the game "pig". A codewalk from Go documentation.
//
// https://golang.org/doc/codewalk/functions/

package main

import (
	"time"
	"log"
	"math/rand"
)

const win = 100 // The winning score in a game of Pig

// A score includes scores accumlated in previous turns for each player, as well
// as the points scored by the current player in this turn.
type score struct {
	player, opponent, thisTurn int
}

// An action type defines the actions that can be taken while playing the game
//
// An action type is a function that takes a score and returns the resulting
// score and whther or not the currrent turn is over.
//
// If the turn is over, the player and opponent fields in the resulting score
// should be swapped, as it is now the other player's turn.
type action func(current score) (result score, turnIsOver bool)

// roll returns the result (result, turnIsOver) outcome of simulating a die
// roll. If the roll value is 1, the thisTurn score is abandoned, and the
// players' roles swap. Otherwise, the roll value is added to thisTurn.
func roll(s score) (score, bool) {
	outcome := rand.Intn(6) + 1 // A die roll: a random int [1, 6]
	if outcome == 1 {
		return score{s.opponent, s.player, 0}, true
	}

	return score{s.player, s.opponent, outcome + s.thisTurn}, false
}

// A player's thisTurn score is added to their total, and their turn ends.
//
// When a turn ends, the scores of the two players swap.
func stay(s score) (score, bool) {
	return score{s.opponent, s.player + s.thisTurn, 0}, true
}

// A strategy chooses and action for any given score
type strategy func(score) action

// stayAtK returns a strategy that rolls until this Turn is at least k, then
// stays.
func stayAtK(k int) strategy {
	// TODO(topher): why do we need to re-define the function signature?
	return func(s score) action {
		if s.thisTurn >= k {
			return stay
		}
		return roll
	}
}

// play simulates a Pig game and returns the winner (0 or 1)
func play(strategy0, strategy1 strategy) int {
	strategies := []strategy{strategy0, strategy1}
	var s score
	var turnIsOver bool
	// Randomly decide who goes first
	currentPlayer := rand.Intn(2)
	for s.player+s.thisTurn < win {
		action := strategies[currentPlayer](s)
		s, turnIsOver = action(s)
		if turnIsOver {
			if currentPlayer == 1 {
				currentPlayer = 0
			} else {
				currentPlayer = 1
			}
			log.Print("Turn is over. scores, new currentPlayer:", s, currentPlayer)
		}
	}
	return currentPlayer
}
<commit_msg>play the Pig game with fixed one strategy<commit_after>// Play the game "pig". A codewalk from Go documentation.
//
// https://golang.org/doc/codewalk/functions/

package main

import (
	"time"
	"log"
	"math/rand"
)

const win = 100 // The winning score in a game of Pig

// A score includes scores accumlated in previous turns for each player, as well
// as the points scored by the current player in this turn.
type score struct {
	player, opponent, thisTurn int
}

// An action type defines the actions that can be taken while playing the game
//
// An action type is a function that takes a score and returns the resulting
// score and whther or not the currrent turn is over.
//
// If the turn is over, the player and opponent fields in the resulting score
// should be swapped, as it is now the other player's turn.
type action func(current score) (result score, turnIsOver bool)

// roll returns the result (result, turnIsOver) outcome of simulating a die
// roll. If the roll value is 1, the thisTurn score is abandoned, and the
// players' roles swap. Otherwise, the roll value is added to thisTurn.
func roll(s score) (score, bool) {
	outcome := rand.Intn(6) + 1 // A die roll: a random int [1, 6]
	if outcome == 1 {
		return score{s.opponent, s.player, 0}, true
	}

	return score{s.player, s.opponent, outcome + s.thisTurn}, false
}

// A player's thisTurn score is added to their total, and their turn ends.
//
// When a turn ends, the scores of the two players swap.
func stay(s score) (score, bool) {
	return score{s.opponent, s.player + s.thisTurn, 0}, true
}

// A strategy chooses and action for any given score
type strategy func(score) action

// stayAtK returns a strategy that rolls until this Turn is at least k, then
// stays.
func stayAtK(k int) strategy {
	// TODO(topher): why do we need to re-define the function signature?
	return func(s score) action {
		if s.thisTurn >= k {
			return stay
		}
		return roll
	}
}

// play simulates a Pig game and returns the winner (0 or 1)
func play(strategy0, strategy1 strategy) int {
	strategies := []strategy{strategy0, strategy1}
	var s score
	var turnIsOver bool
	// Randomly decide who goes first
	currentPlayer := rand.Intn(2)
	for s.player+s.thisTurn < win {
		action := strategies[currentPlayer](s)
		s, turnIsOver = action(s)
		if turnIsOver {
			if currentPlayer == 1 {
				currentPlayer = 0
			} else {
				currentPlayer = 1
			}
			log.Print("Turn is over. scores, new currentPlayer:", s, currentPlayer)
		}
	}
	return currentPlayer
}

func main() {
	rand.Seed(time.Now().UTC().UnixNano())
	log.Print(play(stayAtK(5), stayAtK(49)))
}
<|endoftext|>


<commit_before>package handler

import (
	"encoding/json"
	"errors"
	"io"
	"net/http"

	"github.com/davecgh/go-spew/spew"
	"github.com/tmc/graphql"
	"github.com/tmc/graphql/parser"
	"github.com/tmc/graphql/schema"
)

var (
	ErrMalformedQuery = errors.New("malformed query")
)

type Error struct {
	Message string `json:"message"`
}

type Result struct {
	Data  interface{} `json:"data,omitempty"`
	Error *Error      `json:"error,omitempty"`
}

type SchemaHandler struct {
	schema *schema.Schema
}

func New(schema *schema.Schema) *SchemaHandler {
	return &SchemaHandler{schema: schema}
}

func writeErr(w io.Writer, err error) {
	writeJSON(w, Result{Error: &Error{Message: err.Error()}})
}
func writeJSON(w io.Writer, data interface{}) {
	json.NewEncoder(w).Encode(data)
}

// ServeHTTP provides an entrypoint into a graphql schema. It pulls the query from
// the 'q' GET parameter.
func (h *SchemaHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	//TODO(tmc): reject non-GET requests
	q := r.URL.Query().Get("q")
	spew.Dump(q)
	result, err := parser.Parse("", []byte(q))
	if err != nil {
		writeErr(w, err)
		return
	}
	call := result.(graphql.Call)
	spew.Dump(call)
	result, err = h.schema.HandleCall(call)
	if err != nil {
		writeErr(w, err)
		return
	}
	writeJSON(w, Result{Data: result})
}
<commit_msg>remove debug output<commit_after>package handler

import (
	"encoding/json"
	"errors"
	"io"
	"net/http"

	"github.com/tmc/graphql"
	"github.com/tmc/graphql/parser"
	"github.com/tmc/graphql/schema"
)

var (
	ErrMalformedQuery = errors.New("malformed query")
)

type Error struct {
	Message string `json:"message"`
}

type Result struct {
	Data  interface{} `json:"data,omitempty"`
	Error *Error      `json:"error,omitempty"`
}

type SchemaHandler struct {
	schema *schema.Schema
}

func New(schema *schema.Schema) *SchemaHandler {
	return &SchemaHandler{schema: schema}
}

func writeErr(w io.Writer, err error) {
	writeJSON(w, Result{Error: &Error{Message: err.Error()}})
}
func writeJSON(w io.Writer, data interface{}) {
	json.NewEncoder(w).Encode(data)
}

// ServeHTTP provides an entrypoint into a graphql schema. It pulls the query from
// the 'q' GET parameter.
func (h *SchemaHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	//TODO(tmc): reject non-GET requests
	q := r.URL.Query().Get("q")
	result, err := parser.Parse("", []byte(q))
	if err != nil {
		writeErr(w, err)
		return
	}
	call := result.(graphql.Call)
	result, err = h.schema.HandleCall(call)
	if err != nil {
		writeErr(w, err)
		return
	}
	writeJSON(w, Result{Data: result})
}
<|endoftext|>


<commit_before>package goutils

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
)

type Pkg struct {
	Dir         string   `json:"Dir"`
	ImportPath  string   `json:"ImportPath"`
	Name        string   `json:"Name"`
	Target      string   `json:"Target"`
	Stale       string   `json:"Stale"`
	StaleReason string   `json:"StaleReason"`
	Root        string   `json:"Root"`
	GoFiles     []string `json:"GoFiles"`
	Imports     []string `json:"Imports"`
	Deps        []string `json:"Deps"`
	Incomplete  bool     `json:"Incomplete"`
	DepsErrors  []struct {
		ImportStack []string `json:"ImportStack"`
		Pos         string   `json:"Pos"`
		Err         string   `json:"Err"`
	} `json:"DepsErrors"`
	TestGoFiles []string `json:"TestGoFiles"`
	TestImports []string `json:"TestImports"`

	NoStdDepOKPkgs    []string
	NoStdDepErrPkgs   []NoStdDepErrPkg
	nostdDepErrPkgMap map[string]bool
}

type NoStdDepErrPkg struct {
	PkgName string
	Tips    string
}

func NewPkg(bs []byte) (*Pkg, error) {
	var pkg Pkg
	err := json.Unmarshal(bs, &pkg)
	if CheckErr(err) {
		return nil, err
	}
	pkg.analyse()
	return &pkg, err
}

func (p *Pkg) analyse() {
	p.analyseDepsErrors()
	p.NoStdDepOKPkgs = make([]string, 0, len(p.Deps)-len(p.DepsErrors))
	for _, dep := range p.Deps {
		if p.nostdDepErrPkgMap[dep] {
			continue
		}
		p.NoStdDepOKPkgs = append(p.NoStdDepOKPkgs, dep)
	}
}

func (p *Pkg) analyseDepsErrors() []*NoStdDepErrPkg {
	desize := len(p.DepsErrors)
	if desize <= 0 {
		return nil
	}
	p.nostdDepErrPkgMap = make(map[string]bool)
	ret := make([]*NoStdDepErrPkg, 0, desize)
	for _, depsErr := range p.DepsErrors {
		tips := depsErr.Err
		for _, imp := range depsErr.ImportStack {
			if imp == p.ImportPath {
				continue
			}
			p.nostdDepErrPkgMap[imp] = true
			ret = append(ret, &NoStdDepErrPkg{
				PkgName: imp,
				Tips:    tips,
			})
		}
	}
	return ret
}

func TrimGopath(pkg string) string {
	abspath, err := filepath.Abs(os.Getenv("GOPATH"))
	if CheckErr(err) {
		return pkg
	}
	return strings.TrimPrefix(pkg, abspath+"/src/")
}

/*{
	"Dir": "/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow/src/github.com/toukii/bkg",
	"ImportPath": "github.com/toukii/bkg",
	"Name": "main",
	"Target": "/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow/bin/bkg",
	"Stale": true,
	"StaleReason": "build ID mismatch",
	"Root": "/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow",
	"GoFiles": [
		"main.go"
	],
	"Imports": [
		"time"
	],
	"Deps": [
		"bufio",
		"unicode",
		"unicode/utf16",
		"unicode/utf8",
		"unsafe",
		"vendor/golang_org/x/crypto/chacha20poly1305",
		"vendor/golang_org/x/crypto/chacha20poly1305/internal/chacha20",
		"vendor/golang_org/x/crypto/curve25519",
		"vendor/golang_org/x/crypto/poly1305",
		"vendor/golang_org/x/net/http2/hpack",
		"vendor/golang_org/x/text/unicode/norm"
	],
	"Incomplete": true,
	"DepsErrors": [
		{
			"ImportStack": [
				"github.com/toukii/bkg",
				"github.com/everfore/exc"
			],
			"Pos": "main.go:11:2",
			"Err": "cannot find package \"github.com/everfore/exc\" in any of:\n\t/Users/toukii/PATH/Go/go/src/github.com/everfore/exc (from $GOROOT)\n\t/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow/src/github.com/everfore/exc (from $GOPATH)"
		}
	],
	"TestGoFiles": [
		"jsnm_test.go",
		"zhihu_test.go"
	],
	"TestImports": [
		"encoding/json",
		"fmt",
		"github.com/toukii/goutils",
		"github.com/toukii/membership/pkg3/go-simplejson",
		"io/ioutil",
		"net/http",
		"reflect",
		"testing"
	]
}
*/
<commit_msg>test-user<commit_after>package goutils

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

type Pkg struct {
	Dir         string   `json:"Dir"`
	ImportPath  string   `json:"ImportPath"`
	Name        string   `json:"Name"`
	Target      string   `json:"Target"`
	Stale       string   `json:"Stale"`
	StaleReason string   `json:"StaleReason"`
	Root        string   `json:"Root"`
	GoFiles     []string `json:"GoFiles"`
	Imports     []string `json:"Imports"`
	Deps        []string `json:"Deps"`
	Incomplete  bool     `json:"Incomplete"`
	DepsErrors  []struct {
		ImportStack []string `json:"ImportStack"`
		Pos         string   `json:"Pos"`
		Err         string   `json:"Err"`
	} `json:"DepsErrors"`
	TestGoFiles []string `json:"TestGoFiles"`
	TestImports []string `json:"TestImports"`

	NoStdDepOKPkgs    []string
	NoStdDepErrPkgs   []NoStdDepErrPkg
	nostdDepErrPkgMap map[string]bool
}

type NoStdDepErrPkg struct {
	PkgName string
	Tips    string
}

func NewPkg(bs []byte) (*Pkg, error) {
	var pkg Pkg
	err := json.Unmarshal(bs, &pkg)
	if CheckErr(err) {
		return nil, err
	}
	pkg.analyse()
	return &pkg, err
}

func (p *Pkg) analyse() {
	p.analyseDepsErrors()
	p.NoStdDepOKPkgs = make([]string, 0, len(p.Deps)-len(p.DepsErrors))
	for _, dep := range p.Deps {
		if p.nostdDepErrPkgMap[dep] {
			continue
		}
		p.NoStdDepOKPkgs = append(p.NoStdDepOKPkgs, dep)
	}
}

func (p *Pkg) analyseDepsErrors() []*NoStdDepErrPkg {
	desize := len(p.DepsErrors)
	if desize <= 0 {
		return nil
	}
	p.nostdDepErrPkgMap = make(map[string]bool)
	ret := make([]*NoStdDepErrPkg, 0, desize)
	for _, depsErr := range p.DepsErrors {
		tips := depsErr.Err
		for _, imp := range depsErr.ImportStack {
			if imp == p.ImportPath {
				continue
			}
			p.nostdDepErrPkgMap[imp] = true
			ret = append(ret, &NoStdDepErrPkg{
				PkgName: imp,
				Tips:    tips,
			})
		}
	}
	return ret
}

type PkgNames []interface{}

func (p *PkgNames) DisplayString() string {
	ret := ""
	for i, name := range p {
		ret += fmt.Sprintf("\t%d. %s", i, name)
	}
	return ret
}

func TrimGopath(pkg string) string {
	abspath, err := filepath.Abs(os.Getenv("GOPATH"))
	if CheckErr(err) {
		return pkg
	}
	return strings.TrimPrefix(pkg, abspath+"/src/")
}

/*{
	"Dir": "/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow/src/github.com/toukii/bkg",
	"ImportPath": "github.com/toukii/bkg",
	"Name": "main",
	"Target": "/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow/bin/bkg",
	"Stale": true,
	"StaleReason": "build ID mismatch",
	"Root": "/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow",
	"GoFiles": [
		"main.go"
	],
	"Imports": [
		"time"
	],
	"Deps": [
		"bufio",
		"unicode",
		"unicode/utf16",
		"unicode/utf8",
		"unsafe",
		"vendor/golang_org/x/crypto/chacha20poly1305",
		"vendor/golang_org/x/crypto/chacha20poly1305/internal/chacha20",
		"vendor/golang_org/x/crypto/curve25519",
		"vendor/golang_org/x/crypto/poly1305",
		"vendor/golang_org/x/net/http2/hpack",
		"vendor/golang_org/x/text/unicode/norm"
	],
	"Incomplete": true,
	"DepsErrors": [
		{
			"ImportStack": [
				"github.com/toukii/bkg",
				"github.com/everfore/exc"
			],
			"Pos": "main.go:11:2",
			"Err": "cannot find package \"github.com/everfore/exc\" in any of:\n\t/Users/toukii/PATH/Go/go/src/github.com/everfore/exc (from $GOROOT)\n\t/Users/toukii/PATH/GOPATH/src/gitlab.1dmy.com/ezbuy/goflow/src/github.com/everfore/exc (from $GOPATH)"
		}
	],
	"TestGoFiles": [
		"jsnm_test.go",
		"zhihu_test.go"
	],
	"TestImports": [
		"encoding/json",
		"fmt",
		"github.com/toukii/goutils",
		"github.com/toukii/membership/pkg3/go-simplejson",
		"io/ioutil",
		"net/http",
		"reflect",
		"testing"
	]
}
*/
<|endoftext|>


<commit_before>package nomad

import (
	"testing"
	"time"

	"fmt"

	"github.com/hashicorp/consul/agent/consul/autopilot"
	"github.com/hashicorp/consul/testutil/retry"
	"github.com/hashicorp/nomad/testutil"
	"github.com/hashicorp/raft"
	"github.com/hashicorp/serf/serf"
)

// wantPeers determines whether the server has the given
// number of voting raft peers.
func wantPeers(s *Server, peers int) error {
	future := s.raft.GetConfiguration()
	if err := future.Error(); err != nil {
		return err
	}

	n := autopilot.NumPeers(future.Configuration())
	if got, want := n, peers; got != want {
		return fmt.Errorf("got %d peers want %d", got, want)
	}
	return nil
}

// wantRaft determines if the servers have all of each other in their
// Raft configurations,
func wantRaft(servers []*Server) error {
	// Make sure all the servers are represented in the Raft config,
	// and that there are no extras.
	verifyRaft := func(c raft.Configuration) error {
		want := make(map[raft.ServerID]bool)
		for _, s := range servers {
			want[s.config.RaftConfig.LocalID] = true
		}

		for _, s := range c.Servers {
			if !want[s.ID] {
				return fmt.Errorf("don't want %q", s.ID)
			}
			delete(want, s.ID)
		}

		if len(want) > 0 {
			return fmt.Errorf("didn't find %v", want)
		}
		return nil
	}

	for _, s := range servers {
		future := s.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			return err
		}
		if err := verifyRaft(future.Configuration()); err != nil {
			return err
		}
	}
	return nil
}

func TestAutopilot_CleanupDeadServer(t *testing.T) {
	t.Parallel()
	for i := 1; i <= 3; i++ {
		testCleanupDeadServer(t, i)
	}
}

func testCleanupDeadServer(t *testing.T, raftVersion int) {
	conf := func(c *Config) {
		c.DevDisableBootstrap = true
		c.BootstrapExpect = 3
		c.RaftConfig.ProtocolVersion = raft.ProtocolVersion(raftVersion)
	}
	s1 := TestServer(t, conf)
	defer s1.Shutdown()

	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	servers := []*Server{s1, s2, s3}

	// Try to join
	TestJoin(t, s1, s2, s3)

	for _, s := range servers {
		retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s, 3)) })
	}

	// Bring up a new server
	s4 := TestServer(t, conf)
	defer s4.Shutdown()

	// Kill a non-leader server
	s3.Shutdown()
	retry.Run(t, func(r *retry.R) {
		alive := 0
		for _, m := range s1.Members() {
			if m.Status == serf.StatusAlive {
				alive++
			}
		}
		if alive != 2 {
			r.Fatal(nil)
		}
	})

	// Join the new server
	TestJoin(t, s1, s4)
	servers[2] = s4

	// Make sure the dead server is removed and we're back to 3 total peers
	for _, s := range servers {
		retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s, 3)) })
	}
}

func TestAutopilot_CleanupDeadServerPeriodic(t *testing.T) {
	t.Parallel()
	s1 := TestServer(t, nil)
	defer s1.Shutdown()

	conf := func(c *Config) {
		c.DevDisableBootstrap = true
	}

	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	s4 := TestServer(t, conf)
	defer s4.Shutdown()

	s5 := TestServer(t, conf)
	defer s5.Shutdown()

	servers := []*Server{s1, s2, s3, s4, s5}

	// Join the servers to s1, and wait until they are all promoted to
	// voters.
	TestJoin(t, s1, servers[1:]...)
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 5))
		}
	})

	// Kill a non-leader server
	s4.Shutdown()

	// Should be removed from the peers automatically
	servers = []*Server{s1, s2, s3, s5}
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 4))
		}
	})
}

func TestAutopilot_RollingUpdate(t *testing.T) {
	t.Parallel()
	s1 := TestServer(t, func(c *Config) {
		c.RaftConfig.ProtocolVersion = 3
	})
	defer s1.Shutdown()

	conf := func(c *Config) {
		c.DevDisableBootstrap = true
		c.RaftConfig.ProtocolVersion = 3
	}

	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	// Join the servers to s1, and wait until they are all promoted to
	// voters.
	servers := []*Server{s1, s2, s3}
	TestJoin(t, s1, s2, s3)
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 3))
		}
	})

	// Add one more server like we are doing a rolling update.
	s4 := TestServer(t, conf)
	defer s4.Shutdown()
	TestJoin(t, s1, s4)
	servers = append(servers, s4)
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 3))
		}
	})

	// Now kill one of the "old" nodes like we are doing a rolling update.
	s3.Shutdown()

	isVoter := func() bool {
		future := s1.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			t.Fatalf("err: %v", err)
		}
		for _, s := range future.Configuration().Servers {
			if string(s.ID) == string(s4.config.NodeID) {
				return s.Suffrage == raft.Voter
			}
		}
		t.Fatalf("didn't find s4")
		return false
	}

	// Wait for s4 to stabilize, get promoted to a voter, and for s3 to be
	// removed.
	servers = []*Server{s1, s2, s4}
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 3))
		}
		if !isVoter() {
			r.Fatalf("should be a voter")
		}
	})
}

func TestAutopilot_CleanupStaleRaftServer(t *testing.T) {
	t.Parallel()
	s1 := TestServer(t, nil)
	defer s1.Shutdown()

	conf := func(c *Config) {
		c.DevDisableBootstrap = true
	}
	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	s4 := TestServer(t, conf)
	defer s4.Shutdown()

	servers := []*Server{s1, s2, s3}

	// Join the servers to s1
	TestJoin(t, s1, s2, s3)

	leader := waitForStableLeadership(t, servers)

	// Add s4 to peers directly
	addr := fmt.Sprintf("127.0.0.1:%d", s4.config.RPCAddr.Port)
	future := leader.raft.AddVoter(raft.ServerID(s4.config.NodeID), raft.ServerAddress(addr), 0, 0)
	if err := future.Error(); err != nil {
		t.Fatal(err)
	}

	// Verify we have 4 peers
	peers, err := s1.numPeers()
	if err != nil {
		t.Fatal(err)
	}
	if peers != 4 {
		t.Fatalf("bad: %v", peers)
	}

	// Wait for s4 to be removed
	for _, s := range []*Server{s1, s2, s3} {
		retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s, 3)) })
	}
}

func TestAutopilot_PromoteNonVoter(t *testing.T) {
	t.Parallel()
	s1 := TestServer(t, func(c *Config) {
		c.RaftConfig.ProtocolVersion = 3
	})
	defer s1.Shutdown()
	codec := rpcClient(t, s1)
	defer codec.Close()
	testutil.WaitForLeader(t, s1.RPC)

	s2 := TestServer(t, func(c *Config) {
		c.DevDisableBootstrap = true
		c.RaftConfig.ProtocolVersion = 3
	})
	defer s2.Shutdown()
	TestJoin(t, s1, s2)

	// Make sure we see it as a nonvoter initially. We wait until half
	// the stabilization period has passed.
	retry.Run(t, func(r *retry.R) {
		future := s1.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			r.Fatal(err)
		}

		servers := future.Configuration().Servers
		if len(servers) != 2 {
			r.Fatalf("bad: %v", servers)
		}
		if servers[1].Suffrage != raft.Nonvoter {
			r.Fatalf("bad: %v", servers)
		}
		health := s1.autopilot.GetServerHealth(string(servers[1].ID))
		if health == nil {
			r.Fatalf("nil health, %v", s1.autopilot.GetClusterHealth())
		}
		if !health.Healthy {
			r.Fatalf("bad: %v", health)
		}
		if time.Since(health.StableSince) < s1.config.AutopilotConfig.ServerStabilizationTime/2 {
			r.Fatal("stable period not elapsed")
		}
	})

	// Make sure it ends up as a voter.
	retry.Run(t, func(r *retry.R) {
		future := s1.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			r.Fatal(err)
		}

		servers := future.Configuration().Servers
		if len(servers) != 2 {
			r.Fatalf("bad: %v", servers)
		}
		if servers[1].Suffrage != raft.Voter {
			r.Fatalf("bad: %v", servers)
		}
	})
}
<commit_msg>tests: give up on TestAutopilot_CleanupStaleRaftServer for now<commit_after>package nomad

import (
	"testing"
	"time"

	"fmt"

	"github.com/hashicorp/consul/agent/consul/autopilot"
	"github.com/hashicorp/consul/testutil/retry"
	"github.com/hashicorp/nomad/testutil"
	"github.com/hashicorp/raft"
	"github.com/hashicorp/serf/serf"
)

// wantPeers determines whether the server has the given
// number of voting raft peers.
func wantPeers(s *Server, peers int) error {
	future := s.raft.GetConfiguration()
	if err := future.Error(); err != nil {
		return err
	}

	n := autopilot.NumPeers(future.Configuration())
	if got, want := n, peers; got != want {
		return fmt.Errorf("got %d peers want %d", got, want)
	}
	return nil
}

// wantRaft determines if the servers have all of each other in their
// Raft configurations,
func wantRaft(servers []*Server) error {
	// Make sure all the servers are represented in the Raft config,
	// and that there are no extras.
	verifyRaft := func(c raft.Configuration) error {
		want := make(map[raft.ServerID]bool)
		for _, s := range servers {
			want[s.config.RaftConfig.LocalID] = true
		}

		for _, s := range c.Servers {
			if !want[s.ID] {
				return fmt.Errorf("don't want %q", s.ID)
			}
			delete(want, s.ID)
		}

		if len(want) > 0 {
			return fmt.Errorf("didn't find %v", want)
		}
		return nil
	}

	for _, s := range servers {
		future := s.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			return err
		}
		if err := verifyRaft(future.Configuration()); err != nil {
			return err
		}
	}
	return nil
}

func TestAutopilot_CleanupDeadServer(t *testing.T) {
	t.Parallel()
	for i := 1; i <= 3; i++ {
		testCleanupDeadServer(t, i)
	}
}

func testCleanupDeadServer(t *testing.T, raftVersion int) {
	conf := func(c *Config) {
		c.DevDisableBootstrap = true
		c.BootstrapExpect = 3
		c.RaftConfig.ProtocolVersion = raft.ProtocolVersion(raftVersion)
	}
	s1 := TestServer(t, conf)
	defer s1.Shutdown()

	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	servers := []*Server{s1, s2, s3}

	// Try to join
	TestJoin(t, s1, s2, s3)

	for _, s := range servers {
		retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s, 3)) })
	}

	// Bring up a new server
	s4 := TestServer(t, conf)
	defer s4.Shutdown()

	// Kill a non-leader server
	s3.Shutdown()
	retry.Run(t, func(r *retry.R) {
		alive := 0
		for _, m := range s1.Members() {
			if m.Status == serf.StatusAlive {
				alive++
			}
		}
		if alive != 2 {
			r.Fatal(nil)
		}
	})

	// Join the new server
	TestJoin(t, s1, s4)
	servers[2] = s4

	// Make sure the dead server is removed and we're back to 3 total peers
	for _, s := range servers {
		retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s, 3)) })
	}
}

func TestAutopilot_CleanupDeadServerPeriodic(t *testing.T) {
	t.Parallel()
	s1 := TestServer(t, nil)
	defer s1.Shutdown()

	conf := func(c *Config) {
		c.DevDisableBootstrap = true
	}

	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	s4 := TestServer(t, conf)
	defer s4.Shutdown()

	s5 := TestServer(t, conf)
	defer s5.Shutdown()

	servers := []*Server{s1, s2, s3, s4, s5}

	// Join the servers to s1, and wait until they are all promoted to
	// voters.
	TestJoin(t, s1, servers[1:]...)
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 5))
		}
	})

	// Kill a non-leader server
	s4.Shutdown()

	// Should be removed from the peers automatically
	servers = []*Server{s1, s2, s3, s5}
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 4))
		}
	})
}

func TestAutopilot_RollingUpdate(t *testing.T) {
	t.Parallel()
	s1 := TestServer(t, func(c *Config) {
		c.RaftConfig.ProtocolVersion = 3
	})
	defer s1.Shutdown()

	conf := func(c *Config) {
		c.DevDisableBootstrap = true
		c.RaftConfig.ProtocolVersion = 3
	}

	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	// Join the servers to s1, and wait until they are all promoted to
	// voters.
	servers := []*Server{s1, s2, s3}
	TestJoin(t, s1, s2, s3)
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 3))
		}
	})

	// Add one more server like we are doing a rolling update.
	s4 := TestServer(t, conf)
	defer s4.Shutdown()
	TestJoin(t, s1, s4)
	servers = append(servers, s4)
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 3))
		}
	})

	// Now kill one of the "old" nodes like we are doing a rolling update.
	s3.Shutdown()

	isVoter := func() bool {
		future := s1.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			t.Fatalf("err: %v", err)
		}
		for _, s := range future.Configuration().Servers {
			if string(s.ID) == string(s4.config.NodeID) {
				return s.Suffrage == raft.Voter
			}
		}
		t.Fatalf("didn't find s4")
		return false
	}

	// Wait for s4 to stabilize, get promoted to a voter, and for s3 to be
	// removed.
	servers = []*Server{s1, s2, s4}
	retry.Run(t, func(r *retry.R) {
		r.Check(wantRaft(servers))
		for _, s := range servers {
			r.Check(wantPeers(s, 3))
		}
		if !isVoter() {
			r.Fatalf("should be a voter")
		}
	})
}

func TestAutopilot_CleanupStaleRaftServer(t *testing.T) {
	t.Skip("TestAutopilot_CleanupDeadServer is very flaky, removing it for now")

	t.Parallel()
	s1 := TestServer(t, nil)
	defer s1.Shutdown()

	conf := func(c *Config) {
		c.DevDisableBootstrap = true
	}
	s2 := TestServer(t, conf)
	defer s2.Shutdown()

	s3 := TestServer(t, conf)
	defer s3.Shutdown()

	s4 := TestServer(t, conf)
	defer s4.Shutdown()

	servers := []*Server{s1, s2, s3}

	// Join the servers to s1
	TestJoin(t, s1, s2, s3)

	leader := waitForStableLeadership(t, servers)

	// Add s4 to peers directly
	addr := fmt.Sprintf("127.0.0.1:%d", s4.config.RPCAddr.Port)
	future := leader.raft.AddVoter(raft.ServerID(s4.config.NodeID), raft.ServerAddress(addr), 0, 0)
	if err := future.Error(); err != nil {
		t.Fatal(err)
	}

	// Verify we have 4 peers
	peers, err := s1.numPeers()
	if err != nil {
		t.Fatal(err)
	}
	if peers != 4 {
		t.Fatalf("bad: %v", peers)
	}

	// Wait for s4 to be removed
	for _, s := range []*Server{s1, s2, s3} {
		retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s, 3)) })
	}
}

func TestAutopilot_PromoteNonVoter(t *testing.T) {
	t.Parallel()
	s1 := TestServer(t, func(c *Config) {
		c.RaftConfig.ProtocolVersion = 3
	})
	defer s1.Shutdown()
	codec := rpcClient(t, s1)
	defer codec.Close()
	testutil.WaitForLeader(t, s1.RPC)

	s2 := TestServer(t, func(c *Config) {
		c.DevDisableBootstrap = true
		c.RaftConfig.ProtocolVersion = 3
	})
	defer s2.Shutdown()
	TestJoin(t, s1, s2)

	// Make sure we see it as a nonvoter initially. We wait until half
	// the stabilization period has passed.
	retry.Run(t, func(r *retry.R) {
		future := s1.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			r.Fatal(err)
		}

		servers := future.Configuration().Servers
		if len(servers) != 2 {
			r.Fatalf("bad: %v", servers)
		}
		if servers[1].Suffrage != raft.Nonvoter {
			r.Fatalf("bad: %v", servers)
		}
		health := s1.autopilot.GetServerHealth(string(servers[1].ID))
		if health == nil {
			r.Fatalf("nil health, %v", s1.autopilot.GetClusterHealth())
		}
		if !health.Healthy {
			r.Fatalf("bad: %v", health)
		}
		if time.Since(health.StableSince) < s1.config.AutopilotConfig.ServerStabilizationTime/2 {
			r.Fatal("stable period not elapsed")
		}
	})

	// Make sure it ends up as a voter.
	retry.Run(t, func(r *retry.R) {
		future := s1.raft.GetConfiguration()
		if err := future.Error(); err != nil {
			r.Fatal(err)
		}

		servers := future.Configuration().Servers
		if len(servers) != 2 {
			r.Fatalf("bad: %v", servers)
		}
		if servers[1].Suffrage != raft.Voter {
			r.Fatalf("bad: %v", servers)
		}
	})
}
<|endoftext|>


<commit_before>package lib

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/stats"
	"github.com/robertkrimen/otto"
	"gopkg.in/guregu/null.v3"
	"strconv"
	"sync"
	"time"
)

const (
	TickRate          = 1 * time.Millisecond
	ThresholdTickRate = 2 * time.Second
)

var (
	MetricVUs    = &stats.Metric{Name: "vus", Type: stats.Gauge}
	MetricVUsMax = &stats.Metric{Name: "vus_max", Type: stats.Gauge}
	MetricErrors = &stats.Metric{Name: "errors", Type: stats.Counter}

	ErrTooManyVUs = errors.New("More VUs than the maximum requested")
	ErrMaxTooLow  = errors.New("Can't lower max below current VU count")

	// Special error used to taint a test, without printing an error.
	ErrVUWantsTaint = errors.New("[ErrVUWantsTaint is never logged]")
)

type vuEntry struct {
	VU     VU
	Buffer []stats.Sample
	Cancel context.CancelFunc
}

type Engine struct {
	Runner    Runner
	Status    Status
	Stages    []Stage
	Collector stats.Collector
	Pause     sync.WaitGroup
	Metrics   map[*stats.Metric]stats.Sink

	thresholds  map[string][]Threshold
	thresholdVM *otto.Otto

	ctx    context.Context
	vus    []*vuEntry
	nextID int64

	vuMutex sync.Mutex
}

func NewEngine(r Runner) (*Engine, error) {
	e := &Engine{
		Runner: r,
		Status: Status{
			Running: null.BoolFrom(false),
			Tainted: null.BoolFrom(false),
			VUs:     null.IntFrom(0),
			VUsMax:  null.IntFrom(0),
			AtTime:  null.IntFrom(0),
		},
		Metrics:     make(map[*stats.Metric]stats.Sink),
		thresholds:  make(map[string][]Threshold),
		thresholdVM: otto.New(),
	}
	e.Pause.Add(1)

	return e, nil
}

func (e *Engine) Apply(opts Options) error {
	if opts.Run.Valid {
		e.SetRunning(opts.Run.Bool)
	}
	if opts.VUsMax.Valid {
		if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.VUs.Valid {
		if !opts.VUsMax.Valid {
			if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
				return err
			}
		}
		if err := e.SetVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.Duration.Valid {
		duration, err := time.ParseDuration(opts.Duration.String)
		if err != nil {
			return err
		}
		e.Stages = []Stage{Stage{Duration: null.IntFrom(int64(duration))}}
	}

	if opts.Quit.Valid {
		e.Status.Quit = opts.Quit
	}
	if opts.QuitOnTaint.Valid {
		e.Status.QuitOnTaint = opts.QuitOnTaint
	}

	if opts.Thresholds != nil {
		e.thresholds = opts.Thresholds

		// Make sure all scripts are compiled!
		for m, scripts := range e.thresholds {
			for i, script := range scripts {
				if script.Script != nil {
					continue
				}

				s, err := e.thresholdVM.Compile(fmt.Sprintf("threshold$%s:%i", m, i), script.Source)
				if err != nil {
					return err
				}
				script.Script = s
				scripts[i] = script
			}
		}
	}

	return nil
}

func (e *Engine) Run(ctx context.Context, opts Options) error {
	e.ctx = ctx
	e.nextID = 1

	e.Apply(opts)

	collectorCtx, collectorC := context.WithCancel(context.Background())
	if e.Collector != nil {
		go e.Collector.Run(collectorCtx)
		defer collectorC()
	} else {
		log.Debug("Engine: No Collector")
	}

	go e.runThresholds(ctx)

	e.consumeEngineStats()

	ticker := time.NewTicker(TickRate)
	lastTick := time.Now()

loop:
	for {
		select {
		case now := <-ticker.C:
			timeDelta := now.Sub(lastTick)
			e.Status.AtTime.Int64 += int64(timeDelta)
			lastTick = now

			stage, left, ok := StageAt(e.Stages, time.Duration(e.Status.AtTime.Int64))
			if stage.StartVUs.Valid && stage.EndVUs.Valid {
				progress := (float64(stage.Duration.Int64-int64(left)) / float64(stage.Duration.Int64))
				vus := Lerp(stage.StartVUs.Int64, stage.EndVUs.Int64, progress)
				e.SetVUs(vus)
			}

			for _, vu := range e.vus {
				e.consumeBuffer(vu.Buffer)
			}

			if !ok {
				e.SetRunning(false)

				if e.Status.Quit.Bool {
					break loop
				} else {
					log.Info("Test finished, press Ctrl+C to exit")
					<-ctx.Done()
					break loop
				}
			}

			if e.Status.QuitOnTaint.Bool && e.Status.Tainted.Bool {
				log.Warn("Test tainted, ending early...")
				break loop
			}

			e.consumeEngineStats()
		case <-ctx.Done():
			break loop
		}
	}

	e.vus = nil

	e.Status.Running = null.BoolFrom(false)
	e.Status.VUs = null.IntFrom(0)
	e.Status.VUsMax = null.IntFrom(0)
	e.consumeEngineStats()

	return nil
}

func (e *Engine) IsRunning() bool {
	return e.ctx != nil
}

func (e *Engine) SetRunning(running bool) {
	if running && !e.Status.Running.Bool {
		e.Pause.Done()
		log.Debug("Engine Unpaused")
	} else if !running && e.Status.Running.Bool {
		e.Pause.Add(1)
		log.Debug("Engine Paused")
	}
	e.Status.Running.Bool = running
}

func (e *Engine) SetVUs(v int64) error {
	if e.Status.VUs.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUs.Int64, "to": v}).Debug("Setting VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v > e.Status.VUsMax.Int64 {
		return ErrTooManyVUs
	}

	current := e.Status.VUs.Int64
	for i := current; i < v; i++ {
		entry := e.vus[i]
		if entry.Cancel != nil {
			panic(errors.New("ATTEMPTED TO RESCHEDULE RUNNING VU"))
		}

		ctx, cancel := context.WithCancel(e.ctx)
		entry.Cancel = cancel

		if err := entry.VU.Reconfigure(e.nextID); err != nil {
			return err
		}
		go e.runVU(ctx, e.nextID, entry)
		e.nextID++
	}
	for i := current - 1; i >= v; i-- {
		entry := e.vus[i]
		entry.Cancel()
		entry.Cancel = nil
	}

	e.Status.VUs.Int64 = v
	return nil
}

func (e *Engine) SetMaxVUs(v int64) error {
	if e.Status.VUsMax.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUsMax.Int64, "to": v}).Debug("Setting Max VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v < e.Status.VUs.Int64 {
		return ErrMaxTooLow
	}

	current := e.Status.VUsMax.Int64
	if v > current {
		vus := e.vus
		for i := current; i < v; i++ {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			entry := &vuEntry{VU: vu}
			vus = append(vus, entry)
		}
		e.vus = vus
	} else if v < current {
		e.vus = e.vus[:v]
	}

	e.Status.VUsMax.Int64 = v
	return nil
}

func (e *Engine) runVU(ctx context.Context, id int64, vu *vuEntry) {
	idString := strconv.FormatInt(id, 10)

waitForPause:
	e.Pause.Wait()

	for {
		samples, err := vu.VU.RunOnce(ctx)

		// If the context is cancelled, the iteration is likely to produce erroneous output
		// due to cancelled HTTP requests and whatnot. Discard output from such runs.
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err != nil {
			e.Status.Tainted.Bool = true

			if err != ErrVUWantsTaint {
				if s, ok := err.(fmt.Stringer); ok {
					log.Error(s.String())
				} else {
					log.WithError(err).Error("Runtime Error")
				}

				samples = append(samples, stats.Sample{
					Metric: MetricErrors,
					Time:   time.Now(),
					Tags:   map[string]string{"vu": idString, "error": err.Error()},
					Value:  float64(1),
				})
			}
		}

		vu.Buffer = append(vu.Buffer, samples...)

		if !e.Status.Running.Bool {
			goto waitForPause
		}
	}
}

func (e *Engine) runThresholds(ctx context.Context) {
	ticker := time.NewTicker(ThresholdTickRate)
	for {
		select {
		case <-ticker.C:
			for m, sink := range e.Metrics {
				scripts, ok := e.thresholds[m.Name]
				if !ok {
					continue
				}

				sample := sink.Format()
				for key, value := range sample {
					if m.Contains == stats.Time {
						value = value / float64(time.Millisecond)
					}
					// log.WithFields(log.Fields{"k": key, "v": value}).Debug("setting threshold data")
					e.thresholdVM.Set(key, value)
				}

				taint := false
				for _, script := range scripts {
					v, err := e.thresholdVM.Run(script.Script)
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold Error")
						taint = true
						continue
					}
					// log.WithFields(log.Fields{"metric": m.Name, "v": v, "s": sample}).Debug("threshold tick")
					bV, err := v.ToBoolean()
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold result is invalid")
						taint = true
						continue
					}
					if !bV {
						taint = true
					}
				}

				for key, _ := range sample {
					e.thresholdVM.Set(key, otto.UndefinedValue())
				}

				if taint {
					m.Tainted = true
					e.Status.Tainted.Bool = true
				}
			}
		case <-ctx.Done():
			return
		}
	}
}

// Returns a value sink for a metric, created from the type if unavailable.
func (e *Engine) getSink(m *stats.Metric) stats.Sink {
	s, ok := e.Metrics[m]
	if !ok {
		switch m.Type {
		case stats.Counter:
			s = &stats.CounterSink{}
		case stats.Gauge:
			s = &stats.GaugeSink{}
		case stats.Trend:
			s = &stats.TrendSink{}
		case stats.Rate:
			s = &stats.RateSink{}
		}
		e.Metrics[m] = s
	}
	return s
}

func (e *Engine) consumeBuffer(buffer []stats.Sample) {
	for _, sample := range buffer {
		e.getSink(sample.Metric).Add(sample)
	}
	if e.Collector != nil {
		e.Collector.Collect(buffer)
	}
}

func (e *Engine) consumeEngineStats() {
	t := time.Now()
	e.consumeBuffer([]stats.Sample{
		stats.Sample{Metric: MetricVUs, Time: t, Value: float64(e.Status.VUs.Int64)},
		stats.Sample{Metric: MetricVUsMax, Time: t, Value: float64(e.Status.VUsMax.Int64)},
	})
}
<commit_msg>[fix] Don't re-commit past samples forever >_><commit_after>package lib

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/stats"
	"github.com/robertkrimen/otto"
	"gopkg.in/guregu/null.v3"
	"strconv"
	"sync"
	"time"
)

const (
	TickRate          = 1 * time.Millisecond
	ThresholdTickRate = 2 * time.Second
)

var (
	MetricVUs    = &stats.Metric{Name: "vus", Type: stats.Gauge}
	MetricVUsMax = &stats.Metric{Name: "vus_max", Type: stats.Gauge}
	MetricErrors = &stats.Metric{Name: "errors", Type: stats.Counter}

	ErrTooManyVUs = errors.New("More VUs than the maximum requested")
	ErrMaxTooLow  = errors.New("Can't lower max below current VU count")

	// Special error used to taint a test, without printing an error.
	ErrVUWantsTaint = errors.New("[ErrVUWantsTaint is never logged]")
)

type vuEntry struct {
	VU     VU
	Buffer []stats.Sample
	Cancel context.CancelFunc
}

type Engine struct {
	Runner    Runner
	Status    Status
	Stages    []Stage
	Collector stats.Collector
	Pause     sync.WaitGroup
	Metrics   map[*stats.Metric]stats.Sink

	thresholds  map[string][]Threshold
	thresholdVM *otto.Otto

	ctx    context.Context
	vus    []*vuEntry
	nextID int64

	vuMutex sync.Mutex
}

func NewEngine(r Runner) (*Engine, error) {
	e := &Engine{
		Runner: r,
		Status: Status{
			Running: null.BoolFrom(false),
			Tainted: null.BoolFrom(false),
			VUs:     null.IntFrom(0),
			VUsMax:  null.IntFrom(0),
			AtTime:  null.IntFrom(0),
		},
		Metrics:     make(map[*stats.Metric]stats.Sink),
		thresholds:  make(map[string][]Threshold),
		thresholdVM: otto.New(),
	}
	e.Pause.Add(1)

	return e, nil
}

func (e *Engine) Apply(opts Options) error {
	if opts.Run.Valid {
		e.SetRunning(opts.Run.Bool)
	}
	if opts.VUsMax.Valid {
		if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.VUs.Valid {
		if !opts.VUsMax.Valid {
			if err := e.SetMaxVUs(opts.VUs.Int64); err != nil {
				return err
			}
		}
		if err := e.SetVUs(opts.VUs.Int64); err != nil {
			return err
		}
	}
	if opts.Duration.Valid {
		duration, err := time.ParseDuration(opts.Duration.String)
		if err != nil {
			return err
		}
		e.Stages = []Stage{Stage{Duration: null.IntFrom(int64(duration))}}
	}

	if opts.Quit.Valid {
		e.Status.Quit = opts.Quit
	}
	if opts.QuitOnTaint.Valid {
		e.Status.QuitOnTaint = opts.QuitOnTaint
	}

	if opts.Thresholds != nil {
		e.thresholds = opts.Thresholds

		// Make sure all scripts are compiled!
		for m, scripts := range e.thresholds {
			for i, script := range scripts {
				if script.Script != nil {
					continue
				}

				s, err := e.thresholdVM.Compile(fmt.Sprintf("threshold$%s:%i", m, i), script.Source)
				if err != nil {
					return err
				}
				script.Script = s
				scripts[i] = script
			}
		}
	}

	return nil
}

func (e *Engine) Run(ctx context.Context, opts Options) error {
	e.ctx = ctx
	e.nextID = 1

	e.Apply(opts)

	collectorCtx, collectorC := context.WithCancel(context.Background())
	if e.Collector != nil {
		go e.Collector.Run(collectorCtx)
		defer collectorC()
	} else {
		log.Debug("Engine: No Collector")
	}

	go e.runThresholds(ctx)

	e.consumeEngineStats()

	ticker := time.NewTicker(TickRate)
	lastTick := time.Now()

loop:
	for {
		select {
		case now := <-ticker.C:
			timeDelta := now.Sub(lastTick)
			e.Status.AtTime.Int64 += int64(timeDelta)
			lastTick = now

			stage, left, ok := StageAt(e.Stages, time.Duration(e.Status.AtTime.Int64))
			if stage.StartVUs.Valid && stage.EndVUs.Valid {
				progress := (float64(stage.Duration.Int64-int64(left)) / float64(stage.Duration.Int64))
				vus := Lerp(stage.StartVUs.Int64, stage.EndVUs.Int64, progress)
				e.SetVUs(vus)
			}

			for _, vu := range e.vus {
				buffer := vu.Buffer
				vu.Buffer = nil
				e.consumeBuffer(buffer)
			}

			if !ok {
				e.SetRunning(false)

				if e.Status.Quit.Bool {
					break loop
				} else {
					log.Info("Test finished, press Ctrl+C to exit")
					<-ctx.Done()
					break loop
				}
			}

			if e.Status.QuitOnTaint.Bool && e.Status.Tainted.Bool {
				log.Warn("Test tainted, ending early...")
				break loop
			}

			e.consumeEngineStats()
		case <-ctx.Done():
			break loop
		}
	}

	e.vus = nil

	e.Status.Running = null.BoolFrom(false)
	e.Status.VUs = null.IntFrom(0)
	e.Status.VUsMax = null.IntFrom(0)
	e.consumeEngineStats()

	return nil
}

func (e *Engine) IsRunning() bool {
	return e.ctx != nil
}

func (e *Engine) SetRunning(running bool) {
	if running && !e.Status.Running.Bool {
		e.Pause.Done()
		log.Debug("Engine Unpaused")
	} else if !running && e.Status.Running.Bool {
		e.Pause.Add(1)
		log.Debug("Engine Paused")
	}
	e.Status.Running.Bool = running
}

func (e *Engine) SetVUs(v int64) error {
	if e.Status.VUs.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUs.Int64, "to": v}).Debug("Setting VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v > e.Status.VUsMax.Int64 {
		return ErrTooManyVUs
	}

	current := e.Status.VUs.Int64
	for i := current; i < v; i++ {
		entry := e.vus[i]
		if entry.Cancel != nil {
			panic(errors.New("ATTEMPTED TO RESCHEDULE RUNNING VU"))
		}

		ctx, cancel := context.WithCancel(e.ctx)
		entry.Cancel = cancel

		if err := entry.VU.Reconfigure(e.nextID); err != nil {
			return err
		}
		go e.runVU(ctx, e.nextID, entry)
		e.nextID++
	}
	for i := current - 1; i >= v; i-- {
		entry := e.vus[i]
		entry.Cancel()
		entry.Cancel = nil
	}

	e.Status.VUs.Int64 = v
	return nil
}

func (e *Engine) SetMaxVUs(v int64) error {
	if e.Status.VUsMax.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUsMax.Int64, "to": v}).Debug("Setting Max VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v < e.Status.VUs.Int64 {
		return ErrMaxTooLow
	}

	current := e.Status.VUsMax.Int64
	if v > current {
		vus := e.vus
		for i := current; i < v; i++ {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			entry := &vuEntry{VU: vu}
			vus = append(vus, entry)
		}
		e.vus = vus
	} else if v < current {
		e.vus = e.vus[:v]
	}

	e.Status.VUsMax.Int64 = v
	return nil
}

func (e *Engine) runVU(ctx context.Context, id int64, vu *vuEntry) {
	idString := strconv.FormatInt(id, 10)

waitForPause:
	e.Pause.Wait()

	for {
		samples, err := vu.VU.RunOnce(ctx)

		// If the context is cancelled, the iteration is likely to produce erroneous output
		// due to cancelled HTTP requests and whatnot. Discard output from such runs.
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err != nil {
			e.Status.Tainted.Bool = true

			if err != ErrVUWantsTaint {
				if s, ok := err.(fmt.Stringer); ok {
					log.Error(s.String())
				} else {
					log.WithError(err).Error("Runtime Error")
				}

				samples = append(samples, stats.Sample{
					Metric: MetricErrors,
					Time:   time.Now(),
					Tags:   map[string]string{"vu": idString, "error": err.Error()},
					Value:  float64(1),
				})
			}
		}

		vu.Buffer = append(vu.Buffer, samples...)

		if !e.Status.Running.Bool {
			goto waitForPause
		}
	}
}

func (e *Engine) runThresholds(ctx context.Context) {
	ticker := time.NewTicker(ThresholdTickRate)
	for {
		select {
		case <-ticker.C:
			for m, sink := range e.Metrics {
				scripts, ok := e.thresholds[m.Name]
				if !ok {
					continue
				}

				sample := sink.Format()
				for key, value := range sample {
					if m.Contains == stats.Time {
						value = value / float64(time.Millisecond)
					}
					// log.WithFields(log.Fields{"k": key, "v": value}).Debug("setting threshold data")
					e.thresholdVM.Set(key, value)
				}

				taint := false
				for _, script := range scripts {
					v, err := e.thresholdVM.Run(script.Script)
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold Error")
						taint = true
						continue
					}
					// log.WithFields(log.Fields{"metric": m.Name, "v": v, "s": sample}).Debug("threshold tick")
					bV, err := v.ToBoolean()
					if err != nil {
						log.WithError(err).WithField("metric", m.Name).Error("Threshold result is invalid")
						taint = true
						continue
					}
					if !bV {
						taint = true
					}
				}

				for key, _ := range sample {
					e.thresholdVM.Set(key, otto.UndefinedValue())
				}

				if taint {
					m.Tainted = true
					e.Status.Tainted.Bool = true
				}
			}
		case <-ctx.Done():
			return
		}
	}
}

// Returns a value sink for a metric, created from the type if unavailable.
func (e *Engine) getSink(m *stats.Metric) stats.Sink {
	s, ok := e.Metrics[m]
	if !ok {
		switch m.Type {
		case stats.Counter:
			s = &stats.CounterSink{}
		case stats.Gauge:
			s = &stats.GaugeSink{}
		case stats.Trend:
			s = &stats.TrendSink{}
		case stats.Rate:
			s = &stats.RateSink{}
		}
		e.Metrics[m] = s
	}
	return s
}

func (e *Engine) consumeBuffer(buffer []stats.Sample) {
	for _, sample := range buffer {
		e.getSink(sample.Metric).Add(sample)
	}
	if e.Collector != nil {
		e.Collector.Collect(buffer)
	}
}

func (e *Engine) consumeEngineStats() {
	t := time.Now()
	e.consumeBuffer([]stats.Sample{
		stats.Sample{Metric: MetricVUs, Time: t, Value: float64(e.Status.VUs.Int64)},
		stats.Sample{Metric: MetricVUsMax, Time: t, Value: float64(e.Status.VUsMax.Int64)},
	})
}
<|endoftext|>


<commit_before>package gokeepasslib

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/xml"
	"io"
)

// Header to be put before xml content in kdbx file
var xmlHeader = []byte(`<?xml version="1.0" encoding="utf-8" standalone="yes"?>` + "\n")

// Encoder is used to automaticaly encrypt and write a database to a file, network, etc
type Encoder struct {
	w io.Writer
}

// NewEncoder creates a new encoder with writer w, identical to gokeepasslib.Encoder{w}
func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{w: w}
}

// Encode writes db to e's internal writer
func (e *Encoder) Encode(db *Database) (err error) {
	// Calculate transformed key to make HMAC and encrypt
	transformedKey, err := db.getTransformedKey()
	if err != nil {
		return err
	}

	// Write header then hashes before decode content (necessary to update HeaderHash)
	// db.Header writeTo will change its hash
	if err = db.Header.writeTo(e.w); err != nil {
		return err
	}

	// Update header hash into db.Hashes then write the data
	hash := db.Header.GetSha256()
	if db.Header.IsKdbx4() {
		db.Hashes.Sha256 = hash

		hmacKey := buildHmacKey(db, transformedKey)
		hmacHash := db.Header.GetHmacSha256(hmacKey)
		db.Hashes.Hmac = hmacHash

		if err = db.Hashes.writeTo(e.w); err != nil {
			return err
		}
	} else {
		db.Content.Meta.HeaderHash = base64.StdEncoding.EncodeToString(hash[:])
	}

	var rawContent []byte

	// Encode xml and append header to the top
	rawContent, err = xml.MarshalIndent(db.Content, "", "\t")
	if err != nil {
		return err
	}
	rawContent = append(xmlHeader, rawContent...)

	// Write InnerHeader (Kdbx v4)
	if db.Header.IsKdbx4() {
		var ih bytes.Buffer
		if err = db.Content.InnerHeader.writeTo(&ih); err != nil {
			return err
		}

		rawContent = append(ih.Bytes(), rawContent...)
	}

	// Encode raw content
	var encodedContent []byte
	encodedContent, err = encodeRawContent(db, rawContent, transformedKey)
	if err != nil {
		return err
	}

	// Writes the encrypted database content
	if _, err = e.w.Write(encodedContent); err != nil {
		return err
	}
	return nil
}

func encodeRawContent(db *Database, content []byte, transformedKey []byte) (encoded []byte, err error) {
	// Compress if the header compression flag is 1 (gzip)
	if db.Header.FileHeaders.CompressionFlags == GzipCompressionFlag {
		b := new(bytes.Buffer)
		w := gzip.NewWriter(b)

		if _, err = w.Write(content); err != nil {
			return encoded, err
		}

		// Close() needs to be explicitly called to write Gzip stream footer,
		// Flush() is not enough. some gzip decoders treat missing footer as error
		// while some don't). internally Close() also does flush.
		if err = w.Close(); err != nil {
			return encoded, err
		}

		content = b.Bytes()
	}

	// Compose blocks (Kdbx v3.1)
	if !db.Header.IsKdbx4() {
		var blocks bytes.Buffer
		composeContentBlocks31(&blocks, content)

		// Append blocks to StreamStartBytes
		content = append(db.Header.FileHeaders.StreamStartBytes, blocks.Bytes()...)
	}

	// Adds padding to data as required to encrypt properly
	if len(content)%16 != 0 {
		padding := make([]byte, 16-(len(content)%16))
		for i := 0; i < len(padding); i++ {
			padding[i] = byte(len(padding))
		}
		content = append(content, padding...)
	}

	// Encrypt content
	mode, err := db.Encrypter(transformedKey)
	if err != nil {
		return encoded, err
	}
	encrypted := make([]byte, len(content))
	mode.CryptBlocks(encrypted, content)

	// Compose blocks (Kdbx v4)
	if db.Header.IsKdbx4() {
		var blocks bytes.Buffer
		composeContentBlocks4(&blocks, encrypted, transformedKey)

		encrypted = blocks.Bytes()
	}
	return encrypted, nil
}
<commit_msg>golint fixes<commit_after>package gokeepasslib

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/xml"
	"io"
)

// Header to be put before xml content in kdbx file
var xmlHeader = []byte(`<?xml version="1.0" encoding="utf-8" standalone="yes"?>` + "\n")

// Encoder is used to automaticaly encrypt and write a database to a file, network, etc
type Encoder struct {
	w io.Writer
}

// NewEncoder creates a new encoder with writer w, identical to gokeepasslib.Encoder{w}
func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{w: w}
}

// Encode writes db to e's internal writer
func (e *Encoder) Encode(db *Database) error {
	// Calculate transformed key to make HMAC and encrypt
	transformedKey, err := db.getTransformedKey()
	if err != nil {
		return err
	}

	// Write header then hashes before decode content (necessary to update HeaderHash)
	// db.Header writeTo will change its hash
	if err = db.Header.writeTo(e.w); err != nil {
		return err
	}

	// Update header hash into db.Hashes then write the data
	hash := db.Header.GetSha256()
	if db.Header.IsKdbx4() {
		db.Hashes.Sha256 = hash

		hmacKey := buildHmacKey(db, transformedKey)
		hmacHash := db.Header.GetHmacSha256(hmacKey)
		db.Hashes.Hmac = hmacHash

		if err = db.Hashes.writeTo(e.w); err != nil {
			return err
		}
	} else {
		db.Content.Meta.HeaderHash = base64.StdEncoding.EncodeToString(hash[:])
	}

	// Encode xml and append header to the top
	rawContent, err := xml.MarshalIndent(db.Content, "", "\t")
	if err != nil {
		return err
	}
	rawContent = append(xmlHeader, rawContent...)

	// Write InnerHeader (Kdbx v4)
	if db.Header.IsKdbx4() {
		var ih bytes.Buffer
		if err = db.Content.InnerHeader.writeTo(&ih); err != nil {
			return err
		}

		rawContent = append(ih.Bytes(), rawContent...)
	}

	// Encode raw content
	encodedContent, err := encodeRawContent(db, rawContent, transformedKey)
	if err != nil {
		return err
	}

	// Writes the encrypted database content
	_, err = e.w.Write(encodedContent)
	return err
}

func encodeRawContent(db *Database, content []byte, transformedKey []byte) (encoded []byte, err error) {
	// Compress if the header compression flag is 1 (gzip)
	if db.Header.FileHeaders.CompressionFlags == GzipCompressionFlag {
		b := new(bytes.Buffer)
		w := gzip.NewWriter(b)

		if _, err = w.Write(content); err != nil {
			return encoded, err
		}

		// Close() needs to be explicitly called to write Gzip stream footer,
		// Flush() is not enough. some gzip decoders treat missing footer as error
		// while some don't). internally Close() also does flush.
		if err = w.Close(); err != nil {
			return encoded, err
		}

		content = b.Bytes()
	}

	// Compose blocks (Kdbx v3.1)
	if !db.Header.IsKdbx4() {
		var blocks bytes.Buffer
		composeContentBlocks31(&blocks, content)

		// Append blocks to StreamStartBytes
		content = append(db.Header.FileHeaders.StreamStartBytes, blocks.Bytes()...)
	}

	// Adds padding to data as required to encrypt properly
	if len(content)%16 != 0 {
		padding := make([]byte, 16-(len(content)%16))
		for i := 0; i < len(padding); i++ {
			padding[i] = byte(len(padding))
		}
		content = append(content, padding...)
	}

	// Encrypt content
	mode, err := db.Encrypter(transformedKey)
	if err != nil {
		return encoded, err
	}
	encrypted := make([]byte, len(content))
	mode.CryptBlocks(encrypted, content)

	// Compose blocks (Kdbx v4)
	if db.Header.IsKdbx4() {
		var blocks bytes.Buffer
		composeContentBlocks4(&blocks, encrypted, transformedKey)

		encrypted = blocks.Bytes()
	}
	return encrypted, nil
}
<|endoftext|>


<commit_before>// An example using go-chi (https://github.com/go-chi/chi) middleware and context
package main

import (
	"context"
	"encoding/xml"
	"fmt"
	"github.com/go-chi/chi"
	"github.com/go-chi/chi/middleware"
	"github.com/terminalstatic/go-xsd-validate"
	"io/ioutil"
	"log"
	"net/http"
)

var xsdHandler *xsdvalidate.XsdHandler

func readBody(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)

		if err != nil {
			fmt.Fprintf(w, fmt.Sprintf("%s<error><![CDATA[%s]]></error>", xml.Header, err))
			return
		}

		ctx := context.WithValue(r.Context(), "body", body)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func validateBody(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		body := (ctx.Value("body").([]byte))

		err := xsdHandler.ValidateMem(body, xsdvalidate.ParsErrVerbose)
		if err != nil {
			fmt.Fprintf(w, fmt.Sprintf("%s<error><![CDATA[%s]]></error>", xml.Header, err))
			return
		}

		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func main() {
	addr := ":9999"
	xsdvalidate.Init()
	defer xsdvalidate.Cleanup()
	var err error

	xsdHandler, err = xsdvalidate.NewXsdHandlerUrl("address.xsd", xsdvalidate.ParsErrDefault)
	defer xsdHandler.Free()
	if err != nil {
		panic(err)
	}

	r := chi.NewRouter()
	r.Use(middleware.Logger)
	r.Use(readBody)
	r.Use(validateBody)

	r.Route("/", func(r chi.Router) {
		r.Post("/address", func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("content-type", "application/xml; charset=utf-8")
			ctx := r.Context()
			body := (ctx.Value("body").([]byte))
			fmt.Fprintf(w, fmt.Sprintf("%s<request-content><![CDATA[%s]]></request-content><no-error>No errors</no-error>", xml.Header, body))
		})
	})

	fmt.Printf("Starting http server on %s\n", addr)
	log.Fatal(http.ListenAndServe(addr, r))
}
<commit_msg>Updated remark<commit_after>// An example using go-chi (https://github.com/go-chi/chi) middleware and context
package main

import (
	"context"
	"encoding/xml"
	"fmt"
	"github.com/go-chi/chi"
	"github.com/go-chi/chi/middleware"
	"github.com/terminalstatic/go-xsd-validate"
	"io/ioutil"
	"log"
	"net/http"
)

var xsdHandler *xsdvalidate.XsdHandler

func readBody(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)

		if err != nil {
			fmt.Fprintf(w, fmt.Sprintf("%s<error><![CDATA[%s]]></error>", xml.Header, err))
			return
		}

		ctx := context.WithValue(r.Context(), "body", body)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func validateBody(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		body := (ctx.Value("body").([]byte))

		err := xsdHandler.ValidateMem(body, xsdvalidate.ParsErrVerbose)
		if err != nil {
			fmt.Fprintf(w, fmt.Sprintf("%s<error><![CDATA[%s]]></error>", xml.Header, err))
			return
		}

		next.ServeHTTP(w, r)
	})
}

func main() {
	addr := ":9999"
	xsdvalidate.Init()
	defer xsdvalidate.Cleanup()
	var err error

	xsdHandler, err = xsdvalidate.NewXsdHandlerUrl("address.xsd", xsdvalidate.ParsErrDefault)
	defer xsdHandler.Free()
	if err != nil {
		panic(err)
	}

	r := chi.NewRouter()
	r.Use(middleware.StripSlashes)
	r.Use(readBody)
	r.Use(validateBody)

	r.Route("/", func(r chi.Router) {
		r.Post("/address", func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("content-type", "application/xml; charset=utf-8")
			ctx := r.Context()
			body := (ctx.Value("body").([]byte))
			fmt.Fprintf(w, fmt.Sprintf("%s<request-content><![CDATA[%s]]></request-content><no-error>No errors</no-error>", xml.Header, body))
		})
	})

	fmt.Printf("Starting http server on %s\n", addr)
	log.Fatal(http.ListenAndServe(addr, r))
}
<|endoftext|>


<commit_before>/*
Copyright 2015 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 main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"time"

	broker "cloud-testenv-broker/broker"
)

var (
	// The "wrapper_" prefix should help us avoid flag name collisions.
	checkUrl = flag.String("wrapper_check_url", "",
		"The URL to check for the serving state of the wrapped emulator. "+
			"If unspecified, the value of --wrapper_resolved_target is used.")
	checkRegexp    = flag.String("wrapper_check_regexp", "", "If non-empty, the regular expression used to match content read from --check_url that indicates the emulator is serving")
	resolvedTarget = flag.String("wrapper_resolved_target", "",
		"The address the emulator can be resolved on. "+
			"If unspecified, and a '--port=<PORT>' argument is present in the emulator command, the value 'localhost:<PORT>' is used with that port value.")
	ruleIdFlag = flag.String("wrapper_rule_id", "", "The ResolvedRule id the wrapped emulator is registered as.")
)

func findPort(args []string) int {
	portRegexp, _ := regexp.Compile("^--port=(\\d+)$")
	for _, arg := range args {
		m := portRegexp.FindStringSubmatch(arg)
		if m != nil {
			port, err := strconv.Atoi(m[1])
			if err != nil {
				log.Fatalf("failed to convert to int: %s", m[1])
			}
			return port
		}
	}
	return -1
}

func checkServing() bool {
	resp, err := http.Get(*checkUrl)
	if err == nil {
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			if *checkRegexp == "" {
				log.Printf("check URL responded with 200")
				return true
			}
			content, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return false
			}
			matched, _ := regexp.MatchString(*checkRegexp, string(content))
			if matched {
				log.Printf("check URL has matching text: %s", content)
				return true
			}
		}
	}
	return false
}

func killEmulatorGroupAndExit(cmd *exec.Cmd, code *int) {
	log.Printf("Sending SIGINT to emulator subprocess(es)...")
	err := broker.KillProcessTree(cmd)
	if err != nil {
		log.Printf("failed to kill process tree: %v", err)
	}
	os.Exit(*code)
}

func main() {
	flag.Parse()
	if *ruleIdFlag == "" {
		log.Fatalf("--wrapper_rule_id not specified")
	}
	if len(flag.Args()) == 0 {
		log.Fatalf("emulator command not specified")
	}
	if *resolvedTarget == "" {
		port := findPort(flag.Args())
		if port == -1 {
			log.Fatalf("--wrapper_resolved_target not specified")
		}
		*resolvedTarget = fmt.Sprintf("localhost:%d", port)
	}
	if *checkUrl == "" {
		if strings.HasPrefix(*resolvedTarget, "http") {
			*checkUrl = *resolvedTarget
		} else {
			*checkUrl = fmt.Sprintf("http://%s", *resolvedTarget)
		}
		log.Printf("using check URL: %s", *checkUrl)
	}

	log.Printf("running emulator command: %s", flag.Args())
	cmd := exec.Command(flag.Arg(0), flag.Args()[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	// Use a session to group the child and its subprocesses, if any, for the
	// purpose of terminating them as a group.
	err := broker.StartProcessTree(cmd)
	if err != nil {
		log.Fatalf("failed to start emulator command: %v", err)
	}

	// Kill the wrapped process and any of its children on both normal exit and
	// exit due to a signal. Code following this must set exitCode and return
	// instead of calling os.Exit() directly.
	exitCode := 0
	die := make(chan os.Signal, 1)
	signal.Notify(die, os.Interrupt, os.Kill)
	go func() {
		<-die
		killEmulatorGroupAndExit(cmd, &exitCode)
	}()
	defer killEmulatorGroupAndExit(cmd, &exitCode)

	log.Printf("waiting for emulator serving state check to succeed...")
	for !checkServing() {
		time.Sleep(200 * time.Millisecond)
	}

	err = broker.RegisterWithBroker(*ruleIdFlag, *resolvedTarget, []string{}, 1*time.Second)
	if err != nil {
		exitCode = 1
		time.Sleep(time.Minute)
		return
	}

	// We run until the emulator exits.
	cmd.Wait()
}
<commit_msg>ported test<commit_after>/*
Copyright 2015 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 main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"time"

	broker "cloud-testenv-broker/broker"
)

var (
	// The "wrapper_" prefix should help us avoid flag name collisions.
	checkUrl = flag.String("wrapper_check_url", "",
		"The URL to check for the serving state of the wrapped emulator. "+
			"If unspecified, the value of --wrapper_resolved_target is used.")
	checkRegexp    = flag.String("wrapper_check_regexp", "", "If non-empty, the regular expression used to match content read from --check_url that indicates the emulator is serving")
	resolvedTarget = flag.String("wrapper_resolved_target", "",
		"The address the emulator can be resolved on. "+
			"If unspecified, and a '--port=<PORT>' argument is present in the emulator command, the value 'localhost:<PORT>' is used with that port value.")
	ruleIdFlag = flag.String("wrapper_rule_id", "", "The ResolvedRule id the wrapped emulator is registered as.")
)

func findPort(args []string) int {
	portRegexp, _ := regexp.Compile("^--port=(\\d+)$")
	for _, arg := range args {
		m := portRegexp.FindStringSubmatch(arg)
		if m != nil {
			port, err := strconv.Atoi(m[1])
			if err != nil {
				log.Fatalf("failed to convert to int: %s", m[1])
			}
			return port
		}
	}
	return -1
}

func checkServing() bool {
	resp, err := http.Get(*checkUrl)
	if err == nil {
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			if *checkRegexp == "" {
				log.Printf("check URL responded with 200")
				return true
			}
			content, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return false
			}
			matched, _ := regexp.MatchString(*checkRegexp, string(content))
			if matched {
				log.Printf("check URL has matching text: %s", content)
				return true
			}
		}
	}
	return false
}

func killEmulatorGroupAndExit(cmd *exec.Cmd, code *int) {
	log.Printf("Sending SIGINT to emulator subprocess(es)...")
	err := broker.KillProcessTree(cmd)
	if err != nil {
		log.Printf("failed to kill process tree: %v", err)
	}
	os.Exit(*code)
}

func main() {
	flag.Parse()
	if *ruleIdFlag == "" {
		log.Fatalf("--wrapper_rule_id not specified")
	}
	if len(flag.Args()) == 0 {
		log.Fatalf("emulator command not specified")
	}
	if *resolvedTarget == "" {
		port := findPort(flag.Args())
		if port == -1 {
			log.Fatalf("--wrapper_resolved_target not specified")
		}
		*resolvedTarget = fmt.Sprintf("localhost:%d", port)
	}
	if *checkUrl == "" {
		if strings.HasPrefix(*resolvedTarget, "http") {
			*checkUrl = *resolvedTarget
		} else {
			*checkUrl = fmt.Sprintf("http://%s", *resolvedTarget)
		}
		log.Printf("using check URL: %s", *checkUrl)
	}

	log.Printf("running emulator command: %s", flag.Args())
	cmd := exec.Command(flag.Arg(0), flag.Args()[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	// Use a session to group the child and its subprocesses, if any, for the
	// purpose of terminating them as a group.
	err := broker.StartProcessTree(cmd)
	if err != nil {
		log.Fatalf("failed to start emulator command: %v", err)
	}

	// Kill the wrapped process and any of its children on both normal exit and
	// exit due to a signal. Code following this must set exitCode and return
	// instead of calling os.Exit() directly.
	exitCode := 0
	die := make(chan os.Signal, 1)
	signal.Notify(die, os.Interrupt, os.Kill)
	go func() {
		<-die
		killEmulatorGroupAndExit(cmd, &exitCode)
	}()
	defer killEmulatorGroupAndExit(cmd, &exitCode)

	log.Printf("waiting for emulator serving state check to succeed...")
	for !checkServing() {
		time.Sleep(200 * time.Millisecond)
	}

	_, err = broker.RegisterWithBroker(*ruleIdFlag, *resolvedTarget, []string{}, 1*time.Second)
	if err != nil {
		exitCode = 1
		time.Sleep(time.Minute)
		return
	}

	// We run until the emulator exits.
	cmd.Wait()
}
<|endoftext|>


<commit_before>package vpki

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"time"
)

var (
	day  = 24 * time.Hour
	year = 365 * day

	//DefaultTTL is the default TTL the library will request for certificates
	DefaultTTL = day
	//DefaultStrength is the default strength of RSA keys generated
	DefaultStrength = 2048
)

// SNICertifier abstracts the basic GetCertificate method used in TLSOpts, and
// also implemented by libraries like rsc.io/letsencrypt
type SNICertifier interface {
	GetCertificate(*tls.ClientHelloInfo) (*tls.Certificate, error)
}

// ListenAndServeTLS mostly mirrors the http.ListenAndServeTLS API, but
// generates the certificates for the server automatically via vault, with a
// short TTL. The function only needs an additional Certifier parameter which
// can generate signed certificates in order to work properly.
func ListenAndServeTLS(addr string, handler http.Handler, crt Certifier) error {
	certs := newCertCache(crt)

	tl, err := tls.Listen("tcp", addr, &tls.Config{
		GetCertificate: func(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if crt, ok := crt.(SNICertifier); ok {
				return crt.GetCertificate(h)
			}
			if h.ServerName == "" {
				return nil, fmt.Errorf("Cannot generate certs without TLS SNI (no server name was indicated)")
			}
			return certs.get(h.ServerName)
		},
	})

	if err != nil {
		return err
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}

	return http.Serve(tl, handler)
}
<commit_msg>Add prometheus instrumentation for certificate errors<commit_after>package vpki

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

var (
	day  = 24 * time.Hour
	year = 365 * day

	//DefaultTTL is the default TTL the library will request for certificates
	DefaultTTL = day
	//DefaultStrength is the default strength of RSA keys generated
	DefaultStrength = 2048

	promVPKICertError = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "vpki_cert_error",
		Help: "vpki encountered an error obtaining a certificate",
	}, []string{"server_name"})
)

func init() {
	prometheus.MustRegister(promVPKICertError)
}

// SNICertifier abstracts the basic GetCertificate method used in TLSOpts, and
// also implemented by libraries like rsc.io/letsencrypt
type SNICertifier interface {
	GetCertificate(*tls.ClientHelloInfo) (*tls.Certificate, error)
}

// ListenAndServeTLS mostly mirrors the http.ListenAndServeTLS API, but
// generates the certificates for the server automatically via vault, with a
// short TTL. The function only needs an additional Certifier parameter which
// can generate signed certificates in order to work properly.
func ListenAndServeTLS(addr string, handler http.Handler, crt Certifier) error {
	certs := newCertCache(crt)

	tl, err := tls.Listen("tcp", addr, &tls.Config{
		GetCertificate: func(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if crt, ok := crt.(SNICertifier); ok {
				return crt.GetCertificate(h)
			}
			pl := prometheus.Labels{"server_name": h.ServerName}
			if h.ServerName == "" {
				promVPKICertError.With(pl).Inc()
				return nil, fmt.Errorf("Cannot generate certs without TLS SNI (no server name was indicated)")
			}
			crt, err := certs.get(h.ServerName)
			if err != nil {
				promVPKICertError.With(pl).Inc()
			}
			return crt, err
		},
	})

	if err != nil {
		return err
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}

	return http.Serve(tl, handler)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"appengine"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"appengine/blobstore"
	"appengine/datastore"
	"appengine/taskqueue"
	"appengine/urlfetch"
	mpg "github.com/MiniProfiler/go/miniprofiler_gae"
	"github.com/mjibson/goon"
)

func ImportOpmlTask(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	userid := r.FormValue("user")
	bk := r.FormValue("key")
	fr := blobstore.NewReader(c, appengine.BlobKey(bk))
	data, err := ioutil.ReadAll(fr)
	if err != nil {
		return
	}

	var skip int
	if s, err := strconv.Atoi(r.FormValue("skip")); err == nil {
		skip = s
	}
	c.Debugf("reader import for %v, skip %v", userid, skip)

	var userOpml []*OpmlOutline
	remaining := skip

	var proc func(label string, outlines []*OpmlOutline)
	proc = func(label string, outlines []*OpmlOutline) {
		for _, o := range outlines {
			if o.XmlUrl != "" {
				if remaining > 0 {
					remaining--
				} else if len(userOpml) < IMPORT_LIMIT {
					userOpml = append(userOpml, &OpmlOutline{
						Title:   label,
						Outline: []*OpmlOutline{o},
					})
				}
			}

			if o.Title != "" && len(o.Outline) > 0 {
				proc(o.Title, o.Outline)
			}
		}
	}

	opml := Opml{}
	if err := xml.Unmarshal(data, &opml); err != nil {
		c.Errorf("opml error: %v", err.Error())
		return
	}
	proc("", opml.Outline)

	// todo: refactor below with similar from ImportReaderTask
	wg := sync.WaitGroup{}
	wg.Add(len(userOpml))
	for i := range userOpml {
		go func(i int) {
			o := userOpml[i].Outline[0]
			if err := addFeed(c, userid, userOpml[i]); err != nil {
				c.Warningf("opml import error: %v", err.Error())
				// todo: do something here?
			}
			c.Debugf("opml import: %s, %s", o.Title, o.XmlUrl)
			wg.Done()
		}(i)
	}
	wg.Wait()

	ud := UserData{Id: "data", Parent: gn.Key(&User{Id: userid})}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		gn.Get(&ud)
		mergeUserOpml(&ud, opml.Outline...)
		_, err := gn.Put(&ud)
		return err
	}, nil); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		c.Errorf("ude update error: %v", err.Error())
		return
	}

	if len(userOpml) == IMPORT_LIMIT {
		task := taskqueue.NewPOSTTask(routeUrl("import-opml-task"), url.Values{
			"key":  {bk},
			"user": {userid},
			"skip": {strconv.Itoa(skip + IMPORT_LIMIT)},
		})
		taskqueue.Add(c, task, "import-reader")
	}
}

const IMPORT_LIMIT = 10

func ImportReaderTask(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	userid := r.FormValue("user")
	bk := r.FormValue("key")
	fr := blobstore.NewReader(c, appengine.BlobKey(bk))
	data, err := ioutil.ReadAll(fr)
	if err != nil {
		return
	}

	var skip int
	if s, err := strconv.Atoi(r.FormValue("skip")); err == nil {
		skip = s
	}

	v := struct {
		Subscriptions []struct {
			Id         string `json:"id"`
			Title      string `json:"title"`
			HtmlUrl    string `json:"htmlUrl"`
			Categories []struct {
				Id    string `json:"id"`
				Label string `json:"label"`
			} `json:"categories"`
		} `json:"subscriptions"`
	}{}
	json.Unmarshal(data, &v)
	c.Debugf("reader import for %v, skip %v, len %v", userid, skip, len(v.Subscriptions))

	end := skip + IMPORT_LIMIT
	if end > len(v.Subscriptions) {
		end = len(v.Subscriptions)
	}

	wg := sync.WaitGroup{}
	wg.Add(end - skip)
	userOpml := make([]*OpmlOutline, end-skip)

	for i := range v.Subscriptions[skip:end] {
		go func(i int) {
			sub := v.Subscriptions[skip+i]
			var label string
			if len(sub.Categories) > 0 {
				label = sub.Categories[0].Label
			}
			outline := &OpmlOutline{
				Title: label,
				Outline: []*OpmlOutline{
					&OpmlOutline{
						XmlUrl: sub.Id[5:],
						Title:  sub.Title,
					},
				},
			}
			userOpml[i] = outline
			if err := addFeed(c, userid, outline); err != nil {
				c.Warningf("reader import error: %v", err.Error())
				// todo: do something here?
			}
			c.Debugf("reader import: %s, %s", sub.Title, sub.Id)
			wg.Done()
		}(i)
	}
	wg.Wait()

	ud := UserData{Id: "data", Parent: gn.Key(&User{Id: userid})}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		gn.Get(&ud)
		mergeUserOpml(&ud, userOpml...)
		_, err := gn.Put(&ud)
		return err
	}, nil); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		c.Errorf("ude update error: %v", err.Error())
		return
	}

	if end < len(v.Subscriptions) {
		task := taskqueue.NewPOSTTask(routeUrl("import-reader-task"), url.Values{
			"key":  {bk},
			"user": {userid},
			"skip": {strconv.Itoa(skip + IMPORT_LIMIT)},
		})
		taskqueue.Add(c, task, "import-reader")
	} else {
		blobstore.Delete(c, appengine.BlobKey(bk))
	}
}

func UpdateFeeds(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	q := datastore.NewQuery(gn.Key(&Feed{}).Kind()).KeysOnly()
	q = q.Filter("n <=", time.Now()).Limit(3000)
	it := gn.Run(q)
	i := 0
	for {
		k, err := it.Next(nil)
		if err == datastore.Done {
			break
		} else if err != nil {
			c.Errorf("next error: %v", err.Error())
			break
		}
		t := taskqueue.NewPOSTTask(routeUrl("update-feed"), url.Values{
			"feed": {k.StringID()},
		})
		if _, err := taskqueue.Add(c, t, "update-feed"); err != nil {
			c.Errorf("taskqueue error: %v", err.Error())
		}
		i++
	}
	c.Infof("updating %d feeds", i)
	fmt.Fprintf(w, "updating %d feeds", i)
}

func fetchFeed(c mpg.Context, origUrl, fetchUrl string) (*Feed, []*Story) {
	u, err := url.Parse(fetchUrl)
	if err == nil && u.Scheme == "" {
		u.Scheme = "http"
		origUrl = u.String()
		fetchUrl = origUrl
	}

	cl := &http.Client{
		Transport: &urlfetch.Transport{
			Context:  c,
			Deadline: time.Minute,
		},
	}
	if resp, err := cl.Get(fetchUrl); err == nil && resp.StatusCode == http.StatusOK {
		defer resp.Body.Close()
		b, _ := ioutil.ReadAll(resp.Body)
		if autoUrl, err := Autodiscover(b); err == nil && origUrl == fetchUrl {
			if autoU, err := url.Parse(autoUrl); err == nil {
				if autoU.Scheme == "" {
					autoU.Scheme = u.Scheme
				}
				if autoU.Host == "" {
					autoU.Host = u.Host
				}
				autoUrl = autoU.String()
			}
			return fetchFeed(c, origUrl, autoUrl)
		}
		return ParseFeed(c, origUrl, b)
	} else if err != nil {
		c.Warningf("fetch feed error: %s", err.Error())
	} else {
		c.Warningf("fetch feed error: status code: %s", resp.Status)
	}
	return nil, nil
}

func updateFeed(c mpg.Context, url string, feed *Feed, stories []*Story) error {
	gn := goon.FromContext(c)
	f := Feed{Url: url}
	if err := gn.Get(&f); err != nil {
		return errors.New(fmt.Sprintf("feed not found: %s", url))
	}

	// Compare the feed's listed update to the story's update.
	// Note: these may not be accurate, hence, only compare them to each other,
	// since they should have the same relative error.
	storyDate := f.Updated

	hasUpdated := !feed.Updated.IsZero()
	isFeedUpdated := f.Updated == feed.Updated
	if !hasUpdated {
		feed.Updated = f.Updated
	}
	feed.Date = f.Date
	f = *feed

	if hasUpdated && isFeedUpdated {
		c.Infof("feed %s already updated to %v, putting", url, feed.Updated)
		f.Updated = time.Now()
		gn.Put(&f)
		return nil
	}

	c.Debugf("hasUpdate: %v, isFeedUpdated: %v, storyDate: %v", hasUpdated, isFeedUpdated, storyDate)

	var newStories []*Story
	for _, s := range stories {
		if s.Updated.IsZero() || !s.Updated.Before(storyDate) {
			newStories = append(newStories, s)
		}
	}
	c.Debugf("%v possible stories to update", len(newStories))

	puts := []interface{}{&f}

	// find non existant stories
	fk := gn.Key(&f)
	getStories := make([]*Story, len(newStories))
	for i, s := range newStories {
		getStories[i] = &Story{Id: s.Id, Parent: fk}
	}
	err := gn.GetMulti(getStories)
	var updateStories []*Story
	for i, s := range getStories {
		if goon.NotFound(err, i) {
			updateStories = append(updateStories, newStories[i])
		} else if !newStories[i].Updated.IsZero() && !newStories[i].Updated.Equal(s.Updated) {
			newStories[i].Created = s.Created
			newStories[i].Published = s.Published
			updateStories = append(updateStories, newStories[i])
		}
	}
	c.Debugf("%v update stories", len(updateStories))

	for _, s := range updateStories {
		puts = append(puts, s)
		gn.Put(&StoryContent{
			Id:      1,
			Parent:  gn.Key(s),
			Content: s.content,
		})
	}

	c.Debugf("putting %v entities", len(puts))
	if len(puts) > 1 {
		f.Date = time.Now()
		if !hasUpdated {
			f.Updated = f.Date
		}
	}
	gn.PutMulti(puts)

	return nil
}

func UpdateFeed(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	url := r.FormValue("feed")
	c.Debugf("update feed %s", url)
	f := Feed{Url: url}
	if err := gn.Get(&f); err == datastore.ErrNoSuchEntity {
		return
	} else if time.Now().Before(f.NextUpdate) {
		c.Infof("feed %v already updated", url)
		return
	}
	if feed, stories := fetchFeed(c, url, url); feed != nil {
		updateFeed(c, url, feed, stories)
	} else {
		f.Errors++
		v := f.Errors + 1
		const max = 24 * 7
		if v > max {
			v = max
		} else if f.Errors == 1 {
			v = 0
		}
		f.NextUpdate = time.Now().Add(time.Hour * time.Duration(v))
		gn.Put(&f)
		c.Warningf("error with %v (%v), bump next update to %v", url, f.Errors, f.NextUpdate)
	}
}
<commit_msg>Move back to GetAll since we are using Limit<commit_after>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"appengine"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"appengine/blobstore"
	"appengine/datastore"
	"appengine/taskqueue"
	"appengine/urlfetch"
	mpg "github.com/MiniProfiler/go/miniprofiler_gae"
	"github.com/mjibson/goon"
)

func ImportOpmlTask(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	userid := r.FormValue("user")
	bk := r.FormValue("key")
	fr := blobstore.NewReader(c, appengine.BlobKey(bk))
	data, err := ioutil.ReadAll(fr)
	if err != nil {
		return
	}

	var skip int
	if s, err := strconv.Atoi(r.FormValue("skip")); err == nil {
		skip = s
	}
	c.Debugf("reader import for %v, skip %v", userid, skip)

	var userOpml []*OpmlOutline
	remaining := skip

	var proc func(label string, outlines []*OpmlOutline)
	proc = func(label string, outlines []*OpmlOutline) {
		for _, o := range outlines {
			if o.XmlUrl != "" {
				if remaining > 0 {
					remaining--
				} else if len(userOpml) < IMPORT_LIMIT {
					userOpml = append(userOpml, &OpmlOutline{
						Title:   label,
						Outline: []*OpmlOutline{o},
					})
				}
			}

			if o.Title != "" && len(o.Outline) > 0 {
				proc(o.Title, o.Outline)
			}
		}
	}

	opml := Opml{}
	if err := xml.Unmarshal(data, &opml); err != nil {
		c.Errorf("opml error: %v", err.Error())
		return
	}
	proc("", opml.Outline)

	// todo: refactor below with similar from ImportReaderTask
	wg := sync.WaitGroup{}
	wg.Add(len(userOpml))
	for i := range userOpml {
		go func(i int) {
			o := userOpml[i].Outline[0]
			if err := addFeed(c, userid, userOpml[i]); err != nil {
				c.Warningf("opml import error: %v", err.Error())
				// todo: do something here?
			}
			c.Debugf("opml import: %s, %s", o.Title, o.XmlUrl)
			wg.Done()
		}(i)
	}
	wg.Wait()

	ud := UserData{Id: "data", Parent: gn.Key(&User{Id: userid})}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		gn.Get(&ud)
		mergeUserOpml(&ud, opml.Outline...)
		_, err := gn.Put(&ud)
		return err
	}, nil); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		c.Errorf("ude update error: %v", err.Error())
		return
	}

	if len(userOpml) == IMPORT_LIMIT {
		task := taskqueue.NewPOSTTask(routeUrl("import-opml-task"), url.Values{
			"key":  {bk},
			"user": {userid},
			"skip": {strconv.Itoa(skip + IMPORT_LIMIT)},
		})
		taskqueue.Add(c, task, "import-reader")
	}
}

const IMPORT_LIMIT = 10

func ImportReaderTask(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	userid := r.FormValue("user")
	bk := r.FormValue("key")
	fr := blobstore.NewReader(c, appengine.BlobKey(bk))
	data, err := ioutil.ReadAll(fr)
	if err != nil {
		return
	}

	var skip int
	if s, err := strconv.Atoi(r.FormValue("skip")); err == nil {
		skip = s
	}

	v := struct {
		Subscriptions []struct {
			Id         string `json:"id"`
			Title      string `json:"title"`
			HtmlUrl    string `json:"htmlUrl"`
			Categories []struct {
				Id    string `json:"id"`
				Label string `json:"label"`
			} `json:"categories"`
		} `json:"subscriptions"`
	}{}
	json.Unmarshal(data, &v)
	c.Debugf("reader import for %v, skip %v, len %v", userid, skip, len(v.Subscriptions))

	end := skip + IMPORT_LIMIT
	if end > len(v.Subscriptions) {
		end = len(v.Subscriptions)
	}

	wg := sync.WaitGroup{}
	wg.Add(end - skip)
	userOpml := make([]*OpmlOutline, end-skip)

	for i := range v.Subscriptions[skip:end] {
		go func(i int) {
			sub := v.Subscriptions[skip+i]
			var label string
			if len(sub.Categories) > 0 {
				label = sub.Categories[0].Label
			}
			outline := &OpmlOutline{
				Title: label,
				Outline: []*OpmlOutline{
					&OpmlOutline{
						XmlUrl: sub.Id[5:],
						Title:  sub.Title,
					},
				},
			}
			userOpml[i] = outline
			if err := addFeed(c, userid, outline); err != nil {
				c.Warningf("reader import error: %v", err.Error())
				// todo: do something here?
			}
			c.Debugf("reader import: %s, %s", sub.Title, sub.Id)
			wg.Done()
		}(i)
	}
	wg.Wait()

	ud := UserData{Id: "data", Parent: gn.Key(&User{Id: userid})}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		gn.Get(&ud)
		mergeUserOpml(&ud, userOpml...)
		_, err := gn.Put(&ud)
		return err
	}, nil); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		c.Errorf("ude update error: %v", err.Error())
		return
	}

	if end < len(v.Subscriptions) {
		task := taskqueue.NewPOSTTask(routeUrl("import-reader-task"), url.Values{
			"key":  {bk},
			"user": {userid},
			"skip": {strconv.Itoa(skip + IMPORT_LIMIT)},
		})
		taskqueue.Add(c, task, "import-reader")
	} else {
		blobstore.Delete(c, appengine.BlobKey(bk))
	}
}

func UpdateFeeds(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	q := datastore.NewQuery(gn.Key(&Feed{}).Kind()).KeysOnly()
	q = q.Filter("n <=", time.Now()).Limit(3000)
	keys, _ := gn.GetAll(q, nil)
	for _, k := range keys {
		t := taskqueue.NewPOSTTask(routeUrl("update-feed"), url.Values{
			"feed": {k.StringID()},
		})
		if _, err := taskqueue.Add(c, t, "update-feed"); err != nil {
			c.Errorf("taskqueue error: %v", err.Error())
		}
	}
	c.Infof("updating %d feeds", len(keys))
	fmt.Fprintf(w, "updating %d feeds", len(keys))
}

func fetchFeed(c mpg.Context, origUrl, fetchUrl string) (*Feed, []*Story) {
	u, err := url.Parse(fetchUrl)
	if err == nil && u.Scheme == "" {
		u.Scheme = "http"
		origUrl = u.String()
		fetchUrl = origUrl
	}

	cl := &http.Client{
		Transport: &urlfetch.Transport{
			Context:  c,
			Deadline: time.Minute,
		},
	}
	if resp, err := cl.Get(fetchUrl); err == nil && resp.StatusCode == http.StatusOK {
		defer resp.Body.Close()
		b, _ := ioutil.ReadAll(resp.Body)
		if autoUrl, err := Autodiscover(b); err == nil && origUrl == fetchUrl {
			if autoU, err := url.Parse(autoUrl); err == nil {
				if autoU.Scheme == "" {
					autoU.Scheme = u.Scheme
				}
				if autoU.Host == "" {
					autoU.Host = u.Host
				}
				autoUrl = autoU.String()
			}
			return fetchFeed(c, origUrl, autoUrl)
		}
		return ParseFeed(c, origUrl, b)
	} else if err != nil {
		c.Warningf("fetch feed error: %s", err.Error())
	} else {
		c.Warningf("fetch feed error: status code: %s", resp.Status)
	}
	return nil, nil
}

func updateFeed(c mpg.Context, url string, feed *Feed, stories []*Story) error {
	gn := goon.FromContext(c)
	f := Feed{Url: url}
	if err := gn.Get(&f); err != nil {
		return errors.New(fmt.Sprintf("feed not found: %s", url))
	}

	// Compare the feed's listed update to the story's update.
	// Note: these may not be accurate, hence, only compare them to each other,
	// since they should have the same relative error.
	storyDate := f.Updated

	hasUpdated := !feed.Updated.IsZero()
	isFeedUpdated := f.Updated == feed.Updated
	if !hasUpdated {
		feed.Updated = f.Updated
	}
	feed.Date = f.Date
	f = *feed

	if hasUpdated && isFeedUpdated {
		c.Infof("feed %s already updated to %v, putting", url, feed.Updated)
		f.Updated = time.Now()
		gn.Put(&f)
		return nil
	}

	c.Debugf("hasUpdate: %v, isFeedUpdated: %v, storyDate: %v", hasUpdated, isFeedUpdated, storyDate)

	var newStories []*Story
	for _, s := range stories {
		if s.Updated.IsZero() || !s.Updated.Before(storyDate) {
			newStories = append(newStories, s)
		}
	}
	c.Debugf("%v possible stories to update", len(newStories))

	puts := []interface{}{&f}

	// find non existant stories
	fk := gn.Key(&f)
	getStories := make([]*Story, len(newStories))
	for i, s := range newStories {
		getStories[i] = &Story{Id: s.Id, Parent: fk}
	}
	err := gn.GetMulti(getStories)
	var updateStories []*Story
	for i, s := range getStories {
		if goon.NotFound(err, i) {
			updateStories = append(updateStories, newStories[i])
		} else if !newStories[i].Updated.IsZero() && !newStories[i].Updated.Equal(s.Updated) {
			newStories[i].Created = s.Created
			newStories[i].Published = s.Published
			updateStories = append(updateStories, newStories[i])
		}
	}
	c.Debugf("%v update stories", len(updateStories))

	for _, s := range updateStories {
		puts = append(puts, s)
		gn.Put(&StoryContent{
			Id:      1,
			Parent:  gn.Key(s),
			Content: s.content,
		})
	}

	c.Debugf("putting %v entities", len(puts))
	if len(puts) > 1 {
		f.Date = time.Now()
		if !hasUpdated {
			f.Updated = f.Date
		}
	}
	gn.PutMulti(puts)

	return nil
}

func UpdateFeed(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	gn := goon.FromContext(c)
	url := r.FormValue("feed")
	c.Debugf("update feed %s", url)
	f := Feed{Url: url}
	if err := gn.Get(&f); err == datastore.ErrNoSuchEntity {
		return
	} else if time.Now().Before(f.NextUpdate) {
		c.Infof("feed %v already updated", url)
		return
	}
	if feed, stories := fetchFeed(c, url, url); feed != nil {
		updateFeed(c, url, feed, stories)
	} else {
		f.Errors++
		v := f.Errors + 1
		const max = 24 * 7
		if v > max {
			v = max
		} else if f.Errors == 1 {
			v = 0
		}
		f.NextUpdate = time.Now().Add(time.Hour * time.Duration(v))
		gn.Put(&f)
		c.Warningf("error with %v (%v), bump next update to %v", url, f.Errors, f.NextUpdate)
	}
}
<|endoftext|>