<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
)

func Parse(r io.Reader, name string) (Prog, error) {
	p := &parser{
		r:    bufio.NewReader(r),
		name: name,
		npos: position{
			line: 1,
			col:  1,
		},
	}
	p.push(&p.prog.Stmts)
	p.next()
	p.program()
	return p.prog, p.err
}

type parser struct {
	r    *bufio.Reader
	name string

	err error

	ltok Token
	tok  Token
	lval string
	val  string

	lpos position
	pos  position
	npos position

	prog  Prog
	stack []interface{}
}

type position struct {
	line int
	col  int
}

var reserved = map[rune]bool{
	'\n': true,
	'&':  true,
	'>':  true,
	'<':  true,
	'|':  true,
	';':  true,
	'(':  true,
	')':  true,
}

// like reserved, but these are only reserved if at the start of a word
var starters = map[rune]bool{
	'{': true,
	'}': true,
	'#': true,
}

var space = map[rune]bool{
	' ':  true,
	'\t': true,
}

var quote = map[rune]bool{
	'"':  true,
	'\'': true,
	'`':  true,
}

var identRe = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)

func (p *parser) readRune() (rune, error) {
	r, _, err := p.r.ReadRune()
	if err != nil {
		if err == io.EOF {
			p.setEOF()
		} else {
			p.errPass(err)
		}
		return 0, err
	}
	p.moveWith(r)
	return r, nil
}

func (p *parser) moveWith(r rune) {
	if r == '\n' {
		p.npos.line++
		p.npos.col = 1
	} else {
		p.npos.col++
	}
}

func (p *parser) unreadRune() {
	if err := p.r.UnreadRune(); err != nil {
		panic(err)
	}
}

func (p *parser) readOnly(wanted rune) bool {
	r, _, err := p.r.ReadRune()
	if r == wanted {
		p.moveWith(r)
		return true
	}
	if err == nil {
		p.unreadRune()
	}
	return false
}

func (p *parser) next() {
	p.lpos = p.pos
	r := ' '
	for space[r] {
		var err error
		r, err = p.readRune()
		if err != nil {
			return
		}
	}
	p.pos = p.npos
	p.pos.col--
	switch {
	case r == '\\' && p.readOnly('\n'):
		p.next()
	case reserved[r] || starters[r]:
		switch r {
		case '#':
			p.advance(COMMENT, p.readLine())
		case '\n':
			p.advance('\n', "")
		default:
			p.advance(p.doToken(r), "")
		}
	default:
		p.advance(WORD, p.readWord(r))
	}
}

func (p *parser) readWord(r rune) string {
	var rs []rune
	var q rune
runeLoop:
	for {
		switch {
		case q != '\'' && r == '\\': // quoted rune
			rs = append(rs, r)
			r, _ = p.readRune()
		case q != '\'' && r == '$': // $ continuation
			switch {
			case p.readOnly('{'):
				rs = append(rs, '$', '{')
				rs = append(rs, p.readIncluding('}')...)
				r = '}'
			case p.readOnly('('):
				rs = append(rs, '$', '(')
				rs = append(rs, p.readIncluding(')')...)
				r = ')'
			}
		case q != 0: // rest of quoted cases
			if r == q {
				q = 0
			}
		case quote[r]: // start of a quoted string
			q = r
		case reserved[r] || space[r]: // end of word
			p.npos.col--
			p.unreadRune()
			break runeLoop
		}
		rs = append(rs, r)
		var err error
		r, err = p.readRune()
		if err != nil {
			if q != 0 {
				p.errWanted(Token(q))
			}
			break
		}
	}
	return string(rs)
}

func (p *parser) doToken(r rune) Token {
	switch r {
	case '&':
		if p.readOnly('&') {
			return LAND
		}
		return AND
	case '|':
		if p.readOnly('|') {
			return LOR
		}
		return OR
	case '(':
		return LPAREN
	case '{':
		return LBRACE
	case ')':
		return RPAREN
	case '}':
		return RBRACE
	case ';':
		if p.readOnly(';') {
			return DSEMICOLON
		}
		return SEMICOLON
	case '<':
		return LSS
	case '>':
		if p.readOnly('>') {
			return SHR
		}
		return GTR
	default:
		return Token(r)
	}
}

func (p *parser) advance(tok Token, val string) {
	if p.tok != EOF {
		p.ltok = p.tok
		p.lval = p.val
	}
	p.tok = tok
	p.val = val
}

func (p *parser) setEOF() {
	p.advance(EOF, "EOF")
}

func (p *parser) readUntil(delim rune) ([]rune, bool) {
	var rs []rune
	for {
		r, err := p.readRune()
		if err != nil {
			return rs, false
		}
		if r == delim {
			return rs, true
		}
		rs = append(rs, r)
	}
}

func (p *parser) readIncluding(delim rune) []rune {
	rs, found := p.readUntil(delim)
	if !found {
		p.errWanted(Token(delim))
	}
	return rs
}

func (p *parser) readLine() string {
	rs, _ := p.readUntil('\n')
	return string(rs)
}

// We can't simply have these as tokens as they can sometimes be valid
// words, e.g. `echo if`.
var reservedWords = map[Token]string{
	IF:    "if",
	THEN:  "then",
	ELIF:  "elif",
	ELSE:  "else",
	FI:    "fi",
	WHILE: "while",
	DO:    "do",
	DONE:  "done",
}

func (p *parser) peek(tok Token) bool {
	return p.tok == tok || (p.tok == WORD && p.val == reservedWords[tok])
}

func (p *parser) got(tok Token) bool {
	if p.peek(tok) {
		p.next()
		return true
	}
	return false
}

func (p *parser) want(tok Token) {
	if !p.peek(tok) {
		p.errWanted(tok)
		return
	}
	p.next()
}

func (p *parser) errPass(err error) {
	if p.err == nil {
		p.err = err
	}
	p.setEOF()
}

type lineErr struct {
	fname string
	pos   position
	text  string
}

func (e lineErr) Error() string {
	return fmt.Sprintf("%s:%d:%d: %s", e.fname, e.pos.line, e.pos.col, e.text)
}

func (p *parser) posErr(pos position, format string, v ...interface{}) {
	p.errPass(lineErr{
		fname: p.name,
		pos:   pos,
		text:  fmt.Sprintf(format, v...),
	})
}

func (p *parser) curErr(format string, v ...interface{}) {
	p.posErr(p.pos, format, v...)
}

func (p *parser) errWantedStr(s string) {
	if p.tok == EOF {
		p.pos = p.npos
	}
	p.curErr("unexpected token %s - wanted %s", p.tok, s)
}

func (p *parser) errWanted(tok Token) {
	p.errWantedStr(tok.String())
}

func (p *parser) errAfterStr(s string) {
	p.curErr("unexpected token %s after %s", p.tok, s)
}

func (p *parser) add(n Node) {
	cur := p.stack[len(p.stack)-1]
	switch x := cur.(type) {
	case *[]Node:
		*x = append(*x, n)
	case *Node:
		if *x != nil {
			panic("single node set twice")
		}
		*x = n
	default:
		panic("unknown type in the stack")
	}
}

func (p *parser) pop() {
	p.stack = p.stack[:len(p.stack)-1]
}

func (p *parser) push(v interface{}) {
	p.stack = append(p.stack, v)
}

func (p *parser) popAdd(n Node) {
	p.pop()
	p.add(n)
}

func (p *parser) program() {
	p.commands()
}

func (p *parser) commands(stop ...Token) (count int) {
	for p.tok != EOF {
		for _, tok := range stop {
			if p.peek(tok) {
				return
			}
		}
		p.command()
		count++
	}
	return
}

func (p *parser) command() {
	switch {
	case p.got(COMMENT):
		p.add(Comment{
			Text: p.lval,
		})
	case p.got('\n'), p.got(COMMENT):
		if p.tok != EOF {
			p.command()
		}
	case p.got(LPAREN):
		var sub Subshell
		p.push(&sub.Stmts)
		if p.commands(RPAREN) == 0 {
			p.errWantedStr("command")
		}
		p.want(RPAREN)
		p.popAdd(sub)
	case p.got(LBRACE):
		var bl Block
		p.push(&bl.Stmts)
		if p.commands(RBRACE) == 0 {
			p.errWantedStr("command")
		}
		p.want(RBRACE)
		p.popAdd(bl)
	case p.got(IF):
		var ifs IfStmt
		p.push(&ifs.Cond)
		p.command()
		p.pop()
		p.want(THEN)
		p.push(&ifs.ThenStmts)
		p.commands(FI, ELIF, ELSE)
		p.pop()
		p.push(&ifs.Elifs)
		for p.got(ELIF) {
			var elf Elif
			p.push(&elf.Cond)
			p.command()
			p.pop()
			p.want(THEN)
			p.push(&elf.ThenStmts)
			p.commands(FI, ELIF, ELSE)
			p.popAdd(elf)
		}
		if p.got(ELSE) {
			p.pop()
			p.push(&ifs.ElseStmts)
			p.commands(FI)
		}
		p.want(FI)
		p.popAdd(ifs)
	case p.got(WHILE):
		var whl WhileStmt
		p.push(&whl.Cond)
		p.command()
		p.pop()
		p.want(DO)
		p.push(&whl.DoStmts)
		p.commands(DONE)
		p.want(DONE)
		p.popAdd(whl)
	case p.got(WORD):
		var cmd Command
		p.push(&cmd.Args)
		p.add(Lit{Val: p.lval})
		fpos := p.lpos
		fval := p.lval
		if p.got(LPAREN) {
			p.want(RPAREN)
			if !identRe.MatchString(fval) {
				p.posErr(fpos, "invalid func name %q", fval)
			}
			fun := FuncDecl{
				Name: Lit{Val: fval},
			}
			p.push(&fun.Body)
			p.command()
			p.pop()
			p.popAdd(fun)
			return
		}
	args:
		for p.tok != EOF {
			switch {
			case p.got(WORD):
				p.add(Lit{Val: p.lval})
			case p.got(LAND):
				p.binaryExpr(LAND, cmd)
				return
			case p.got(OR):
				p.binaryExpr(OR, cmd)
				return
			case p.got(LOR):
				p.binaryExpr(LOR, cmd)
				return
			case p.gotRedirect():
			case p.got(AND):
				cmd.Background = true
				fallthrough
			case p.got(SEMICOLON), p.got('\n'):
				break args
			default:
				p.errAfterStr("command")
			}
		}
		p.popAdd(cmd)
	default:
		p.errWantedStr("command")
	}
}

func (p *parser) binaryExpr(op Token, left Node) {
	b := BinaryExpr{Op: op}
	p.push(&b.Y)
	p.command()
	p.pop()
	b.X = left
	p.popAdd(b)
}

func (p *parser) gotRedirect() bool {
	var r Redirect
	switch {
	case p.got(GTR):
		r.Op = GTR
		if p.got(AND) {
			p.want(WORD)
			r.Obj = Lit{Val: "&" + p.lval}
		} else {
			p.want(WORD)
			r.Obj = Lit{Val: p.lval}
		}
	case p.got(SHR):
		r.Op = SHR
		p.want(WORD)
		r.Obj = Lit{Val: p.lval}
	case p.got(LSS):
		r.Op = LSS
		p.want(WORD)
		r.Obj = Lit{Val: p.lval}
	default:
		return false
	}
	p.add(r)
	return true
}
<commit_msg>Wording fix<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
)

func Parse(r io.Reader, name string) (Prog, error) {
	p := &parser{
		r:    bufio.NewReader(r),
		name: name,
		npos: position{
			line: 1,
			col:  1,
		},
	}
	p.push(&p.prog.Stmts)
	p.next()
	p.program()
	return p.prog, p.err
}

type parser struct {
	r    *bufio.Reader
	name string

	err error

	ltok Token
	tok  Token
	lval string
	val  string

	lpos position
	pos  position
	npos position

	prog  Prog
	stack []interface{}
}

type position struct {
	line int
	col  int
}

var reserved = map[rune]bool{
	'\n': true,
	'&':  true,
	'>':  true,
	'<':  true,
	'|':  true,
	';':  true,
	'(':  true,
	')':  true,
}

// like reserved, but these are only reserved if at the start of a word
var starters = map[rune]bool{
	'{': true,
	'}': true,
	'#': true,
}

var space = map[rune]bool{
	' ':  true,
	'\t': true,
}

var quote = map[rune]bool{
	'"':  true,
	'\'': true,
	'`':  true,
}

var identRe = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)

func (p *parser) readRune() (rune, error) {
	r, _, err := p.r.ReadRune()
	if err != nil {
		if err == io.EOF {
			p.setEOF()
		} else {
			p.errPass(err)
		}
		return 0, err
	}
	p.moveWith(r)
	return r, nil
}

func (p *parser) moveWith(r rune) {
	if r == '\n' {
		p.npos.line++
		p.npos.col = 1
	} else {
		p.npos.col++
	}
}

func (p *parser) unreadRune() {
	if err := p.r.UnreadRune(); err != nil {
		panic(err)
	}
}

func (p *parser) readOnly(wanted rune) bool {
	r, _, err := p.r.ReadRune()
	if r == wanted {
		p.moveWith(r)
		return true
	}
	if err == nil {
		p.unreadRune()
	}
	return false
}

func (p *parser) next() {
	p.lpos = p.pos
	r := ' '
	for space[r] {
		var err error
		r, err = p.readRune()
		if err != nil {
			return
		}
	}
	p.pos = p.npos
	p.pos.col--
	switch {
	case r == '\\' && p.readOnly('\n'):
		p.next()
	case reserved[r] || starters[r]:
		switch r {
		case '#':
			p.advance(COMMENT, p.readLine())
		case '\n':
			p.advance('\n', "")
		default:
			p.advance(p.doToken(r), "")
		}
	default:
		p.advance(WORD, p.readWord(r))
	}
}

func (p *parser) readWord(r rune) string {
	var rs []rune
	var q rune
runeLoop:
	for {
		switch {
		case q != '\'' && r == '\\': // escaped rune
			rs = append(rs, r)
			r, _ = p.readRune()
		case q != '\'' && r == '$': // $ continuation
			switch {
			case p.readOnly('{'):
				rs = append(rs, '$', '{')
				rs = append(rs, p.readIncluding('}')...)
				r = '}'
			case p.readOnly('('):
				rs = append(rs, '$', '(')
				rs = append(rs, p.readIncluding(')')...)
				r = ')'
			}
		case q != 0: // rest of quoted cases
			if r == q {
				q = 0
			}
		case quote[r]: // start of a quoted string
			q = r
		case reserved[r] || space[r]: // end of word
			p.npos.col--
			p.unreadRune()
			break runeLoop
		}
		rs = append(rs, r)
		var err error
		r, err = p.readRune()
		if err != nil {
			if q != 0 {
				p.errWanted(Token(q))
			}
			break
		}
	}
	return string(rs)
}

func (p *parser) doToken(r rune) Token {
	switch r {
	case '&':
		if p.readOnly('&') {
			return LAND
		}
		return AND
	case '|':
		if p.readOnly('|') {
			return LOR
		}
		return OR
	case '(':
		return LPAREN
	case '{':
		return LBRACE
	case ')':
		return RPAREN
	case '}':
		return RBRACE
	case ';':
		if p.readOnly(';') {
			return DSEMICOLON
		}
		return SEMICOLON
	case '<':
		return LSS
	case '>':
		if p.readOnly('>') {
			return SHR
		}
		return GTR
	default:
		return Token(r)
	}
}

func (p *parser) advance(tok Token, val string) {
	if p.tok != EOF {
		p.ltok = p.tok
		p.lval = p.val
	}
	p.tok = tok
	p.val = val
}

func (p *parser) setEOF() {
	p.advance(EOF, "EOF")
}

func (p *parser) readUntil(delim rune) ([]rune, bool) {
	var rs []rune
	for {
		r, err := p.readRune()
		if err != nil {
			return rs, false
		}
		if r == delim {
			return rs, true
		}
		rs = append(rs, r)
	}
}

func (p *parser) readIncluding(delim rune) []rune {
	rs, found := p.readUntil(delim)
	if !found {
		p.errWanted(Token(delim))
	}
	return rs
}

func (p *parser) readLine() string {
	rs, _ := p.readUntil('\n')
	return string(rs)
}

// We can't simply have these as tokens as they can sometimes be valid
// words, e.g. `echo if`.
var reservedWords = map[Token]string{
	IF:    "if",
	THEN:  "then",
	ELIF:  "elif",
	ELSE:  "else",
	FI:    "fi",
	WHILE: "while",
	DO:    "do",
	DONE:  "done",
}

func (p *parser) peek(tok Token) bool {
	return p.tok == tok || (p.tok == WORD && p.val == reservedWords[tok])
}

func (p *parser) got(tok Token) bool {
	if p.peek(tok) {
		p.next()
		return true
	}
	return false
}

func (p *parser) want(tok Token) {
	if !p.peek(tok) {
		p.errWanted(tok)
		return
	}
	p.next()
}

func (p *parser) errPass(err error) {
	if p.err == nil {
		p.err = err
	}
	p.setEOF()
}

type lineErr struct {
	fname string
	pos   position
	text  string
}

func (e lineErr) Error() string {
	return fmt.Sprintf("%s:%d:%d: %s", e.fname, e.pos.line, e.pos.col, e.text)
}

func (p *parser) posErr(pos position, format string, v ...interface{}) {
	p.errPass(lineErr{
		fname: p.name,
		pos:   pos,
		text:  fmt.Sprintf(format, v...),
	})
}

func (p *parser) curErr(format string, v ...interface{}) {
	p.posErr(p.pos, format, v...)
}

func (p *parser) errWantedStr(s string) {
	if p.tok == EOF {
		p.pos = p.npos
	}
	p.curErr("unexpected token %s - wanted %s", p.tok, s)
}

func (p *parser) errWanted(tok Token) {
	p.errWantedStr(tok.String())
}

func (p *parser) errAfterStr(s string) {
	p.curErr("unexpected token %s after %s", p.tok, s)
}

func (p *parser) add(n Node) {
	cur := p.stack[len(p.stack)-1]
	switch x := cur.(type) {
	case *[]Node:
		*x = append(*x, n)
	case *Node:
		if *x != nil {
			panic("single node set twice")
		}
		*x = n
	default:
		panic("unknown type in the stack")
	}
}

func (p *parser) pop() {
	p.stack = p.stack[:len(p.stack)-1]
}

func (p *parser) push(v interface{}) {
	p.stack = append(p.stack, v)
}

func (p *parser) popAdd(n Node) {
	p.pop()
	p.add(n)
}

func (p *parser) program() {
	p.commands()
}

func (p *parser) commands(stop ...Token) (count int) {
	for p.tok != EOF {
		for _, tok := range stop {
			if p.peek(tok) {
				return
			}
		}
		p.command()
		count++
	}
	return
}

func (p *parser) command() {
	switch {
	case p.got(COMMENT):
		p.add(Comment{
			Text: p.lval,
		})
	case p.got('\n'), p.got(COMMENT):
		if p.tok != EOF {
			p.command()
		}
	case p.got(LPAREN):
		var sub Subshell
		p.push(&sub.Stmts)
		if p.commands(RPAREN) == 0 {
			p.errWantedStr("command")
		}
		p.want(RPAREN)
		p.popAdd(sub)
	case p.got(LBRACE):
		var bl Block
		p.push(&bl.Stmts)
		if p.commands(RBRACE) == 0 {
			p.errWantedStr("command")
		}
		p.want(RBRACE)
		p.popAdd(bl)
	case p.got(IF):
		var ifs IfStmt
		p.push(&ifs.Cond)
		p.command()
		p.pop()
		p.want(THEN)
		p.push(&ifs.ThenStmts)
		p.commands(FI, ELIF, ELSE)
		p.pop()
		p.push(&ifs.Elifs)
		for p.got(ELIF) {
			var elf Elif
			p.push(&elf.Cond)
			p.command()
			p.pop()
			p.want(THEN)
			p.push(&elf.ThenStmts)
			p.commands(FI, ELIF, ELSE)
			p.popAdd(elf)
		}
		if p.got(ELSE) {
			p.pop()
			p.push(&ifs.ElseStmts)
			p.commands(FI)
		}
		p.want(FI)
		p.popAdd(ifs)
	case p.got(WHILE):
		var whl WhileStmt
		p.push(&whl.Cond)
		p.command()
		p.pop()
		p.want(DO)
		p.push(&whl.DoStmts)
		p.commands(DONE)
		p.want(DONE)
		p.popAdd(whl)
	case p.got(WORD):
		var cmd Command
		p.push(&cmd.Args)
		p.add(Lit{Val: p.lval})
		fpos := p.lpos
		fval := p.lval
		if p.got(LPAREN) {
			p.want(RPAREN)
			if !identRe.MatchString(fval) {
				p.posErr(fpos, "invalid func name %q", fval)
			}
			fun := FuncDecl{
				Name: Lit{Val: fval},
			}
			p.push(&fun.Body)
			p.command()
			p.pop()
			p.popAdd(fun)
			return
		}
	args:
		for p.tok != EOF {
			switch {
			case p.got(WORD):
				p.add(Lit{Val: p.lval})
			case p.got(LAND):
				p.binaryExpr(LAND, cmd)
				return
			case p.got(OR):
				p.binaryExpr(OR, cmd)
				return
			case p.got(LOR):
				p.binaryExpr(LOR, cmd)
				return
			case p.gotRedirect():
			case p.got(AND):
				cmd.Background = true
				fallthrough
			case p.got(SEMICOLON), p.got('\n'):
				break args
			default:
				p.errAfterStr("command")
			}
		}
		p.popAdd(cmd)
	default:
		p.errWantedStr("command")
	}
}

func (p *parser) binaryExpr(op Token, left Node) {
	b := BinaryExpr{Op: op}
	p.push(&b.Y)
	p.command()
	p.pop()
	b.X = left
	p.popAdd(b)
}

func (p *parser) gotRedirect() bool {
	var r Redirect
	switch {
	case p.got(GTR):
		r.Op = GTR
		if p.got(AND) {
			p.want(WORD)
			r.Obj = Lit{Val: "&" + p.lval}
		} else {
			p.want(WORD)
			r.Obj = Lit{Val: p.lval}
		}
	case p.got(SHR):
		r.Op = SHR
		p.want(WORD)
		r.Obj = Lit{Val: p.lval}
	case p.got(LSS):
		r.Op = LSS
		p.want(WORD)
		r.Obj = Lit{Val: p.lval}
	default:
		return false
	}
	p.add(r)
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"

	"github.com/graphql-go/graphql"
	gqlhandler "github.com/graphql-go/graphql-go-handler"
)

var queryType = graphql.NewObject(graphql.ObjectConfig{
	Name: "Query",
	Fields: graphql.Fields{
		"latestPost": &graphql.Field{
			Type: graphql.String,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				return "Hello World!", nil
			},
		},
	},
})

var Schema, _ = graphql.NewSchema(graphql.SchemaConfig{
	Query: queryType,
})

func main() {
	h := gqlhandler.New(&gqlhandler.Config{
		Schema: &Schema,
		Pretty: true,
	})

	http.Handle("/graphql", h)

	http.ListenAndServe(":8080", nil)
}
<commit_msg>Route implemented with vestigo<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"

	"net/url"

	"github.com/graphql-go/graphql"
	"github.com/husobee/vestigo"
)

type user struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

var data map[string]user

/*
   Create User object type with fields "id" and "name" by using GraphQLObjectTypeConfig:
       - Name: name of object type
       - Fields: a map of fields by using GraphQLFields
   Setup type of field use GraphQLFieldConfig
*/
var userType = graphql.NewObject(
	graphql.ObjectConfig{
		Name: "User",
		Fields: graphql.Fields{
			"id": &graphql.Field{
				Type: graphql.String,
			},
			"name": &graphql.Field{
				Type: graphql.String,
			},
		},
	},
)

/*
   Create Query object type with fields "user" has type [userType] by using GraphQLObjectTypeConfig:
       - Name: name of object type
       - Fields: a map of fields by using GraphQLFields
   Setup type of field use GraphQLFieldConfig to define:
       - Type: type of field
       - Args: arguments to query with current field
       - Resolve: function to query data using params from [Args] and return value with current type
*/
var queryType = graphql.NewObject(
	graphql.ObjectConfig{
		Name: "Query",
		Fields: graphql.Fields{
			"user": &graphql.Field{
				Type: userType,
				Args: graphql.FieldConfigArgument{
					"id": &graphql.ArgumentConfig{
						Type: graphql.String,
					},
				},
				Resolve: func(p graphql.ResolveParams) (interface{}, error) {
					idQuery, isOK := p.Args["id"].(string)
					if isOK {
						return data[idQuery], nil
					}
					return nil, nil
				},
			},
		},
	},
)

var schema, _ = graphql.NewSchema(
	graphql.SchemaConfig{
		Query: queryType,
	},
)

func main() {

	// Router
	router := vestigo.NewRouter()
	vestigo.AllowTrace = true //Trace is on

	// Handlers
	router.HandleFunc("/", Index)

	router.Get("/graphql", GraphQLGETHandler)
	router.Post("/graphql", GraphQLPOSTHandler)

	// Serve
	log.Fatal(http.ListenAndServe(":8080", router))
}

func executeQuery(query string, schema graphql.Schema) *graphql.Result {
	result := graphql.Do(graphql.Params{
		Schema:        schema,
		RequestString: query,
	})
	if len(result.Errors) > 0 {
		fmt.Printf("wrong result, unexpected errors: %v", result.Errors)
	}
	return result
}

func Index(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(200)
	w.Write([]byte("POST/GET on /graphql"))
}

func GraphQLGETHandler(w http.ResponseWriter, r *http.Request) {
	query, err := url.ParseQuery(r.URL.RawQuery)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Errorf("Error parsing query: ", err.Error())
		return
	}

	w.WriteHeader(200)
	for key, value := range query {
		fmt.Fprintf(w, "Query: %s => %s ", key, value)
	}
}
func GraphQLPOSTHandler(w http.ResponseWriter, r *http.Request) {

}
<|endoftext|>


<commit_before>/*
   Copyright 2016 Simon J Mudd <sjmudd@pobox.com>

   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 discovery manages a queue as described: an ordered
// queue with no duplicates.
// The configured processor function is called on an entry that's
// received and this happens in parallel until we reach maxCapacity
// go-routines are running in parallel.  Any new requests get queued
// and processed as a go routine becomes free.  The queue is processed
// in FIFO order.  If a request is received for an element that is
// already in the queue or being processed then it will be silently
// ignored.
package discovery

import (
	"errors"
	"sync"

	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/go/inst"
)

type Queue struct {
	concurrency    uint                      // current concurrency
	done           chan inst.InstanceKey     // for synchronising completed discoveries
	inputChan      <-chan inst.InstanceKey   // input channel we are reading from
	knownKeys      map[inst.InstanceKey]bool // pending instances so we don't queue anything up more than one
	lock           sync.Mutex                // lock while making changes
	maxConcurrency uint                      // maximum concurrency of the queue
	processor      func(i inst.InstanceKey)  // process to run on each received key
	queue          []inst.InstanceKey        // instances in fifo order so we process in the order received.
}

var emptyKey = inst.InstanceKey{}

// provide a channel to read from and a function to run on the instance to be processed
func NewQueue(maxConcurrency uint, inputChan chan inst.InstanceKey, processor func(i inst.InstanceKey)) *Queue {
	log.Infof("Queue.NewQueue()")
	q := new(Queue)

	q.concurrency = 0                    // explicitly
	q.done = make(chan inst.InstanceKey) // Do I need this to be larger?
	q.inputChan = inputChan
	q.knownKeys = make(map[inst.InstanceKey]bool)
	q.maxConcurrency = maxConcurrency
	q.processor = processor
	q.queue = make([]inst.InstanceKey, 0)

	return q
}

// add the key to the slice if it does not exist in known keys
// - goroutine safe as only called inside the mutex
func (q *Queue) push(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.push(%v) is empty", key)
	}
	// log.Debugf("Queue.push(%+v)", key)

	if _, found := q.knownKeys[key]; !found {
		// log.Debugf("Queue.push() adding %+v to knownKeys", key)
		// add to the items that are being processed
		q.knownKeys[key] = true
		q.queue = append(q.queue, key)
	} else {
		// If key already there we just ignore it as the request is in the queue.
		// the known key also records stuff in the queue, so pending + active jobs.
		// log.Debugf("Queue.push() ignoring knownKey %+v", key)
	}
}

// remove the entry and remove it from known keys
func (q *Queue) pop() (inst.InstanceKey, error) {
	if len(q.queue) == 0 {
		return inst.InstanceKey{}, errors.New("q.pop() on empty queue")
	}
	key := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.knownKeys, key)
	// log.Debugf("Queue.pop() returns %+v", key)
	return key, nil
}

// dispatch a job from the queue (assumes we are in a locked state)
func (q *Queue) dispatch() {
	key, err := q.pop() // should never give an error but let's check anyway
	if err != nil {
		log.Fatal("Queue.dispatch() q.pop() returns: %+v", err)
		return
	}
	if key == emptyKey {
		log.Fatal("Queue.dispatch() key is empty")
	}

	q.concurrency++
	q.knownKeys[key] = true

	// log.Debugf("Queue.dispatch() key: %q, concurrency: %d", key, q.concurrency)

	// dispatch a discoverInstance() but tell us when we're done (to limit concurrency)
	go func() { // discover asynchronously
		q.processor(key)
		q.done <- key
	}()
}

// acknowledge a job has finished
// - we deal with the locking inside
func (q *Queue) acknowledgeJob(key inst.InstanceKey) {
	q.lock.Lock()
	delete(q.knownKeys, key)
	q.concurrency--
	// log.Debugf("Queue.acknowledgeJob(%+v) q.concurrency: %d", key, q.concurrency)
	q.lock.Unlock()
}

// drain queue by dispatching any jobs we have still
func (q *Queue) maybeDispatch() {
	q.lock.Lock()
	// log.Debugf("Queue.maybeDispatch() q.concurrency: %d, q.maxConcurrency: %d, len(q.queue): %d", q.concurrency, q.maxConcurrency, len(q.queue))
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// add an entry to the queue and dispatch something if concurrency is low enough
// - we deal with locking inside
func (q *Queue) queueAndMaybeDispatch(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.queueAndMaybeDispatch(%v) is empty", key)
	}
	q.lock.Lock()
	// log.Debugf("Queue.queueAndMaybeDispatch(%+v) concurency: %d", key, q.concurrency)
	q.push(key)
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// cleanup is called when the input channel closes.
// we can not sit in the loop so we have to wait for running go-routines to finish
// but also to dispatch anything left in the queue until finally everything is done.
func (q *Queue) cleanup() {
	log.Infof("Queue.cleanup()")
	for q.concurrency > 0 && len(q.queue) > 0 {
		q.maybeDispatch()
		if key, closed := <-q.done; closed {
			return
		} else {
			q.acknowledgeJob(key)
		}
	}
}

// Ends when all elements in the queue have been handled.
// we read from inputChan and call processor up to maxConcurrency times in parallel
func (q *Queue) HandleRequests() {
	if q == nil {
		log.Infof("Queue.HandleRequests() q == nil ??. Should not happen")

		// no queue, nothing to do
		return
	}
	log.Infof("Queue.NewQueue() processing requests")
	for {
		select {
		case key, ok := <-q.inputChan:
			if ok {
				if key != emptyKey {
					q.queueAndMaybeDispatch(key)
				} else {
					log.Warningf("Queue.HandleRequests() q.inputChan received empty key %+v, ignoring (fix the upstream code to prevent this)", key)
				}
			} else {
				q.cleanup()
				log.Infof("Queue.HandleRequests() q.inputChan is closed. returning")
				return
			}
		case key, ok := <-q.done:
			if ok {
				q.acknowledgeJob(key)
			} else {
				log.Infof("Queue.HandleRequests() q.done is closed. returning (shouldn't get here)")
				return // we shouldn't get here as the return above should get triggered first
			}
		}
	}
}
<commit_msg>Improve description of queuing and how variables are used<commit_after>/*
   Copyright 2016 Simon J Mudd <sjmudd@pobox.com>

   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 discovery manages a queue of discovery requests: an ordered
queue with no duplicates.

The configured processor function is called on an entry that's
received and this happens in parallel until maxCapacity go-routines
are running in parallel.

Any new requests get queued and processed as a go routine becomes
free.  The queue is processed in FIFO order.  If a request is
received for an instance that is already in the queue or already
being processed then it will be silently ignored.

*/

package discovery

import (
	"errors"
	"sync"

	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/go/inst"
)

// Queue is a container for processing the orchestrator discovery requests.
type Queue struct {
	concurrency    uint                      // The number of active go routines processing discovery requests.
	done           chan inst.InstanceKey     // Channel used by the active go routines to say they have finished processing.
	inputChan      <-chan inst.InstanceKey   // Input channel we are reading the discovery requests from.
	knownKeys      map[inst.InstanceKey]bool // This has 2 uses: to indicate if there is a request in the
						 // queue (so not being processed) or to indicate that the request
                                                 // is actively being dealth with.  That state is not explicitly
                                                 // stored as it is not really needed.
	lock           sync.Mutex                // lock while making changes
	maxConcurrency uint                      // The maximum number of go routines allowed to handle the queue at once.
                                                 // This is a configuration parameter provided when creating the Queue.
	processor      func(i inst.InstanceKey)  // process to run on each received key
	queue          []inst.InstanceKey        // This holds the discover requests (the instance name) which need to be
                                                 // processed. but which are not currently being processed. All requests
                                                 // are initially added to the end of this queue, and then the top element
                                                 // will be popped off if the number of active go routines (defined by
                                                 // concurrency) is less than the maximum specified value at which point
                                                 // it will be processed by a new go routine.
}

var emptyKey = inst.InstanceKey{}

// NewQueue creates a new Queue entry.
func NewQueue(maxConcurrency uint, inputChan chan inst.InstanceKey, processor func(i inst.InstanceKey)) *Queue {
	log.Infof("Queue.NewQueue()")
	q := new(Queue)

	q.concurrency = 0                    // explicitly
	q.done = make(chan inst.InstanceKey) // Do I need this to be larger?
	q.inputChan = inputChan
	q.knownKeys = make(map[inst.InstanceKey]bool)
	q.maxConcurrency = maxConcurrency
	q.processor = processor
	q.queue = make([]inst.InstanceKey, 0)

	return q
}

// add the key to the slice if it does not exist in known keys
// - goroutine safe as only called inside the mutex
func (q *Queue) push(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.push(%v) is empty", key)
	}
	// log.Debugf("Queue.push(%+v)", key)

	if _, found := q.knownKeys[key]; !found {
		// log.Debugf("Queue.push() adding %+v to knownKeys", key)
		// add to the items that are being processed
		q.knownKeys[key] = true
		q.queue = append(q.queue, key)
	} else {
		// If key already there we just ignore it as the request is in the queue.
		// the known key also records stuff in the queue, so pending + active jobs.
		// log.Debugf("Queue.push() ignoring knownKey %+v", key)
	}
}

// remove the entry and remove it from known keys
func (q *Queue) pop() (inst.InstanceKey, error) {
	if len(q.queue) == 0 {
		return inst.InstanceKey{}, errors.New("q.pop() on empty queue")
	}
	key := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.knownKeys, key)
	// log.Debugf("Queue.pop() returns %+v", key)
	return key, nil
}

// dispatch a job from the queue (assumes we are in a locked state)
func (q *Queue) dispatch() {
	key, err := q.pop() // should never give an error but let's check anyway
	if err != nil {
		log.Fatal("Queue.dispatch() q.pop() returns: %+v", err)
		return
	}
	if key == emptyKey {
		log.Fatal("Queue.dispatch() key is empty")
	}

	q.concurrency++
	q.knownKeys[key] = true

	// log.Debugf("Queue.dispatch() key: %q, concurrency: %d", key, q.concurrency)

	// dispatch a discoverInstance() but tell us when we're done (to limit concurrency)
	go func() { // discover asynchronously
		q.processor(key)
		q.done <- key
	}()
}

// acknowledge a job has finished
// - we deal with the locking inside
func (q *Queue) acknowledgeJob(key inst.InstanceKey) {
	q.lock.Lock()
	delete(q.knownKeys, key)
	q.concurrency--
	// log.Debugf("Queue.acknowledgeJob(%+v) q.concurrency: %d", key, q.concurrency)
	q.lock.Unlock()
}

// drain queue by dispatching any jobs we have still
func (q *Queue) maybeDispatch() {
	q.lock.Lock()
	// log.Debugf("Queue.maybeDispatch() q.concurrency: %d, q.maxConcurrency: %d, len(q.queue): %d", q.concurrency, q.maxConcurrency, len(q.queue))
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// add an entry to the queue and dispatch something if concurrency is low enough
// - we deal with locking inside
func (q *Queue) queueAndMaybeDispatch(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.queueAndMaybeDispatch(%v) is empty", key)
	}
	q.lock.Lock()
	// log.Debugf("Queue.queueAndMaybeDispatch(%+v) concurency: %d", key, q.concurrency)
	q.push(key)
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// cleanup is called when the input channel closes.
// we can not sit in the loop so we have to wait for running go-routines to finish
// but also to dispatch anything left in the queue until finally everything is done.
func (q *Queue) cleanup() {
	log.Infof("Queue.cleanup()")
	for q.concurrency > 0 && len(q.queue) > 0 {
		q.maybeDispatch()
		if key, closed := <-q.done; closed {
			return
		} else {
			q.acknowledgeJob(key)
		}
	}
}

// Ends when all elements in the queue have been handled.
// we read from inputChan and call processor up to maxConcurrency times in parallel
func (q *Queue) HandleRequests() {
	if q == nil {
		log.Infof("Queue.HandleRequests() q == nil ??. Should not happen")

		// no queue, nothing to do
		return
	}
	log.Infof("Queue.NewQueue() processing requests")
	for {
		select {
		case key, ok := <-q.inputChan:
			if ok {
				if key != emptyKey {
					q.queueAndMaybeDispatch(key)
				} else {
					log.Warningf("Queue.HandleRequests() q.inputChan received empty key %+v, ignoring (fix the upstream code to prevent this)", key)
				}
			} else {
				q.cleanup()
				log.Infof("Queue.HandleRequests() q.inputChan is closed. returning")
				return
			}
		case key, ok := <-q.done:
			if ok {
				q.acknowledgeJob(key)
			} else {
				log.Infof("Queue.HandleRequests() q.done is closed. returning (shouldn't get here)")
				return // we shouldn't get here as the return above should get triggered first
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// MerkleAudit runs a merkle tree audit in the background once in a while.
// It verifies the skips of a randomly chosen merkle tree root, making
// sure that the server is not tampering with the merkle trees.

package engine

import (
	"crypto/rand"
	"errors"
	"fmt"
	"math/big"
	"sync"
	"time"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

var (
	errAuditOffline    = errors.New("Merkle audit failed to run due to the lack of connectivity.")
	errAuditNoLastRoot = errors.New("Merkle audit failed to run due to not being able to get the last root.")
)

var MerkleAuditSettings = BackgroundTaskSettings{
	Start:        5 * time.Minute,
	StartStagger: 1 * time.Hour,
	Interval:     6 * time.Hour,
	Limit:        1 * time.Minute,
}

// MerkleAudit is an engine.
type MerkleAudit struct {
	libkb.Contextified
	sync.Mutex

	args *MerkleAuditArgs
	task *BackgroundTask
}

type MerkleAuditArgs struct {
	// Channels used for testing. Normally nil.
	testingMetaCh     chan<- string
	testingRoundResCh chan<- error
}

type merkleAuditState struct {
	RetrySeqno *keybase1.Seqno `json:"retrySeqno"`
}

// NewMerkleAudit creates a new MerkleAudit engine.
func NewMerkleAudit(g *libkb.GlobalContext, args *MerkleAuditArgs) *MerkleAudit {
	task := NewBackgroundTask(g, &BackgroundTaskArgs{
		Name:     "MerkleAudit",
		F:        MerkleAuditRound,
		Settings: MerkleAuditSettings,

		testingMetaCh:     args.testingMetaCh,
		testingRoundResCh: args.testingRoundResCh,
	})
	return &MerkleAudit{
		Contextified: libkb.NewContextified(g),
		args:         args,
		task:         task,
	}
}

// Name is the unique engine name.
func (e *MerkleAudit) Name() string {
	return "MerkleAudit"
}

// GetPrereqs returns the engine prereqs.
func (e *MerkleAudit) Prereqs() Prereqs {
	return Prereqs{}
}

// RequiredUIs returns the required UIs.
func (e *MerkleAudit) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *MerkleAudit) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{}
}

// Run starts the engine.
// Returns immediately, kicks off a background goroutine.
func (e *MerkleAudit) Run(m libkb.MetaContext) (err error) {
	if m.G().GetEnv().GetDisableMerkleAuditor() {
		m.G().Log.CDebugf(m.Ctx(), "merkle audit disabled, aborting run")
	}
	return RunEngine2(m, e.task)
}

func (e *MerkleAudit) Shutdown() {
	e.task.Shutdown()
}

// randSeqno picks a random number between [low, high).
func randSeqno(lo keybase1.Seqno, hi keybase1.Seqno) (keybase1.Seqno, error) {
	rangeBig := big.NewInt(int64(hi - lo))
	n, err := rand.Int(rand.Reader, rangeBig)
	if err != nil {
		return keybase1.Seqno(0), err
	}
	return keybase1.Seqno(n.Int64()) + lo, nil
}

var merkleAuditKey = libkb.DbKey{
	Typ: libkb.DBMerkleAudit,
	Key: "root",
}

func lookupMerkleAuditRetryFromState(m libkb.MetaContext) (*keybase1.Seqno, error) {
	var state merkleAuditState
	found, err := m.G().LocalDb.GetInto(&state, merkleAuditKey)
	if err != nil {
		return nil, err
	}
	if !found {
		// Nothing found, no error
		return nil, nil
	}

	// Can still be nil
	return state.RetrySeqno, nil
}

func saveMerkleAuditState(m libkb.MetaContext, state merkleAuditState) error {
	return m.G().LocalDb.PutObj(merkleAuditKey, nil, state)
}

func performMerkleAudit(m libkb.MetaContext, startSeqno keybase1.Seqno) error {
	if m.G().ConnectivityMonitor.IsConnected(m.Ctx()) == libkb.ConnectivityMonitorNo {
		m.Debug("MerkleAudit giving up offline")
		return errAuditOffline
	}

	// Acquire the most recent merkle tree root
	lastRoot := m.G().MerkleClient.LastRoot()
	if lastRoot == nil {
		m.Debug("MerkleAudit unable to retrieve the last root")
		return errAuditNoLastRoot
	}

	// We can copy the pointer's value as it can only return nil if root == nil.
	lastSeqno := *lastRoot.Seqno()

	// Acquire the first root and calculate its hash
	startRoot, err := m.G().MerkleClient.LookupRootAtSeqno(m, startSeqno)
	if err != nil {
		return err
	}
	startHash := startRoot.ShortHash()

	// Traverse the merkle tree seqnos
	currentSeqno := startSeqno + 1
	step := 1
	for {
		// Proceed until the last known root
		if currentSeqno > lastSeqno {
			break
		}

		currentRoot, err := m.G().MerkleClient.LookupRootAtSeqno(m, currentSeqno)
		if err != nil {
			return err
		}
		currentHash := currentRoot.SkipToSeqno(startSeqno)
		if currentHash == nil {
			return libkb.NewClientMerkleSkipMissingError(
				fmt.Sprintf("Root %d missing skip hash to %d", currentSeqno, startSeqno),
			)
		}

		if !startHash.Eq(currentHash) {
			// Warn the user about the possibility of the server tampering with the roots.
			return libkb.NewClientMerkleSkipHashMismatchError(
				fmt.Sprintf("Invalid skip hash from %d to %d", currentSeqno, startSeqno),
			)
		}

		// We're doing this exponentially to make use of the skips.
		currentSeqno += keybase1.Seqno(step)
		step *= 2
	}

	return nil
}

func MerkleAuditRound(m libkb.MetaContext) (err error) {
	m = m.WithLogTag("MAUDT")
	defer m.TraceTimed("MerkleAuditRound", func() error { return err })()

	// Look up any previously requested retries
	startSeqno, err := lookupMerkleAuditRetryFromState(m)
	if err != nil {
		m.Debug("MerkleAudit unable to acquire saved state from localdb")
		return nil
	}

	// If no retry was requested
	if startSeqno == nil {
		// nil seqno, generate a new one:
		// 1. Acquire the most recent merkle tree root
		lastRoot := m.G().MerkleClient.LastRoot()
		if lastRoot == nil {
			m.Debug("MerkleAudit unable to retrieve the last root")
			return nil
		}
		lastSeqno := *lastRoot.Seqno()

		// 2. Figure out the first merkle root seqno with skips, fall back to 1
		firstSeqno := m.G().MerkleClient.FirstSeqnoWithSkips(m)
		if firstSeqno == nil {
			val := keybase1.Seqno(1)
			firstSeqno = &val
		}

		// 3. Generate a random seqno for the starting root in the audit.
		randomSeqno, err := randSeqno(*firstSeqno, lastSeqno)
		if err != nil {
			return err
		}

		startSeqno = &randomSeqno
	}

	// If this time it fails, save it
	err = performMerkleAudit(m, *startSeqno)
	if err == nil {
		// Early return for fewer ifs
		return saveMerkleAuditState(m, merkleAuditState{
			RetrySeqno: nil,
		})
	}

	// All MerkleClientErrors would suggest that the server is tampering with the roots
	if _, ok := err.(libkb.MerkleClientError); ok {
		m.Error("MerkleAudit fatally failed: %s", err)
		// Send the notification to the client
		m.G().NotifyRouter.HandleRootAuditError(fmt.Sprintf(
			"Merkle tree audit from %d failed: %s",
			startSeqno, err.Error(),
		))
	} else {
		m.Debug("MerkleAudit could not complete: %s", err)
	}

	// Use another error variable to prevent shadowing
	if serr := saveMerkleAuditState(m, merkleAuditState{
		RetrySeqno: startSeqno,
	}); serr != nil {
		return serr
	}

	return err
}
<commit_msg>fix for TestMerkleAuditRetry flake (#16669)<commit_after>// Copyright 2019 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// MerkleAudit runs a merkle tree audit in the background once in a while.
// It verifies the skips of a randomly chosen merkle tree root, making
// sure that the server is not tampering with the merkle trees.

package engine

import (
	"crypto/rand"
	"errors"
	"fmt"
	"math/big"
	"sync"
	"time"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

var (
	errAuditOffline    = errors.New("Merkle audit failed to run due to the lack of connectivity.")
	errAuditNoLastRoot = errors.New("Merkle audit failed to run due to not being able to get the last root.")
)

var MerkleAuditSettings = BackgroundTaskSettings{
	Start:        5 * time.Minute,
	StartStagger: 1 * time.Hour,
	Interval:     6 * time.Hour,
	Limit:        1 * time.Minute,
}

// MerkleAudit is an engine.
type MerkleAudit struct {
	libkb.Contextified
	sync.Mutex

	args *MerkleAuditArgs
	task *BackgroundTask
}

type MerkleAuditArgs struct {
	// Channels used for testing. Normally nil.
	testingMetaCh     chan<- string
	testingRoundResCh chan<- error
}

type merkleAuditState struct {
	RetrySeqno *keybase1.Seqno `json:"retrySeqno"`
	LastSeqno  *keybase1.Seqno `json:"lastSeqno"`
}

// NewMerkleAudit creates a new MerkleAudit engine.
func NewMerkleAudit(g *libkb.GlobalContext, args *MerkleAuditArgs) *MerkleAudit {
	task := NewBackgroundTask(g, &BackgroundTaskArgs{
		Name:     "MerkleAudit",
		F:        MerkleAuditRound,
		Settings: MerkleAuditSettings,

		testingMetaCh:     args.testingMetaCh,
		testingRoundResCh: args.testingRoundResCh,
	})
	return &MerkleAudit{
		Contextified: libkb.NewContextified(g),
		args:         args,
		task:         task,
	}
}

// Name is the unique engine name.
func (e *MerkleAudit) Name() string {
	return "MerkleAudit"
}

// GetPrereqs returns the engine prereqs.
func (e *MerkleAudit) Prereqs() Prereqs {
	return Prereqs{}
}

// RequiredUIs returns the required UIs.
func (e *MerkleAudit) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *MerkleAudit) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{}
}

// Run starts the engine.
// Returns immediately, kicks off a background goroutine.
func (e *MerkleAudit) Run(m libkb.MetaContext) (err error) {
	if m.G().GetEnv().GetDisableMerkleAuditor() {
		m.G().Log.CDebugf(m.Ctx(), "merkle audit disabled, aborting run")
	}
	return RunEngine2(m, e.task)
}

func (e *MerkleAudit) Shutdown() {
	e.task.Shutdown()
}

// randSeqno picks a random number between [low, high) that's different from prev.
func randSeqno(lo keybase1.Seqno, hi keybase1.Seqno, prev *keybase1.Seqno) (keybase1.Seqno, error) {
	// Prevent an infinite loop if [0,1) and prev = 0
	if hi-lo == 1 && prev != nil && *prev == lo {
		return keybase1.Seqno(0), fmt.Errorf("unable to generate a non-duplicate seqno other than %d", *prev)
	}
	for {
		rangeBig := big.NewInt(int64(hi - lo))
		n, err := rand.Int(rand.Reader, rangeBig)
		if err != nil {
			return keybase1.Seqno(0), err
		}
		newSeqno := keybase1.Seqno(n.Int64()) + lo
		if prev == nil || *prev != newSeqno {
			return newSeqno, nil
		}
	}
}

var merkleAuditKey = libkb.DbKey{
	Typ: libkb.DBMerkleAudit,
	Key: "root",
}

func lookupMerkleAuditRetryFromState(m libkb.MetaContext) (*keybase1.Seqno, *keybase1.Seqno, error) {
	var state merkleAuditState
	found, err := m.G().LocalDb.GetInto(&state, merkleAuditKey)
	if err != nil {
		return nil, nil, err
	}
	if !found {
		// Nothing found, no error
		return nil, nil, nil
	}

	// Can still be nil
	return state.RetrySeqno, state.LastSeqno, nil
}

func saveMerkleAuditState(m libkb.MetaContext, state merkleAuditState) error {
	return m.G().LocalDb.PutObj(merkleAuditKey, nil, state)
}

func performMerkleAudit(m libkb.MetaContext, startSeqno keybase1.Seqno) error {
	if m.G().ConnectivityMonitor.IsConnected(m.Ctx()) == libkb.ConnectivityMonitorNo {
		m.Debug("MerkleAudit giving up offline")
		return errAuditOffline
	}

	// Acquire the most recent merkle tree root
	lastRoot := m.G().MerkleClient.LastRoot()
	if lastRoot == nil {
		m.Debug("MerkleAudit unable to retrieve the last root")
		return errAuditNoLastRoot
	}

	// We can copy the pointer's value as it can only return nil if root == nil.
	lastSeqno := *lastRoot.Seqno()

	// Acquire the first root and calculate its hash
	startRoot, err := m.G().MerkleClient.LookupRootAtSeqno(m, startSeqno)
	if err != nil {
		return err
	}
	startHash := startRoot.ShortHash()

	// Traverse the merkle tree seqnos
	currentSeqno := startSeqno + 1
	step := 1
	for {
		// Proceed until the last known root
		if currentSeqno > lastSeqno {
			break
		}

		currentRoot, err := m.G().MerkleClient.LookupRootAtSeqno(m, currentSeqno)
		if err != nil {
			return err
		}
		currentHash := currentRoot.SkipToSeqno(startSeqno)
		if currentHash == nil {
			return libkb.NewClientMerkleSkipMissingError(
				fmt.Sprintf("Root %d missing skip hash to %d", currentSeqno, startSeqno),
			)
		}

		if !startHash.Eq(currentHash) {
			// Warn the user about the possibility of the server tampering with the roots.
			return libkb.NewClientMerkleSkipHashMismatchError(
				fmt.Sprintf("Invalid skip hash from %d to %d", currentSeqno, startSeqno),
			)
		}

		// We're doing this exponentially to make use of the skips.
		currentSeqno += keybase1.Seqno(step)
		step *= 2
	}

	return nil
}

func MerkleAuditRound(m libkb.MetaContext) (err error) {
	m = m.WithLogTag("MAUDT")
	defer m.TraceTimed("MerkleAuditRound", func() error { return err })()

	// Look up any previously requested retries
	startSeqno, prevSeqno, err := lookupMerkleAuditRetryFromState(m)
	if err != nil {
		m.Debug("MerkleAudit unable to acquire saved state from localdb")
		return nil
	}

	// If no retry was requested
	if startSeqno == nil {
		// nil seqno, generate a new one:
		// 1. Acquire the most recent merkle tree root
		lastRoot := m.G().MerkleClient.LastRoot()
		if lastRoot == nil {
			m.Debug("MerkleAudit unable to retrieve the last root")
			return nil
		}
		lastSeqno := *lastRoot.Seqno()

		// 2. Figure out the first merkle root seqno with skips, fall back to 1
		firstSeqno := m.G().MerkleClient.FirstSeqnoWithSkips(m)
		if firstSeqno == nil {
			val := keybase1.Seqno(1)
			firstSeqno = &val
		}

		// 3. Generate a random seqno for the starting root in the audit.
		randomSeqno, err := randSeqno(*firstSeqno, lastSeqno, prevSeqno)
		if err != nil {
			return err
		}
		startSeqno = &randomSeqno
	}

	// If this time it fails, save it
	err = performMerkleAudit(m, *startSeqno)
	if err == nil {
		// Early return for fewer ifs
		return saveMerkleAuditState(m, merkleAuditState{
			RetrySeqno: nil,
			LastSeqno:  startSeqno,
		})
	}

	// All MerkleClientErrors would suggest that the server is tampering with the roots
	if _, ok := err.(libkb.MerkleClientError); ok {
		m.Error("MerkleAudit fatally failed: %s", err)
		// Send the notification to the client
		m.G().NotifyRouter.HandleRootAuditError(fmt.Sprintf(
			"Merkle tree audit from %d failed: %s",
			startSeqno, err.Error(),
		))
	} else {
		m.Debug("MerkleAudit could not complete: %s", err)
	}

	// Use another error variable to prevent shadowing
	if serr := saveMerkleAuditState(m, merkleAuditState{
		RetrySeqno: startSeqno,
		LastSeqno:  prevSeqno,
	}); serr != nil {
		return serr
	}

	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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 pools provides functionality to manage and reuse resources
// like connections.
package pools

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/timer"
	"vitess.io/vitess/go/trace"
)

var (
	// ErrClosed is returned if ResourcePool is used when it's closed.
	ErrClosed = errors.New("resource pool is closed")

	// ErrTimeout is returned if a resource get times out.
	ErrTimeout = errors.New("resource pool timed out")

	// ErrCtxTimeout is returned if a ctx is already expired by the time the resource pool is used
	ErrCtxTimeout = errors.New("resource pool context already expired")

	prefillTimeout = 30 * time.Second
)

// Factory is a function that can be used to create a resource.
type Factory func(context.Context) (Resource, error)

// Resource defines the interface that every resource must provide.
// Thread synchronization between Close() and IsClosed()
// is the responsibility of the caller.
type Resource interface {
	Close()
}

// ResourcePool allows you to use a pool of resources.
type ResourcePool struct {
	// stats. Atomic fields must remain at the top in order to prevent panics on certain architectures.
	available  sync2.AtomicInt64
	active     sync2.AtomicInt64
	inUse      sync2.AtomicInt64
	waitCount  sync2.AtomicInt64
	waitTime   sync2.AtomicDuration
	idleClosed sync2.AtomicInt64
	exhausted  sync2.AtomicInt64

	capacity    sync2.AtomicInt64
	idleTimeout sync2.AtomicDuration

	resources chan resourceWrapper
	factory   Factory
	idleTimer *timer.Timer
	logWait   func(time.Time)
}

type resourceWrapper struct {
	resource Resource
	timeUsed time.Time
}

// NewResourcePool creates a new ResourcePool pool.
// capacity is the number of possible resources in the pool:
// there can be up to 'capacity' of these at a given time.
// maxCap specifies the extent to which the pool can be resized
// in the future through the SetCapacity function.
// You cannot resize the pool beyond maxCap.
// If a resource is unused beyond idleTimeout, it's replaced
// with a new one.
// An idleTimeout of 0 means that there is no timeout.
// A non-zero value of prefillParallelism causes the pool to be pre-filled.
// The value specifies how many resources can be opened in parallel.
func NewResourcePool(factory Factory, capacity, maxCap int, idleTimeout time.Duration, prefillParallelism int, logWait func(time.Time)) *ResourcePool {
	if capacity <= 0 || maxCap <= 0 || capacity > maxCap {
		panic(errors.New("invalid/out of range capacity"))
	}
	rp := &ResourcePool{
		resources:   make(chan resourceWrapper, maxCap),
		factory:     factory,
		available:   sync2.NewAtomicInt64(int64(capacity)),
		capacity:    sync2.NewAtomicInt64(int64(capacity)),
		idleTimeout: sync2.NewAtomicDuration(idleTimeout),
		logWait:     logWait,
	}
	for i := 0; i < capacity; i++ {
		rp.resources <- resourceWrapper{}
	}

	ctx, cancel := context.WithTimeout(context.TODO(), prefillTimeout)
	defer cancel()
	if prefillParallelism != 0 {
		sem := sync2.NewSemaphore(prefillParallelism, 0 /* timeout */)
		var wg sync.WaitGroup
		for i := 0; i < capacity; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				_ = sem.Acquire()
				defer sem.Release()

				// If context has expired, give up.
				select {
				case <-ctx.Done():
					return
				default:
				}

				r, err := rp.Get(ctx)
				if err != nil {
					return
				}
				rp.Put(r)
			}()
		}
		wg.Wait()
	}

	if idleTimeout != 0 {
		rp.idleTimer = timer.NewTimer(idleTimeout / 10)
		rp.idleTimer.Start(rp.closeIdleResources)
	}
	return rp
}

// Close empties the pool calling Close on all its resources.
// You can call Close while there are outstanding resources.
// It waits for all resources to be returned (Put).
// After a Close, Get is not allowed.
func (rp *ResourcePool) Close() {
	if rp.idleTimer != nil {
		rp.idleTimer.Stop()
	}
	_ = rp.SetCapacity(0)
}

// IsClosed returns true if the resource pool is closed.
func (rp *ResourcePool) IsClosed() (closed bool) {
	return rp.capacity.Get() == 0
}

// closeIdleResources scans the pool for idle resources
func (rp *ResourcePool) closeIdleResources() {
	available := int(rp.Available())
	idleTimeout := rp.IdleTimeout()

	for i := 0; i < available; i++ {
		var wrapper resourceWrapper
		select {
		case wrapper = <-rp.resources:
		default:
			// stop early if we don't get anything new from the pool
			return
		}

		func() {
			defer func() { rp.resources <- wrapper }()

			if wrapper.resource != nil && idleTimeout > 0 && time.Until(wrapper.timeUsed.Add(idleTimeout)) < 0 {
				wrapper.resource.Close()
				rp.idleClosed.Add(1)
				rp.reopenResource(&wrapper)
			}
		}()

	}
}

// Get will return the next available resource. If capacity
// has not been reached, it will create a new one using the factory. Otherwise,
// it will wait till the next resource becomes available or a timeout.
// A timeout of 0 is an indefinite wait.
func (rp *ResourcePool) Get(ctx context.Context) (resource Resource, err error) {
	span, ctx := trace.NewSpan(ctx, "ResourcePool.Get")
	span.Annotate("capacity", rp.capacity.Get())
	span.Annotate("in_use", rp.inUse.Get())
	span.Annotate("available", rp.available.Get())
	span.Annotate("active", rp.active.Get())
	defer span.Finish()
	return rp.get(ctx)
}

func (rp *ResourcePool) get(ctx context.Context) (resource Resource, err error) {
	// If ctx has already expired, avoid racing with rp's resource channel.
	select {
	case <-ctx.Done():
		return nil, ErrCtxTimeout
	default:
	}

	// Fetch
	var wrapper resourceWrapper
	var ok bool
	select {
	case wrapper, ok = <-rp.resources:
	default:
		startTime := time.Now()
		select {
		case wrapper, ok = <-rp.resources:
		case <-ctx.Done():
			return nil, ErrTimeout
		}
		rp.recordWait(startTime)
	}
	if !ok {
		return nil, ErrClosed
	}

	// Unwrap
	if wrapper.resource == nil {
		span, _ := trace.NewSpan(ctx, "ResourcePool.factory")
		wrapper.resource, err = rp.factory(ctx)
		span.Finish()
		if err != nil {
			rp.resources <- resourceWrapper{}
			return nil, err
		}
		rp.active.Add(1)
	}
	if rp.available.Add(-1) <= 0 {
		rp.exhausted.Add(1)
	}
	rp.inUse.Add(1)
	return wrapper.resource, err
}

// Put will return a resource to the pool. For every successful Get,
// a corresponding Put is required. If you no longer need a resource,
// you will need to call Put(nil) instead of returning the closed resource.
// This will cause a new resource to be created in its place.
func (rp *ResourcePool) Put(resource Resource) {
	var wrapper resourceWrapper
	if resource != nil {
		wrapper = resourceWrapper{
			resource: resource,
			timeUsed: time.Now(),
		}
	} else {
		rp.reopenResource(&wrapper)
	}
	select {
	case rp.resources <- wrapper:
	default:
		panic(errors.New("attempt to Put into a full ResourcePool"))
	}
	rp.inUse.Add(-1)
	rp.available.Add(1)
}

func (rp *ResourcePool) reopenResource(wrapper *resourceWrapper) {
	if r, err := rp.factory(context.TODO()); err == nil {
		wrapper.resource = r
		wrapper.timeUsed = time.Now()
	} else {
		wrapper.resource = nil
		rp.active.Add(-1)
	}
}

// SetCapacity changes the capacity of the pool.
// You can use it to shrink or expand, but not beyond
// the max capacity. If the change requires the pool
// to be shrunk, SetCapacity waits till the necessary
// number of resources are returned to the pool.
// A SetCapacity of 0 is equivalent to closing the ResourcePool.
func (rp *ResourcePool) SetCapacity(capacity int) error {
	if capacity < 0 || capacity > cap(rp.resources) {
		return fmt.Errorf("capacity %d is out of range", capacity)
	}

	// Atomically swap new capacity with old, but only
	// if old capacity is non-zero.
	var oldcap int
	for {
		oldcap = int(rp.capacity.Get())
		if oldcap == 0 {
			return ErrClosed
		}
		if oldcap == capacity {
			return nil
		}
		if rp.capacity.CompareAndSwap(int64(oldcap), int64(capacity)) {
			break
		}
	}

	if capacity < oldcap {
		for i := 0; i < oldcap-capacity; i++ {
			wrapper := <-rp.resources
			if wrapper.resource != nil {
				wrapper.resource.Close()
				rp.active.Add(-1)
			}
			rp.available.Add(-1)
		}
	} else {
		for i := 0; i < capacity-oldcap; i++ {
			rp.resources <- resourceWrapper{}
			rp.available.Add(1)
		}
	}
	if capacity == 0 {
		close(rp.resources)
	}
	return nil
}

func (rp *ResourcePool) recordWait(start time.Time) {
	rp.waitCount.Add(1)
	rp.waitTime.Add(time.Since(start))
	if rp.logWait != nil {
		rp.logWait(start)
	}
}

// SetIdleTimeout sets the idle timeout. It can only be used if there was an
// idle timeout set when the pool was created.
func (rp *ResourcePool) SetIdleTimeout(idleTimeout time.Duration) {
	if rp.idleTimer == nil {
		panic("SetIdleTimeout called when timer not initialized")
	}

	rp.idleTimeout.Set(idleTimeout)
	rp.idleTimer.SetInterval(idleTimeout / 10)
}

// StatsJSON returns the stats in JSON format.
func (rp *ResourcePool) StatsJSON() string {
	return fmt.Sprintf(`{"Capacity": %v, "Available": %v, "Active": %v, "InUse": %v, "MaxCapacity": %v, "WaitCount": %v, "WaitTime": %v, "IdleTimeout": %v, "IdleClosed": %v, "Exhausted": %v}`,
		rp.Capacity(),
		rp.Available(),
		rp.Active(),
		rp.InUse(),
		rp.MaxCap(),
		rp.WaitCount(),
		rp.WaitTime().Nanoseconds(),
		rp.IdleTimeout().Nanoseconds(),
		rp.IdleClosed(),
		rp.Exhausted(),
	)
}

// Capacity returns the capacity.
func (rp *ResourcePool) Capacity() int64 {
	return rp.capacity.Get()
}

// Available returns the number of currently unused and available resources.
func (rp *ResourcePool) Available() int64 {
	return rp.available.Get()
}

// Active returns the number of active (i.e. non-nil) resources either in the
// pool or claimed for use
func (rp *ResourcePool) Active() int64 {
	return rp.active.Get()
}

// InUse returns the number of claimed resources from the pool
func (rp *ResourcePool) InUse() int64 {
	return rp.inUse.Get()
}

// MaxCap returns the max capacity.
func (rp *ResourcePool) MaxCap() int64 {
	return int64(cap(rp.resources))
}

// WaitCount returns the total number of waits.
func (rp *ResourcePool) WaitCount() int64 {
	return rp.waitCount.Get()
}

// WaitTime returns the total wait time.
func (rp *ResourcePool) WaitTime() time.Duration {
	return rp.waitTime.Get()
}

// IdleTimeout returns the idle timeout.
func (rp *ResourcePool) IdleTimeout() time.Duration {
	return rp.idleTimeout.Get()
}

// IdleClosed returns the count of resources closed due to idle timeout.
func (rp *ResourcePool) IdleClosed() int64 {
	return rp.idleClosed.Get()
}

// Exhausted returns the number of times Available dropped below 1
func (rp *ResourcePool) Exhausted() int64 {
	return rp.exhausted.Get()
}
<commit_msg>Adds better error when there are timeouts in resource pools<commit_after>/*
Copyright 2019 The Vitess Authors.

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 pools provides functionality to manage and reuse resources
// like connections.
package pools

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/timer"
	"vitess.io/vitess/go/trace"
	"vitess.io/vitess/go/vt/vterrors"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

var (
	// ErrClosed is returned if ResourcePool is used when it's closed.
	ErrClosed = errors.New("resource pool is closed")

	// ErrTimeout is returned if a resource get times out.
	ErrTimeout = vterrors.New(vtrpcpb.Code_DEADLINE_EXCEEDED, "resource pool timed out")

	// ErrCtxTimeout is returned if a ctx is already expired by the time the resource pool is used
	ErrCtxTimeout = vterrors.New(vtrpcpb.Code_DEADLINE_EXCEEDED, "resource pool context already expired")

	prefillTimeout = 30 * time.Second
)

// Factory is a function that can be used to create a resource.
type Factory func(context.Context) (Resource, error)

// Resource defines the interface that every resource must provide.
// Thread synchronization between Close() and IsClosed()
// is the responsibility of the caller.
type Resource interface {
	Close()
}

// ResourcePool allows you to use a pool of resources.
type ResourcePool struct {
	// stats. Atomic fields must remain at the top in order to prevent panics on certain architectures.
	available  sync2.AtomicInt64
	active     sync2.AtomicInt64
	inUse      sync2.AtomicInt64
	waitCount  sync2.AtomicInt64
	waitTime   sync2.AtomicDuration
	idleClosed sync2.AtomicInt64
	exhausted  sync2.AtomicInt64

	capacity    sync2.AtomicInt64
	idleTimeout sync2.AtomicDuration

	resources chan resourceWrapper
	factory   Factory
	idleTimer *timer.Timer
	logWait   func(time.Time)
}

type resourceWrapper struct {
	resource Resource
	timeUsed time.Time
}

// NewResourcePool creates a new ResourcePool pool.
// capacity is the number of possible resources in the pool:
// there can be up to 'capacity' of these at a given time.
// maxCap specifies the extent to which the pool can be resized
// in the future through the SetCapacity function.
// You cannot resize the pool beyond maxCap.
// If a resource is unused beyond idleTimeout, it's replaced
// with a new one.
// An idleTimeout of 0 means that there is no timeout.
// A non-zero value of prefillParallelism causes the pool to be pre-filled.
// The value specifies how many resources can be opened in parallel.
func NewResourcePool(factory Factory, capacity, maxCap int, idleTimeout time.Duration, prefillParallelism int, logWait func(time.Time)) *ResourcePool {
	if capacity <= 0 || maxCap <= 0 || capacity > maxCap {
		panic(errors.New("invalid/out of range capacity"))
	}
	rp := &ResourcePool{
		resources:   make(chan resourceWrapper, maxCap),
		factory:     factory,
		available:   sync2.NewAtomicInt64(int64(capacity)),
		capacity:    sync2.NewAtomicInt64(int64(capacity)),
		idleTimeout: sync2.NewAtomicDuration(idleTimeout),
		logWait:     logWait,
	}
	for i := 0; i < capacity; i++ {
		rp.resources <- resourceWrapper{}
	}

	ctx, cancel := context.WithTimeout(context.TODO(), prefillTimeout)
	defer cancel()
	if prefillParallelism != 0 {
		sem := sync2.NewSemaphore(prefillParallelism, 0 /* timeout */)
		var wg sync.WaitGroup
		for i := 0; i < capacity; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				_ = sem.Acquire()
				defer sem.Release()

				// If context has expired, give up.
				select {
				case <-ctx.Done():
					return
				default:
				}

				r, err := rp.Get(ctx)
				if err != nil {
					return
				}
				rp.Put(r)
			}()
		}
		wg.Wait()
	}

	if idleTimeout != 0 {
		rp.idleTimer = timer.NewTimer(idleTimeout / 10)
		rp.idleTimer.Start(rp.closeIdleResources)
	}
	return rp
}

// Close empties the pool calling Close on all its resources.
// You can call Close while there are outstanding resources.
// It waits for all resources to be returned (Put).
// After a Close, Get is not allowed.
func (rp *ResourcePool) Close() {
	if rp.idleTimer != nil {
		rp.idleTimer.Stop()
	}
	_ = rp.SetCapacity(0)
}

// IsClosed returns true if the resource pool is closed.
func (rp *ResourcePool) IsClosed() (closed bool) {
	return rp.capacity.Get() == 0
}

// closeIdleResources scans the pool for idle resources
func (rp *ResourcePool) closeIdleResources() {
	available := int(rp.Available())
	idleTimeout := rp.IdleTimeout()

	for i := 0; i < available; i++ {
		var wrapper resourceWrapper
		select {
		case wrapper = <-rp.resources:
		default:
			// stop early if we don't get anything new from the pool
			return
		}

		func() {
			defer func() { rp.resources <- wrapper }()

			if wrapper.resource != nil && idleTimeout > 0 && time.Until(wrapper.timeUsed.Add(idleTimeout)) < 0 {
				wrapper.resource.Close()
				rp.idleClosed.Add(1)
				rp.reopenResource(&wrapper)
			}
		}()

	}
}

// Get will return the next available resource. If capacity
// has not been reached, it will create a new one using the factory. Otherwise,
// it will wait till the next resource becomes available or a timeout.
// A timeout of 0 is an indefinite wait.
func (rp *ResourcePool) Get(ctx context.Context) (resource Resource, err error) {
	span, ctx := trace.NewSpan(ctx, "ResourcePool.Get")
	span.Annotate("capacity", rp.capacity.Get())
	span.Annotate("in_use", rp.inUse.Get())
	span.Annotate("available", rp.available.Get())
	span.Annotate("active", rp.active.Get())
	defer span.Finish()
	return rp.get(ctx)
}

func (rp *ResourcePool) get(ctx context.Context) (resource Resource, err error) {
	// If ctx has already expired, avoid racing with rp's resource channel.
	select {
	case <-ctx.Done():
		return nil, ErrCtxTimeout
	default:
	}

	// Fetch
	var wrapper resourceWrapper
	var ok bool
	select {
	case wrapper, ok = <-rp.resources:
	default:
		startTime := time.Now()
		select {
		case wrapper, ok = <-rp.resources:
		case <-ctx.Done():
			return nil, ErrTimeout
		}
		rp.recordWait(startTime)
	}
	if !ok {
		return nil, ErrClosed
	}

	// Unwrap
	if wrapper.resource == nil {
		span, _ := trace.NewSpan(ctx, "ResourcePool.factory")
		wrapper.resource, err = rp.factory(ctx)
		span.Finish()
		if err != nil {
			rp.resources <- resourceWrapper{}
			return nil, err
		}
		rp.active.Add(1)
	}
	if rp.available.Add(-1) <= 0 {
		rp.exhausted.Add(1)
	}
	rp.inUse.Add(1)
	return wrapper.resource, err
}

// Put will return a resource to the pool. For every successful Get,
// a corresponding Put is required. If you no longer need a resource,
// you will need to call Put(nil) instead of returning the closed resource.
// This will cause a new resource to be created in its place.
func (rp *ResourcePool) Put(resource Resource) {
	var wrapper resourceWrapper
	if resource != nil {
		wrapper = resourceWrapper{
			resource: resource,
			timeUsed: time.Now(),
		}
	} else {
		rp.reopenResource(&wrapper)
	}
	select {
	case rp.resources <- wrapper:
	default:
		panic(errors.New("attempt to Put into a full ResourcePool"))
	}
	rp.inUse.Add(-1)
	rp.available.Add(1)
}

func (rp *ResourcePool) reopenResource(wrapper *resourceWrapper) {
	if r, err := rp.factory(context.TODO()); err == nil {
		wrapper.resource = r
		wrapper.timeUsed = time.Now()
	} else {
		wrapper.resource = nil
		rp.active.Add(-1)
	}
}

// SetCapacity changes the capacity of the pool.
// You can use it to shrink or expand, but not beyond
// the max capacity. If the change requires the pool
// to be shrunk, SetCapacity waits till the necessary
// number of resources are returned to the pool.
// A SetCapacity of 0 is equivalent to closing the ResourcePool.
func (rp *ResourcePool) SetCapacity(capacity int) error {
	if capacity < 0 || capacity > cap(rp.resources) {
		return fmt.Errorf("capacity %d is out of range", capacity)
	}

	// Atomically swap new capacity with old, but only
	// if old capacity is non-zero.
	var oldcap int
	for {
		oldcap = int(rp.capacity.Get())
		if oldcap == 0 {
			return ErrClosed
		}
		if oldcap == capacity {
			return nil
		}
		if rp.capacity.CompareAndSwap(int64(oldcap), int64(capacity)) {
			break
		}
	}

	if capacity < oldcap {
		for i := 0; i < oldcap-capacity; i++ {
			wrapper := <-rp.resources
			if wrapper.resource != nil {
				wrapper.resource.Close()
				rp.active.Add(-1)
			}
			rp.available.Add(-1)
		}
	} else {
		for i := 0; i < capacity-oldcap; i++ {
			rp.resources <- resourceWrapper{}
			rp.available.Add(1)
		}
	}
	if capacity == 0 {
		close(rp.resources)
	}
	return nil
}

func (rp *ResourcePool) recordWait(start time.Time) {
	rp.waitCount.Add(1)
	rp.waitTime.Add(time.Since(start))
	if rp.logWait != nil {
		rp.logWait(start)
	}
}

// SetIdleTimeout sets the idle timeout. It can only be used if there was an
// idle timeout set when the pool was created.
func (rp *ResourcePool) SetIdleTimeout(idleTimeout time.Duration) {
	if rp.idleTimer == nil {
		panic("SetIdleTimeout called when timer not initialized")
	}

	rp.idleTimeout.Set(idleTimeout)
	rp.idleTimer.SetInterval(idleTimeout / 10)
}

// StatsJSON returns the stats in JSON format.
func (rp *ResourcePool) StatsJSON() string {
	return fmt.Sprintf(`{"Capacity": %v, "Available": %v, "Active": %v, "InUse": %v, "MaxCapacity": %v, "WaitCount": %v, "WaitTime": %v, "IdleTimeout": %v, "IdleClosed": %v, "Exhausted": %v}`,
		rp.Capacity(),
		rp.Available(),
		rp.Active(),
		rp.InUse(),
		rp.MaxCap(),
		rp.WaitCount(),
		rp.WaitTime().Nanoseconds(),
		rp.IdleTimeout().Nanoseconds(),
		rp.IdleClosed(),
		rp.Exhausted(),
	)
}

// Capacity returns the capacity.
func (rp *ResourcePool) Capacity() int64 {
	return rp.capacity.Get()
}

// Available returns the number of currently unused and available resources.
func (rp *ResourcePool) Available() int64 {
	return rp.available.Get()
}

// Active returns the number of active (i.e. non-nil) resources either in the
// pool or claimed for use
func (rp *ResourcePool) Active() int64 {
	return rp.active.Get()
}

// InUse returns the number of claimed resources from the pool
func (rp *ResourcePool) InUse() int64 {
	return rp.inUse.Get()
}

// MaxCap returns the max capacity.
func (rp *ResourcePool) MaxCap() int64 {
	return int64(cap(rp.resources))
}

// WaitCount returns the total number of waits.
func (rp *ResourcePool) WaitCount() int64 {
	return rp.waitCount.Get()
}

// WaitTime returns the total wait time.
func (rp *ResourcePool) WaitTime() time.Duration {
	return rp.waitTime.Get()
}

// IdleTimeout returns the idle timeout.
func (rp *ResourcePool) IdleTimeout() time.Duration {
	return rp.idleTimeout.Get()
}

// IdleClosed returns the count of resources closed due to idle timeout.
func (rp *ResourcePool) IdleClosed() int64 {
	return rp.idleClosed.Get()
}

// Exhausted returns the number of times Available dropped below 1
func (rp *ResourcePool) Exhausted() int64 {
	return rp.exhausted.Get()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"strings"
)

type Build struct {
	Module  ModuleConfig
	Success bool
	Output  string
}

type Builds []Build

func (build Build) String() string {
	if build.Success {
		return fmt.Sprintf("%s: OK", build.Module.Name)
	} else {
		return fmt.Sprintf("%s: ERROR", build.Module.Name)
	}
}

func (builds Builds) Success() bool {
	for _, build := range builds {
		if !build.Success {
			return false
		}
	}
	return true
}

func (builds Builds) String() string {
	if builds.Success() {
		return "SUCCESS"
	} else {
		return "FAILURE"
	}
}

func BuildModule(module ModuleConfig, directory string) Build {
	fmt.Printf("Building '%s'... ", module.Name)
	moduleDir := path.Join(directory, module.Name)
	// go in build directory
	currentDir, err := os.Getwd()
	defer os.Chdir(currentDir)
	err = os.Chdir(directory)
	if err != nil {
		return Build{
			Module:  module,
			Success: false,
			Output:  err.Error(),
		}
	}
	// delete module directory if it already exists
	if _, err := os.Stat(moduleDir); err == nil {
		os.RemoveAll(moduleDir)
	}
	// git clone the module repository
	cmd := exec.Command("git", "clone", module.Url)
	output, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Println("ERROR")
		return Build{
			Module:  module,
			Success: false,
			Output:  string(output),
		}
	} else {
		defer os.RemoveAll(moduleDir)
		os.Chdir(moduleDir)
		// run the build command
		cmd := exec.Command("bash", "-c", module.Command)
		output, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("ERROR")
			return Build{
				Module:  module,
				Success: false,
				Output:  strings.TrimSpace(string(output)),
			}
		} else {
			fmt.Println("OK")
			return Build{
				Module:  module,
				Success: true,
				Output:  string(output),
			}
		}
	}
}

func BuildModules(config Config) Builds {
	builds := make(Builds, len(config.Modules))
	for index, module := range config.Modules {
		builds[index] = BuildModule(module, config.Directory)
	}
	return builds
}
<commit_msg>Refacto<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"strings"
)

type Build struct {
	Module  ModuleConfig
	Success bool
	Output  string
}

type Builds []Build

func (build Build) String() string {
	if build.Success {
		return fmt.Sprintf("%s: OK", build.Module.Name)
	} else {
		return fmt.Sprintf("%s: ERROR", build.Module.Name)
	}
}

func (builds Builds) Success() bool {
	for _, build := range builds {
		if !build.Success {
			return false
		}
	}
	return true
}

func (builds Builds) String() string {
	if builds.Success() {
		return "SUCCESS"
	} else {
		return "FAILURE"
	}
}

func GetModule(module ModuleConfig) (string, error) {
	cmd := exec.Command("git", "clone", module.Url)
	output, err := cmd.CombinedOutput()
	return string(output), err
}

func BuildModule(module ModuleConfig, directory string) Build {
	fmt.Printf("Building '%s'... ", module.Name)
	moduleDir := path.Join(directory, module.Name)
	// go in build directory
	currentDir, err := os.Getwd()
	defer os.Chdir(currentDir)
	err = os.Chdir(directory)
	if err != nil {
		return Build{
			Module:  module,
			Success: false,
			Output:  err.Error(),
		}
	}
	// delete module directory if it already exists
	if _, err := os.Stat(moduleDir); err == nil {
		os.RemoveAll(moduleDir)
	}
	// git clone the module repository
	output, err := GetModule(module)
	if err != nil {
		fmt.Println("ERROR")
		return Build{
			Module:  module,
			Success: false,
			Output:  string(output),
		}
	} else {
		defer os.RemoveAll(moduleDir)
		os.Chdir(moduleDir)
		// run the build command
		cmd := exec.Command("bash", "-c", module.Command)
		output, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("ERROR")
			return Build{
				Module:  module,
				Success: false,
				Output:  strings.TrimSpace(string(output)),
			}
		} else {
			fmt.Println("OK")
			return Build{
				Module:  module,
				Success: true,
				Output:  string(output),
			}
		}
	}
}

func BuildModules(config Config) Builds {
	builds := make(Builds, len(config.Modules))
	for index, module := range config.Modules {
		builds[index] = BuildModule(module, config.Directory)
	}
	return builds
}
<|endoftext|>


<commit_before>package tpm2tools

import (
	"testing"

	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpmutil"

	"github.com/google/go-tpm-tools/simulator"
)

func getEK(t *testing.T) (*simulator.Simulator, *Key) {
	t.Helper()
	simulator, err := simulator.Get()
	if err != nil {
		t.Fatal(err)
	}
	key, err := EndorsementKeyRSA(simulator)
	if err != nil {
		simulator.Close()
		t.Fatal(err)
	}
	return simulator, key
}

func TestNameMatchesPublicArea(t *testing.T) {
	s, ek := getEK(t)
	defer s.Close()
	defer ek.Close()

	matches, err := ek.Name().MatchesPublic(ek.pubArea)
	if err != nil {
		t.Fatal(err)
	}
	if !matches {
		t.Fatal("Returned name and computed name do not match")
	}
}

func TestCreateSigningKeysInAllHierarchies(t *testing.T) {
	simulator, err := simulator.Get()
	if err != nil {
		t.Fatal(err)
	}

	template := AIKTemplateRSA([256]byte{})
	for _, hierarchy := range []tpmutil.Handle{tpm2.HandleOwner, tpm2.HandleEndorsement, tpm2.HandlePlatform, tpm2.HandleNull} {
		key, err := NewKey(simulator, hierarchy, template)
		if err != nil {
			t.Errorf("Hierarchy %+v: %s", hierarchy, err)
		} else {
			key.Close()
		}
	}
}
<commit_msg>Have tpm2tools use the internal testing functions<commit_after>package tpm2tools

import (
	"testing"

	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpmutil"

	"github.com/google/go-tpm-tools/internal"
)

func TestNameMatchesPublicArea(t *testing.T) {
	rwc := internal.GetTPM(t)
	defer rwc.Close()
	ek, err := EndorsementKeyRSA(rwc)
	if err != nil {
		t.Fatal(err)
	}
	defer ek.Close()

	matches, err := ek.Name().MatchesPublic(ek.pubArea)
	if err != nil {
		t.Fatal(err)
	}
	if !matches {
		t.Fatal("Returned name and computed name do not match")
	}
}

func TestCreateSigningKeysInHierarchies(t *testing.T) {
	rwc := internal.GetTPM(t)
	defer rwc.Close()
	template := AIKTemplateRSA([256]byte{})

	// We are not authorized to create keys in the Platform Hierarchy
	for _, hierarchy := range []tpmutil.Handle{tpm2.HandleOwner, tpm2.HandleEndorsement, tpm2.HandleNull} {
		key, err := NewKey(rwc, hierarchy, template)
		if err != nil {
			t.Errorf("Hierarchy %+v: %s", hierarchy, err)
		} else {
			key.Close()
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

// Package stats implements a means of tracking processing statistics for a
// BitTorrent tracker.
package stats

import "time"

const (
	Announce = iota
	Scrape
	Completed
	NewPeer
	DeletedPeer
	NewSeeder
	DeletedSeeder
	NewTorrent
	DeletedTorrent
)

type Stats struct {
	Start     time.Time
	Announces uint64
	Scrapes   uint64
	Completed uint64
	Peers     uint64
	Seeders   uint64
	Torrents  uint64

	events chan int
}

func New(chanSize int) *Stats {
	s := &Stats{
		Start:  time.Now(),
		events: make(chan int, chanSize),
	}

	go s.handleEvents()

	return s
}

func (s *Stats) Close() {
	close(s.events)
}

func (s *Stats) Uptime() time.Duration {
	return time.Since(s.Start)
}

func (s *Stats) RecordEvent(event int) {
	s.events <- event
}

func (s *Stats) handleEvents() {
	for event := range s.events {
		switch event {
		case Announce:
			s.Announces++

		case Scrape:
			s.Scrapes++

		case Completed:
			s.Completed++

		case NewPeer:
			s.Peers++

		case DeletedPeer:
			s.Peers--

		case NewSeeder:
			s.Seeders++

		case DeletedSeeder:
			s.Seeders--

		case NewTorrent:
			s.Torrents++

		case DeletedTorrent:
			s.Torrents--

		default:
			panic("stats: RecordEvent called with an unknown event")
		}
	}
}
<commit_msg>Flesh out available stats more<commit_after>// Copyright 2014 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

// Package stats implements a means of tracking processing statistics for a
// BitTorrent tracker.
package stats

import "time"

const (
	Announce = iota
	Scrape

	CompletedIPv4
	NewLeechIPv4
	DeletedLeechIPv4
	ReapedLeechIPv4
	NewSeedIPv4
	DeletedSeedIPv4
	ReapedSeedIPv4

	CompletedIPv6
	NewLeechIPv6
	DeletedLeechIPv6
	ReapedLeechIPv6
	NewSeedIPv6
	DeletedSeedIPv6
	ReapedSeedIPv6

	NewTorrent
	DeletedTorrent
	ReapedTorrent

	AcceptedConnection
	ClosedConnection

	HandledRequest
	ErroredRequest
)

type PeerStats struct {
	// Stats for all peers.
	Completed uint64
	Joined    uint64
	Left      uint64
	Reaped    uint64

	// Stats for seeds only (subset of total).
	SeedsJoined uint64
	SeedsLeft   uint64
	SeedsReaped uint64
}

type Stats struct {
	Start time.Time

	Announces uint64
	Scrapes   uint64

	IPv4Peers PeerStats
	IPv6Peers PeerStats

	TorrentsAdded   uint64
	TorrentsRemoved uint64
	TorrentsReaped  uint64

	ActiveConnections   uint64
	ConnectionsAccepted uint64
	BytesTransmitted    uint64

	RequestsHandled uint64
	RequestsErrored uint64

	events chan int
}

func New(chanSize int) *Stats {
	s := &Stats{
		Start:  time.Now(),
		events: make(chan int, chanSize),
	}

	go s.handleEvents()

	return s
}

func (s *Stats) Close() {
	close(s.events)
}

func (s *Stats) Uptime() time.Duration {
	return time.Since(s.Start)
}

func (s *Stats) RecordEvent(event int) {
	s.events <- event
}

func (s *Stats) handleEvents() {
	for event := range s.events {
		switch event {
		case Announce:
			s.Announces++
		case Scrape:
			s.Scrapes++

		case CompletedIPv4:
			s.IPv4Peers.Completed++
		case NewLeechIPv4:
			s.IPv4Peers.Joined++
		case DeletedLeechIPv4:
			s.IPv4Peers.Left++
		case ReapedLeechIPv4:
			s.IPv4Peers.Reaped++
		case NewSeedIPv4:
			s.IPv4Peers.SeedsJoined++
		case DeletedSeedIPv4:
			s.IPv4Peers.SeedsLeft++
		case ReapedSeedIPv4:
			s.IPv4Peers.SeedsReaped++

		case CompletedIPv6:
			s.IPv6Peers.Completed++
		case NewLeechIPv6:
			s.IPv6Peers.Joined++
		case DeletedLeechIPv6:
			s.IPv6Peers.Left++
		case ReapedLeechIPv6:
			s.IPv6Peers.Reaped++
		case NewSeedIPv6:
			s.IPv6Peers.SeedsJoined++
		case DeletedSeedIPv6:
			s.IPv6Peers.SeedsLeft++
		case ReapedSeedIPv6:
			s.IPv6Peers.SeedsReaped++

		case NewTorrent:
			s.TorrentsAdded++
		case DeletedTorrent:
			s.TorrentsRemoved++
		case ReapedTorrent:
			s.TorrentsReaped++

		case AcceptedConnection:
			s.ConnectionsAccepted++
			s.ActiveConnections++

		case ClosedConnection:
			s.ActiveConnections--

		case HandledRequest:
			s.RequestsHandled++

		case ErroredRequest:
			s.RequestsErrored++

		default:
			panic("stats: RecordEvent called with an unknown event")
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"os"

	"github.com/foomo/petze/config"
	"github.com/foomo/petze/service"
)

func main() {
	flagConfigDir := flag.String("config-dir", "", "config-dir")
	flag.Parse()

	if *flagConfigDir == "" {
		flag.Usage()
		os.Exit(1)
	}

	serverConfig, err := config.LoadServer(*flagConfigDir)
	if err != nil {
		log.Fatal(err)
	}
	log.Println(service.Run(serverConfig, *flagConfigDir))
}
<commit_msg>Remove optional arguments from petze CLI<commit_after>package main

import (
	"flag"
	"log"
	"os"

	"github.com/foomo/petze/service"
	"fmt"
	"github.com/foomo/petze/config"
)

func main() {
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}

	configurationDirectory := os.Args[1]
	if stat, err := os.Stat(configurationDirectory); err == nil && stat.IsDir() {
		runServer(configurationDirectory)
	} else {
		log.Fatal("specified configuration directory does not exist or is not a directory")
	}
}

func runServer(configurationDirectory string) {
	serverConfig, err := config.LoadServer(configurationDirectory)
	if err != nil {
		log.Fatal(err)
	}
	log.Println(service.Run(serverConfig, configurationDirectory))
}

func usage() {
	fmt.Printf("Usage: %s configuration-directory \n", os.Args[0])
	flag.PrintDefaults()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"go/parser"
	"go/token"
	"os"
	"strings"
)

func init() {
	// declare sub-command "gourd gen service"
	appendCmds("gen", cli.Command{
		Name:    "service",
		Aliases: []string{"s"},
		Usage:   "generate service from a type",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "type, t",
				Value: "",
				Usage: "comma-separated list of type names; must be set",
			},
			cli.StringFlag{
				Name:  "output, o",
				Value: "",
				Usage: "output file name; default srcdir/<type>_service.go",
			},
		},
		Action: genService,
	})
}

// generate the service go file
func genService(c *cli.Context) {

	// files to parse
	fns := make([]string, 0)
	if len(c.Args()) == 0 {
		// TODO: find all .go in the current folder
		fmt.Println("Please provide files to parse")
		os.Exit(1)
	} else {
		fns = c.Args()
	}

	// target type
	if c.String("type") == "" {
		fmt.Println("Please provide the target type(s)")
		os.Exit(1)
	}
	tn := c.String("type")

	// output file
	var o string
	if c.String("output") == "" {
		o = strings.ToLower(tn) + "_service.go"
	} else {
		o = c.String("output")
	}

	// read type of type name from given file(s)
	t := readTypeFile(fns[0], tn)

	// dummy output
	fmt.Printf("t: %#v ==> output: %s\n", t, o)
}

// read typeSpec from files
func readTypeFile(inputPath string, tn string) (spec genTypeSpec) {
	fset := token.NewFileSet()

	// TODO: get to know what inputPath can be (e.g. dir?)
	f, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments)
	if err != nil {
		panic(err)
	}

	// read package name
	var pkg string
	if f.Name == nil {
		fmt.Println("Error. Unknown package name. Exit.")
		os.Exit(1)
	} else {
		pkg = f.Name.Name
	}

	// read types name and details
	for ts := range filterTypeSpec(filterGenDecl(chanDecls(f.Decls))) {
		pts := parseTypeSpec(ts)
		if pts.Name == tn {
			spec = pts
			break
		}
	}
	spec.Pkg = pkg
	return
}
<commit_msg>Improve code style with golint suggestion<commit_after>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"go/parser"
	"go/token"
	"os"
	"strings"
)

func init() {
	// declare sub-command "gourd gen service"
	appendCmds("gen", cli.Command{
		Name:    "service",
		Aliases: []string{"s"},
		Usage:   "generate service from a type",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "type, t",
				Value: "",
				Usage: "comma-separated list of type names; must be set",
			},
			cli.StringFlag{
				Name:  "output, o",
				Value: "",
				Usage: "output file name; default srcdir/<type>_service.go",
			},
		},
		Action: genService,
	})
}

// generate the service go file
func genService(c *cli.Context) {

	// files to parse
	var fns []string
	if len(c.Args()) == 0 {
		// TODO: find all .go in the current folder
		fmt.Println("Please provide files to parse")
		os.Exit(1)
	} else {
		fns = c.Args()
	}

	// target type
	if c.String("type") == "" {
		fmt.Println("Please provide the target type(s)")
		os.Exit(1)
	}
	tn := c.String("type")

	// output file
	var o string
	if c.String("output") == "" {
		o = strings.ToLower(tn) + "_service.go"
	} else {
		o = c.String("output")
	}

	// read type of type name from given file(s)
	t := readTypeFile(fns[0], tn)

	// dummy output
	fmt.Printf("t: %#v ==> output: %s\n", t, o)
}

// read typeSpec from files
func readTypeFile(inputPath string, tn string) (spec genTypeSpec) {
	fset := token.NewFileSet()

	// TODO: get to know what inputPath can be (e.g. dir?)
	f, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments)
	if err != nil {
		panic(err)
	}

	// read package name
	var pkg string
	if f.Name == nil {
		fmt.Println("Error. Unknown package name. Exit.")
		os.Exit(1)
	} else {
		pkg = f.Name.Name
	}

	// read types name and details
	for ts := range filterTypeSpec(filterGenDecl(chanDecls(f.Decls))) {
		pts := parseTypeSpec(ts)
		if pts.Name == tn {
			spec = pts
			break
		}
	}
	spec.Pkg = pkg
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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 spdy

import (
	"fmt"
	"net/http"
	"net/url"
	"time"

	"k8s.io/apimachinery/pkg/util/httpstream"
	"k8s.io/apimachinery/pkg/util/httpstream/spdy"
	restclient "k8s.io/client-go/rest"
)

// Upgrader validates a response from the server after a SPDY upgrade.
type Upgrader interface {
	// NewConnection validates the response and creates a new Connection.
	NewConnection(resp *http.Response) (httpstream.Connection, error)
}

// RoundTripperFor returns a round tripper and upgrader to use with SPDY.
func RoundTripperFor(config *restclient.Config) (http.RoundTripper, Upgrader, error) {
	tlsConfig, err := restclient.TLSConfigFor(config)
	if err != nil {
		return nil, nil, err
	}
	proxy := http.ProxyFromEnvironment
	if config.Proxy != nil {
		proxy = config.Proxy
	}
	upgradeRoundTripper := spdy.NewRoundTripperWithConfig(spdy.RoundTripperConfig{
		TLS:                      tlsConfig,
		FollowRedirects:          true,
		RequireSameHostRedirects: false,
		Proxier:                  proxy,
		PingPeriod:               time.Second * 5,
	})
	wrapper, err := restclient.HTTPWrappersForConfig(config, upgradeRoundTripper)
	if err != nil {
		return nil, nil, err
	}
	return wrapper, upgradeRoundTripper, nil
}

// dialer implements the httpstream.Dialer interface.
type dialer struct {
	client   *http.Client
	upgrader Upgrader
	method   string
	url      *url.URL
}

var _ httpstream.Dialer = &dialer{}

// NewDialer will create a dialer that connects to the provided URL and upgrades the connection to SPDY.
func NewDialer(upgrader Upgrader, client *http.Client, method string, url *url.URL) httpstream.Dialer {
	return &dialer{
		client:   client,
		upgrader: upgrader,
		method:   method,
		url:      url,
	}
}

func (d *dialer) Dial(protocols ...string) (httpstream.Connection, string, error) {
	req, err := http.NewRequest(d.method, d.url.String(), nil)
	if err != nil {
		return nil, "", fmt.Errorf("error creating request: %v", err)
	}
	return Negotiate(d.upgrader, d.client, req, protocols...)
}

// Negotiate opens a connection to a remote server and attempts to negotiate
// a SPDY connection. Upon success, it returns the connection and the protocol selected by
// the server. The client transport must use the upgradeRoundTripper - see RoundTripperFor.
func Negotiate(upgrader Upgrader, client *http.Client, req *http.Request, protocols ...string) (httpstream.Connection, string, error) {
	for i := range protocols {
		req.Header.Add(httpstream.HeaderProtocolVersion, protocols[i])
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, "", fmt.Errorf("error sending request: %v", err)
	}
	defer resp.Body.Close()
	conn, err := upgrader.NewConnection(resp)
	if err != nil {
		return nil, "", err
	}
	return conn, resp.Header.Get(httpstream.HeaderProtocolVersion), nil
}
<commit_msg>Don't follow redirects with spdy<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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 spdy

import (
	"fmt"
	"net/http"
	"net/url"
	"time"

	"k8s.io/apimachinery/pkg/util/httpstream"
	"k8s.io/apimachinery/pkg/util/httpstream/spdy"
	restclient "k8s.io/client-go/rest"
)

// Upgrader validates a response from the server after a SPDY upgrade.
type Upgrader interface {
	// NewConnection validates the response and creates a new Connection.
	NewConnection(resp *http.Response) (httpstream.Connection, error)
}

// RoundTripperFor returns a round tripper and upgrader to use with SPDY.
func RoundTripperFor(config *restclient.Config) (http.RoundTripper, Upgrader, error) {
	tlsConfig, err := restclient.TLSConfigFor(config)
	if err != nil {
		return nil, nil, err
	}
	proxy := http.ProxyFromEnvironment
	if config.Proxy != nil {
		proxy = config.Proxy
	}
	upgradeRoundTripper := spdy.NewRoundTripperWithConfig(spdy.RoundTripperConfig{
		TLS:        tlsConfig,
		Proxier:    proxy,
		PingPeriod: time.Second * 5,
	})
	wrapper, err := restclient.HTTPWrappersForConfig(config, upgradeRoundTripper)
	if err != nil {
		return nil, nil, err
	}
	return wrapper, upgradeRoundTripper, nil
}

// dialer implements the httpstream.Dialer interface.
type dialer struct {
	client   *http.Client
	upgrader Upgrader
	method   string
	url      *url.URL
}

var _ httpstream.Dialer = &dialer{}

// NewDialer will create a dialer that connects to the provided URL and upgrades the connection to SPDY.
func NewDialer(upgrader Upgrader, client *http.Client, method string, url *url.URL) httpstream.Dialer {
	return &dialer{
		client:   client,
		upgrader: upgrader,
		method:   method,
		url:      url,
	}
}

func (d *dialer) Dial(protocols ...string) (httpstream.Connection, string, error) {
	req, err := http.NewRequest(d.method, d.url.String(), nil)
	if err != nil {
		return nil, "", fmt.Errorf("error creating request: %v", err)
	}
	return Negotiate(d.upgrader, d.client, req, protocols...)
}

// Negotiate opens a connection to a remote server and attempts to negotiate
// a SPDY connection. Upon success, it returns the connection and the protocol selected by
// the server. The client transport must use the upgradeRoundTripper - see RoundTripperFor.
func Negotiate(upgrader Upgrader, client *http.Client, req *http.Request, protocols ...string) (httpstream.Connection, string, error) {
	for i := range protocols {
		req.Header.Add(httpstream.HeaderProtocolVersion, protocols[i])
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, "", fmt.Errorf("error sending request: %v", err)
	}
	defer resp.Body.Close()
	conn, err := upgrader.NewConnection(resp)
	if err != nil {
		return nil, "", err
	}
	return conn, resp.Header.Get(httpstream.HeaderProtocolVersion), nil
}
<|endoftext|>


<commit_before>package main

import (
	"strconv"
	"time"

	"github.com/everpcpc/pixiv"
)

func initPixiv(cfg *PixivConfig) error {
	tokenKey := "pixiv:" + cfg.Username + ":auth"
	token := redisClient.HGet(tokenKey, "token").String()
	refreshToken := redisClient.HGet(tokenKey, "refresh_token").String()
	tokenDeadline, _ := time.Parse(time.RFC3339, redisClient.HGet(tokenKey, "token_deadline").String())
	pixiv.HookAuth(func(t, rt string, td time.Time) error {
		v := map[string]interface{}{
			"token":          t,
			"refresh_token":  rt,
			"token_deadline": td.Format(time.RFC3339),
		}
		return redisClient.HMSet(tokenKey, v).Err()
	})

	var account *pixiv.Account
	var err error
	if token+refreshToken == "" {
		account, err = pixiv.Login(cfg.Username, cfg.Password)
	} else {
		account, err = pixiv.LoadAuth(token, refreshToken, tokenDeadline)
	}
	if err == nil {
		logger.Debugf("pixiv: %+v", account)
	}
	return err
}

func pixivFollow(cfg *PixivConfig, ttl int) {
	if err := initPixiv(cfg); err != nil {
		logger.Fatal(err)
	}
	if ttl < 10 {
		ttl = 10
	}
	papp := pixiv.NewApp()
	ticker := time.Tick(time.Duration(ttl) * time.Second)
	maxIDKey := "pixiv:"
	for {
		<-ticker
		illusts, err := papp.IllustFollow("public", 0)
		if err != nil {
			logger.Error(err)
			continue
		}
		maxID, err := redisClient.Get(maxIDKey).Uint64()
		if err != nil {
			logger.Error(err)
			continue
		}
		if len(illusts) == 0 {
			continue
		}
		for i := range illusts {
			if maxID >= illusts[i].ID {
				break
			}
			go telegramBot.send(telegramBot.SelfChatID, pixivURL(illusts[i].ID))
		}
		if err := redisClient.Set(maxIDKey, illusts[0].ID, 0).Err(); err != nil {
			logger.Error(err)
		}
	}
}

func pixivURL(id uint64) string {
	return "https://www.pixiv.net/member_illust.php?mode=medium&illust_id=" + strconv.FormatUint(id, 10)
}
<commit_msg>log auth<commit_after>package main

import (
	"strconv"
	"time"

	"github.com/everpcpc/pixiv"
)

func initPixiv(cfg *PixivConfig) error {
	tokenKey := "pixiv:" + cfg.Username + ":auth"
	token := redisClient.HGet(tokenKey, "token").String()
	refreshToken := redisClient.HGet(tokenKey, "refresh_token").String()
	tokenDeadline, _ := time.Parse(time.RFC3339, redisClient.HGet(tokenKey, "token_deadline").String())
	pixiv.HookAuth(func(t, rt string, td time.Time) error {
		v := map[string]interface{}{
			"token":          t,
			"refresh_token":  rt,
			"token_deadline": td.Format(time.RFC3339),
		}
		return redisClient.HMSet(tokenKey, v).Err()
	})

	var account *pixiv.Account
	var err error
	if token+refreshToken == "" {
		logger.Debugf("login with %s", cfg.Username)
		account, err = pixiv.Login(cfg.Username, cfg.Password)
	} else {
		logger.Debugf("load auth with %+v", tokenDeadline)
		account, err = pixiv.LoadAuth(token, refreshToken, tokenDeadline)
	}
	if err == nil {
		logger.Debugf("pixiv: %+v", account)
	}
	return err
}

func pixivFollow(cfg *PixivConfig, ttl int) {
	if err := initPixiv(cfg); err != nil {
		logger.Fatal(err)
	}
	if ttl < 10 {
		ttl = 10
	}
	papp := pixiv.NewApp()
	ticker := time.Tick(time.Duration(ttl) * time.Second)
	maxIDKey := "pixiv:"
	for {
		<-ticker
		illusts, err := papp.IllustFollow("public", 0)
		if err != nil {
			logger.Error(err)
			continue
		}
		maxID, err := redisClient.Get(maxIDKey).Uint64()
		if err != nil {
			logger.Error(err)
			continue
		}
		if len(illusts) == 0 {
			continue
		}
		for i := range illusts {
			if maxID >= illusts[i].ID {
				break
			}
			go telegramBot.send(telegramBot.SelfChatID, pixivURL(illusts[i].ID))
		}
		if err := redisClient.Set(maxIDKey, illusts[0].ID, 0).Err(); err != nil {
			logger.Error(err)
		}
	}
}

func pixivURL(id uint64) string {
	return "https://www.pixiv.net/member_illust.php?mode=medium&illust_id=" + strconv.FormatUint(id, 10)
}
<|endoftext|>