<commit_before>package peco

import (
	"fmt"
	"unicode/utf8"

	"github.com/mattn/go-runewidth"
	"github.com/nsf/termbox-go"
)

type View struct {
	*Ctx
}

type PagingRequest int

const (
	ToNextLine PagingRequest = iota
	ToNextPage
	ToPrevLine
	ToPrevPage
)

func (u *View) Loop() {
	defer u.ReleaseWaitGroup()
	for {
		select {
		case <-u.LoopCh():
			return
		case r := <-u.PagingCh():
			u.movePage(r)
		case lines := <-u.DrawCh():
			u.drawScreen(lines)
		}
	}
}

func (v *View) printStatus() {
	w, h := termbox.Size()

	msg := v.statusMessage
	width := runewidth.StringWidth(msg)

	pad := make([]byte, w - width)
	for i := 0; i < w - width; i++ {
		pad[i] = ' '
	}

	printTB(0, h - 2, termbox.ColorDefault, termbox.ColorDefault, string(pad))
	if width > 0 {
		printTB(w - width, h - 2, termbox.AttrReverse|termbox.ColorDefault|termbox.AttrBold, termbox.AttrReverse|termbox.ColorDefault, msg)
	}
}

func printTB(x, y int, fg, bg termbox.Attribute, msg string) {
	for len(msg) > 0 {
		c, w := utf8.DecodeRuneInString(msg)
		if c == utf8.RuneError {
			c = '?'
			w = 1
		}
		msg = msg[w:]
		termbox.SetCell(x, y, c, fg, bg)
		x += runewidth.RuneWidth(c)
	}

	width, _ := termbox.Size()
	for ; x < width; x++ {
		termbox.SetCell(x, y, ' ', fg, bg)
	}
}

func (v *View) movePage(p PagingRequest) {
	_, height := termbox.Size()
	perPage := height - 4

	switch p {
	case ToPrevLine:
		v.currentLine--
	case ToNextLine:
		v.currentLine++
	case ToPrevPage, ToNextPage:
		if p == ToPrevPage {
			v.currentLine -= perPage
		} else {
			v.currentLine += perPage
		}
	}

	if v.currentLine < 1 {
		if v.current != nil {
			// Go to last page, if possible
			v.currentLine = len(v.current)
		} else {
			v.currentLine = 1
		}
	} else if v.current != nil && v.currentLine > len(v.current) {
		v.currentLine = 1
	}
	v.drawScreen(nil)
}

func (u *View) drawScreen(targets []Match) {
	u.mutex.Lock()
	defer u.mutex.Unlock()

	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	if targets == nil {
		if current := u.Ctx.current; current != nil {
			targets = u.Ctx.current
		} else {
			targets = u.Ctx.lines
		}
	}

	width, height := termbox.Size()
	perPage := height - 4

CALCULATE_PAGE:
	currentPage := ((u.Ctx.currentLine - 1) / perPage) + 1
	if currentPage <= 0 {
		currentPage = 1
	}
	offset := (currentPage - 1) * perPage
	var maxPage int
	if len(targets) == 0 {
		maxPage = 1
	} else {
		maxPage = ((len(targets) + perPage - 1) / perPage)
	}

	if maxPage < currentPage {
		u.Ctx.currentLine = offset
		goto CALCULATE_PAGE
	}

	prompt := "QUERY>"
	promptLen := runewidth.StringWidth(prompt)
	printTB(0, 0, termbox.ColorDefault, termbox.ColorDefault, prompt)

	if u.caretPos <= 0 {
		u.caretPos = 0 // sanity
	}
	if u.caretPos > len(u.query) {
		u.caretPos = len(u.query)
	}

	if u.caretPos == len(u.query) {
		// the entire string + the caret after the string
		printTB(promptLen+1, 0, termbox.ColorDefault, termbox.ColorDefault, string(u.query))
		termbox.SetCell(promptLen+1+runewidth.StringWidth(string(u.query)), 0, ' ', termbox.ColorDefault|termbox.AttrReverse, termbox.ColorDefault|termbox.AttrReverse)
	} else {
		// the caret is in the middle of the string
		prev := 0
		for i, r := range u.query {
			fg := termbox.ColorDefault
			bg := termbox.ColorDefault
			if i == u.caretPos {
				fg |= termbox.AttrReverse
				bg |= termbox.AttrReverse
			}
			termbox.SetCell(promptLen+1+prev, 0, r, fg, bg)
			prev += runewidth.RuneWidth(r)
		}
	}

	pmsg := fmt.Sprintf("%s [%d/%d]", u.Ctx.Matcher().String(), currentPage, maxPage)

	printTB(width-runewidth.StringWidth(pmsg), 0, termbox.ColorDefault, termbox.ColorDefault, pmsg)

	for n := 1; n <= perPage; n++ {
		fgAttr := u.config.Style.Basic.fg
		bgAttr := u.config.Style.Basic.bg
		if n+offset == u.currentLine {
			fgAttr = u.config.Style.Selected.fg
			bgAttr = u.config.Style.Selected.bg
		} else if u.selection.Has(n+offset) {
			fgAttr = u.config.Style.SavedSelection.fg
			bgAttr = u.config.Style.SavedSelection.bg
		}

		targetIdx := offset + n - 1
		if targetIdx >= len(targets) {
			break
		}

		target := targets[targetIdx]
		line   := target.Line()
		matches := target.Indices()
		if matches == nil {
			printTB(0, n, fgAttr, bgAttr, line)
		} else {
			prev := 0
			index := 0
			for _, m := range matches {
				if m[0] > index {
					c := line[index:m[0]]
					printTB(prev, n, fgAttr, bgAttr, c)
					prev += runewidth.StringWidth(c)
					index += len(c)
				}
				c := line[m[0]:m[1]]
				printTB(prev, n, u.config.Style.Query.fg, bgAttr|u.config.Style.Query.bg, c)
				prev += runewidth.StringWidth(c)
				index += len(c)
			}

			m := matches[len(matches)-1]
			if m[0] > prev {
				printTB(prev, n, u.config.Style.Query.fg, bgAttr|u.config.Style.Query.bg, line[m[0]:m[1]])
			} else if len(line) > m[1] {
				printTB(prev, n, fgAttr, bgAttr, line[m[1]:len(line)])
			}
		}
	}

	u.printStatus()
	termbox.Flush()

	// FIXME
	u.current = targets
}
<commit_msg>some linting... I mean, juuuust a little bit<commit_after>package peco

import (
	"fmt"
	"unicode/utf8"

	"github.com/mattn/go-runewidth"
	"github.com/nsf/termbox-go"
)

// View handles the drawing/updating the screen
type View struct {
	*Ctx
}

// PagingRequest can be sent to move the selection cursor
type PagingRequest int

const (
	// ToNextLine moves the selection to the next line
	ToNextLine PagingRequest = iota
	// ToNextPage moves the selection to the next page
	ToNextPage
	// ToPrevLine moves the selection to the previous line
	ToPrevLine
	// ToPrevPage moves the selection to the previous page
	ToPrevPage
)

// Loop receives requests to update the screen
func (v *View) Loop() {
	defer u.ReleaseWaitGroup()
	for {
		select {
		case <-u.LoopCh():
			return
		case r := <-u.PagingCh():
			u.movePage(r)
		case lines := <-u.DrawCh():
			u.drawScreen(lines)
		}
	}
}

func (v *View) printStatus() {
	w, h := termbox.Size()

	msg := v.statusMessage
	width := runewidth.StringWidth(msg)

	pad := make([]byte, w - width)
	for i := 0; i < w - width; i++ {
		pad[i] = ' '
	}

	printTB(0, h - 2, termbox.ColorDefault, termbox.ColorDefault, string(pad))
	if width > 0 {
		printTB(w - width, h - 2, termbox.AttrReverse|termbox.ColorDefault|termbox.AttrBold, termbox.AttrReverse|termbox.ColorDefault, msg)
	}
}

func printTB(x, y int, fg, bg termbox.Attribute, msg string) {
	for len(msg) > 0 {
		c, w := utf8.DecodeRuneInString(msg)
		if c == utf8.RuneError {
			c = '?'
			w = 1
		}
		msg = msg[w:]
		termbox.SetCell(x, y, c, fg, bg)
		x += runewidth.RuneWidth(c)
	}

	width, _ := termbox.Size()
	for ; x < width; x++ {
		termbox.SetCell(x, y, ' ', fg, bg)
	}
}

func (v *View) movePage(p PagingRequest) {
	_, height := termbox.Size()
	perPage := height - 4

	switch p {
	case ToPrevLine:
		v.currentLine--
	case ToNextLine:
		v.currentLine++
	case ToPrevPage, ToNextPage:
		if p == ToPrevPage {
			v.currentLine -= perPage
		} else {
			v.currentLine += perPage
		}
	}

	if v.currentLine < 1 {
		if v.current != nil {
			// Go to last page, if possible
			v.currentLine = len(v.current)
		} else {
			v.currentLine = 1
		}
	} else if v.current != nil && v.currentLine > len(v.current) {
		v.currentLine = 1
	}
	v.drawScreen(nil)
}

func (v *View) drawScreen(targets []Match) {
	u.mutex.Lock()
	defer u.mutex.Unlock()

	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	if targets == nil {
		if current := u.Ctx.current; current != nil {
			targets = u.Ctx.current
		} else {
			targets = u.Ctx.lines
		}
	}

	width, height := termbox.Size()
	perPage := height - 4

CALCULATE_PAGE:
	currentPage := ((u.Ctx.currentLine - 1) / perPage) + 1
	if currentPage <= 0 {
		currentPage = 1
	}
	offset := (currentPage - 1) * perPage
	var maxPage int
	if len(targets) == 0 {
		maxPage = 1
	} else {
		maxPage = ((len(targets) + perPage - 1) / perPage)
	}

	if maxPage < currentPage {
		u.Ctx.currentLine = offset
		goto CALCULATE_PAGE
	}

	prompt := "QUERY>"
	promptLen := runewidth.StringWidth(prompt)
	printTB(0, 0, termbox.ColorDefault, termbox.ColorDefault, prompt)

	if u.caretPos <= 0 {
		u.caretPos = 0 // sanity
	}
	if u.caretPos > len(u.query) {
		u.caretPos = len(u.query)
	}

	if u.caretPos == len(u.query) {
		// the entire string + the caret after the string
		printTB(promptLen+1, 0, termbox.ColorDefault, termbox.ColorDefault, string(u.query))
		termbox.SetCell(promptLen+1+runewidth.StringWidth(string(u.query)), 0, ' ', termbox.ColorDefault|termbox.AttrReverse, termbox.ColorDefault|termbox.AttrReverse)
	} else {
		// the caret is in the middle of the string
		prev := 0
		for i, r := range u.query {
			fg := termbox.ColorDefault
			bg := termbox.ColorDefault
			if i == u.caretPos {
				fg |= termbox.AttrReverse
				bg |= termbox.AttrReverse
			}
			termbox.SetCell(promptLen+1+prev, 0, r, fg, bg)
			prev += runewidth.RuneWidth(r)
		}
	}

	pmsg := fmt.Sprintf("%s [%d/%d]", u.Ctx.Matcher().String(), currentPage, maxPage)

	printTB(width-runewidth.StringWidth(pmsg), 0, termbox.ColorDefault, termbox.ColorDefault, pmsg)

	for n := 1; n <= perPage; n++ {
		fgAttr := u.config.Style.Basic.fg
		bgAttr := u.config.Style.Basic.bg
		if n+offset == u.currentLine {
			fgAttr = u.config.Style.Selected.fg
			bgAttr = u.config.Style.Selected.bg
		} else if u.selection.Has(n+offset) {
			fgAttr = u.config.Style.SavedSelection.fg
			bgAttr = u.config.Style.SavedSelection.bg
		}

		targetIdx := offset + n - 1
		if targetIdx >= len(targets) {
			break
		}

		target := targets[targetIdx]
		line   := target.Line()
		matches := target.Indices()
		if matches == nil {
			printTB(0, n, fgAttr, bgAttr, line)
		} else {
			prev := 0
			index := 0
			for _, m := range matches {
				if m[0] > index {
					c := line[index:m[0]]
					printTB(prev, n, fgAttr, bgAttr, c)
					prev += runewidth.StringWidth(c)
					index += len(c)
				}
				c := line[m[0]:m[1]]
				printTB(prev, n, u.config.Style.Query.fg, bgAttr|u.config.Style.Query.bg, c)
				prev += runewidth.StringWidth(c)
				index += len(c)
			}

			m := matches[len(matches)-1]
			if m[0] > prev {
				printTB(prev, n, u.config.Style.Query.fg, bgAttr|u.config.Style.Query.bg, line[m[0]:m[1]])
			} else if len(line) > m[1] {
				printTB(prev, n, fgAttr, bgAttr, line[m[1]:len(line)])
			}
		}
	}

	u.printStatus()
	termbox.Flush()

	// FIXME
	u.current = targets
}
<|endoftext|>


<commit_before>package gateway

import (
	"net"
	"net/http"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

type webServer struct {
	gw *Gateway

	name       string
	maxClients int

	router *httprouter.Router

	httpListener net.Listener
	httpServer   *http.Server

	httpsListener net.Listener
	httpsServer   *http.Server

	waitExitFunc    waitExitFunc
	connStateFunc   connStateFunc
	onConnNewFunc   onConnNewFunc
	onConnCloseFunc onConnCloseFunc

	onStop        func()
	mu            sync.Mutex
	waiterStarted bool

	// FIXME if http/https listener both enabled, must able to tell them apart
	activeConnN int32

	// TODO channel performance is frustrating, no better than mutex/map use ring buffer
	stateIdleCh, stateRemoveCh, stateActiveCh chan net.Conn
}

func newWebServer(name string, httpAddr, httpsAddr string, maxClients int,
	gw *Gateway) *webServer {
	const initialConnBuckets = 200
	this := &webServer{
		name:          name,
		gw:            gw,
		maxClients:    maxClients,
		stateActiveCh: make(chan net.Conn, initialConnBuckets),
		stateIdleCh:   make(chan net.Conn, initialConnBuckets),
		stateRemoveCh: make(chan net.Conn, initialConnBuckets),
		router:        httprouter.New(),
	}

	if Options.EnableHttpPanicRecover {
		this.router.PanicHandler = func(w http.ResponseWriter,
			r *http.Request, err interface{}) {
			log.Error("%s[%s] %s %s: %+v", this.name, r.RemoteAddr, r.Method, r.RequestURI, err)
		}
	}

	if httpAddr != "" {
		this.httpServer = &http.Server{
			Addr:           httpAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	if httpsAddr != "" {
		this.httpsServer = &http.Server{
			Addr:           httpsAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	return this
}

func (this *webServer) Router() *httprouter.Router {
	return this.router
}

func (this *webServer) Start() {
	if this.waitExitFunc == nil {
		this.waitExitFunc = this.defaultWaitExit
	}
	if this.connStateFunc == nil {
		this.connStateFunc = this.defaultConnStateHook

		go this.manageIdleConns()
	}

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
		this.startServer(true)
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
		this.startServer(false)
	}

}

func (this *webServer) startServer(https bool) {
	var err error
	waitListenerUp := make(chan struct{})
	go func() {
		if Options.CpuAffinity {
			runtime.LockOSThread()
		}

		var (
			retryDelay         time.Duration
			theListener        net.Listener
			waitListenerUpOnce sync.Once
		)
		for {
			if https {
				this.httpsListener, err = net.Listen("tcp", this.httpsServer.Addr)
				if err != nil {
					if retryDelay == 0 {
						retryDelay = 50 * time.Millisecond
					} else {
						retryDelay = 2 * retryDelay
					}
					if maxDelay := time.Second; retryDelay > maxDelay {
						retryDelay = maxDelay
					}
					log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
					time.Sleep(retryDelay)
					continue
				}

				this.httpsListener, err = setupHttpsListener(this.httpsListener,
					this.gw.certFile, this.gw.keyFile)
				if err != nil {
					panic(err)
				}

				theListener = this.httpsListener
			} else {
				this.httpListener, err = net.Listen("tcp", this.httpServer.Addr)
				theListener = this.httpListener
			}

			if err != nil {
				if retryDelay == 0 {
					retryDelay = 50 * time.Millisecond
				} else {
					retryDelay = 2 * retryDelay
				}
				if maxDelay := time.Second; retryDelay > maxDelay {
					retryDelay = maxDelay
				}
				log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
				time.Sleep(retryDelay)
				continue
			}

			theListener = LimitListener(this.name, this.gw, theListener, this.maxClients)
			waitListenerUpOnce.Do(func() {
				close(waitListenerUp)
			})

			// on non-temporary err, net/http will close the listener
			if https {
				err = this.httpsServer.Serve(theListener)
			} else {
				err = this.httpServer.Serve(theListener)
			}

			select {
			case <-this.gw.shutdownCh:
				return

			default:
				log.Error("%s server: %v", this.name, err)
			}
		}
	}()

	// FIXME if net.Listen fails, kateway will not be able to stop
	// e,g. start a kateway, then start another, the 2nd will not be able to stop

	// wait for the listener up
	<-waitListenerUp

	this.mu.Lock()
	if !this.waiterStarted {
		this.waiterStarted = true

		this.gw.wg.Add(1)
		go this.waitExitFunc(this.gw.shutdownCh)
	}
	this.mu.Unlock()

	if https {
		log.Info("%s https server ready on %s", this.name, this.httpsServer.Addr)
	} else {
		log.Info("%s http server ready on %s", this.name, this.httpServer.Addr)
	}
}

func (this *webServer) defaultConnStateHook(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		atomic.AddInt32(&this.activeConnN, 1)

		if this.onConnNewFunc != nil {
			this.onConnNewFunc(c)
		}

	case http.StateIdle:
		this.stateIdleCh <- c

	case http.StateActive:
		this.stateActiveCh <- c

	case http.StateClosed, http.StateHijacked:
		atomic.AddInt32(&this.activeConnN, -1)
		this.stateRemoveCh <- c

		if this.onConnCloseFunc != nil {
			this.onConnCloseFunc(c)
		}
	}
}

func (this *webServer) manageIdleConns() {
	var (
		idleConns     = make(map[net.Conn]struct{}, 200)
		c             net.Conn
		waitNextRound = make(chan struct{})
	)
	defer close(waitNextRound)

	for {
		select {
		case <-this.gw.shutdownCh:
			if len(idleConns) == 0 {
				// happy ending
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case <-waitNextRound:
			if len(idleConns) == 0 {
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case c = <-this.stateActiveCh:
			delete(idleConns, c)

		case c = <-this.stateIdleCh:
			idleConns[c] = struct{}{}

		case c = <-this.stateRemoveCh:
			delete(idleConns, c)
		}
	}
}

func (this *webServer) defaultWaitExit(exit <-chan struct{}) {
	<-exit

	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpsListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	// wait for all established http/https conns close
	waitStart := time.Now()
	var prompt sync.Once
	for {
		activeConnN := atomic.LoadInt32(&this.activeConnN)
		if activeConnN == 0 {
			// good luck, all connections finished
			break
		}

		prompt.Do(func() {
			log.Trace("%s waiting for %d clients shutdown...", this.name, activeConnN)
		})

		// timeout mechanism
		if time.Since(waitStart) > Options.SubTimeout+time.Second {
			log.Warn("%s still left %d conns after %s, forced to shutdown",
				this.name, activeConnN, Options.SubTimeout)
			break
		}

		time.Sleep(time.Millisecond * 50)
	}
	log.Trace("%s all connections finished", this.name)

	if this.httpsServer != nil {
		this.httpsServer.ConnState = nil
	}
	if this.httpServer != nil {
		this.httpServer.ConnState = nil
	}

	close(this.stateActiveCh)
	close(this.stateIdleCh)
	close(this.stateRemoveCh)

	if this.onStop != nil {
		this.onStop()
	}

	this.gw.wg.Done()
}

func (this *webServer) notFoundHandler(w http.ResponseWriter, r *http.Request) {
	log.Error("%s: not found %s", this.name, r.RequestURI)

	writeNotFound(w)
}
<commit_msg>handle nil pointer issue<commit_after>package gateway

import (
	"net"
	"net/http"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

type webServer struct {
	gw *Gateway

	name       string
	maxClients int

	router *httprouter.Router

	httpListener net.Listener
	httpServer   *http.Server

	httpsListener net.Listener
	httpsServer   *http.Server

	waitExitFunc    waitExitFunc
	connStateFunc   connStateFunc
	onConnNewFunc   onConnNewFunc
	onConnCloseFunc onConnCloseFunc

	onStop        func()
	mu            sync.Mutex
	waiterStarted bool

	// FIXME if http/https listener both enabled, must able to tell them apart
	activeConnN int32

	// TODO channel performance is frustrating, no better than mutex/map use ring buffer
	stateIdleCh, stateRemoveCh, stateActiveCh chan net.Conn
}

func newWebServer(name string, httpAddr, httpsAddr string, maxClients int,
	gw *Gateway) *webServer {
	const initialConnBuckets = 200
	this := &webServer{
		name:          name,
		gw:            gw,
		maxClients:    maxClients,
		stateActiveCh: make(chan net.Conn, initialConnBuckets),
		stateIdleCh:   make(chan net.Conn, initialConnBuckets),
		stateRemoveCh: make(chan net.Conn, initialConnBuckets),
		router:        httprouter.New(),
	}

	if Options.EnableHttpPanicRecover {
		this.router.PanicHandler = func(w http.ResponseWriter,
			r *http.Request, err interface{}) {
			log.Error("%s[%s] %s %s: %+v", this.name, r.RemoteAddr, r.Method, r.RequestURI, err)
		}
	}

	if httpAddr != "" {
		this.httpServer = &http.Server{
			Addr:           httpAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	if httpsAddr != "" {
		this.httpsServer = &http.Server{
			Addr:           httpsAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	return this
}

func (this *webServer) Router() *httprouter.Router {
	return this.router
}

func (this *webServer) Start() {
	if this.waitExitFunc == nil {
		this.waitExitFunc = this.defaultWaitExit
	}
	if this.connStateFunc == nil {
		this.connStateFunc = this.defaultConnStateHook

		go this.manageIdleConns()
	}

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
		this.startServer(true)
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
		this.startServer(false)
	}

}

func (this *webServer) startServer(https bool) {
	var err error
	waitListenerUp := make(chan struct{})
	go func() {
		if Options.CpuAffinity {
			runtime.LockOSThread()
		}

		var (
			retryDelay         time.Duration
			theListener        net.Listener
			waitListenerUpOnce sync.Once
		)
		for {
			if https {
				this.httpsListener, err = net.Listen("tcp", this.httpsServer.Addr)
				if err != nil {
					if retryDelay == 0 {
						retryDelay = 50 * time.Millisecond
					} else {
						retryDelay = 2 * retryDelay
					}
					if maxDelay := time.Second; retryDelay > maxDelay {
						retryDelay = maxDelay
					}
					log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
					time.Sleep(retryDelay)
					continue
				}

				this.httpsListener, err = setupHttpsListener(this.httpsListener,
					this.gw.certFile, this.gw.keyFile)
				if err != nil {
					panic(err)
				}

				theListener = this.httpsListener
			} else {
				this.httpListener, err = net.Listen("tcp", this.httpServer.Addr)
				theListener = this.httpListener
			}

			if err != nil {
				if retryDelay == 0 {
					retryDelay = 50 * time.Millisecond
				} else {
					retryDelay = 2 * retryDelay
				}
				if maxDelay := time.Second; retryDelay > maxDelay {
					retryDelay = maxDelay
				}
				log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
				time.Sleep(retryDelay)
				continue
			}

			theListener = LimitListener(this.name, this.gw, theListener, this.maxClients)
			waitListenerUpOnce.Do(func() {
				close(waitListenerUp)
			})

			// on non-temporary err, net/http will close the listener
			if https {
				err = this.httpsServer.Serve(theListener)
			} else {
				err = this.httpServer.Serve(theListener)
			}

			select {
			case <-this.gw.shutdownCh:
				return

			default:
				log.Error("%s server: %v", this.name, err)
			}
		}
	}()

	// FIXME if net.Listen fails, kateway will not be able to stop
	// e,g. start a kateway, then start another, the 2nd will not be able to stop

	// wait for the listener up
	<-waitListenerUp

	this.mu.Lock()
	if !this.waiterStarted {
		this.waiterStarted = true

		this.gw.wg.Add(1)
		go this.waitExitFunc(this.gw.shutdownCh)
	}
	this.mu.Unlock()

	if https {
		log.Info("%s https server ready on %s", this.name, this.httpsServer.Addr)
	} else {
		log.Info("%s http server ready on %s", this.name, this.httpServer.Addr)
	}
}

func (this *webServer) defaultConnStateHook(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		atomic.AddInt32(&this.activeConnN, 1)

		if this.onConnNewFunc != nil {
			this.onConnNewFunc(c)
		}

	case http.StateIdle:
		this.stateIdleCh <- c

	case http.StateActive:
		this.stateActiveCh <- c

	case http.StateClosed, http.StateHijacked:
		atomic.AddInt32(&this.activeConnN, -1)
		this.stateRemoveCh <- c

		if this.onConnCloseFunc != nil {
			this.onConnCloseFunc(c)
		}
	}
}

func (this *webServer) manageIdleConns() {
	var (
		idleConns     = make(map[net.Conn]struct{}, 200)
		c             net.Conn
		waitNextRound = make(chan struct{})
	)
	defer close(waitNextRound)

	for {
		select {
		case <-this.gw.shutdownCh:
			if len(idleConns) == 0 {
				// happy ending
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				if conn == nil {
					continue
				}

				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case <-waitNextRound:
			if len(idleConns) == 0 {
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				if conn == nil {
					continue
				}

				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case c = <-this.stateActiveCh:
			delete(idleConns, c)

		case c = <-this.stateIdleCh:
			idleConns[c] = struct{}{}

		case c = <-this.stateRemoveCh:
			delete(idleConns, c)
		}
	}
}

func (this *webServer) defaultWaitExit(exit <-chan struct{}) {
	<-exit

	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpsListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	// wait for all established http/https conns close
	waitStart := time.Now()
	var prompt sync.Once
	for {
		activeConnN := atomic.LoadInt32(&this.activeConnN)
		if activeConnN == 0 {
			// good luck, all connections finished
			break
		}

		prompt.Do(func() {
			log.Trace("%s waiting for %d clients shutdown...", this.name, activeConnN)
		})

		// timeout mechanism
		if time.Since(waitStart) > Options.SubTimeout+time.Second {
			log.Warn("%s still left %d conns after %s, forced to shutdown",
				this.name, activeConnN, Options.SubTimeout)
			break
		}

		time.Sleep(time.Millisecond * 50)
	}
	log.Trace("%s all connections finished", this.name)

	if this.httpsServer != nil {
		this.httpsServer.ConnState = nil
	}
	if this.httpServer != nil {
		this.httpServer.ConnState = nil
	}

	close(this.stateActiveCh)
	close(this.stateIdleCh)
	close(this.stateRemoveCh)

	if this.onStop != nil {
		this.onStop()
	}

	this.gw.wg.Done()
}

func (this *webServer) notFoundHandler(w http.ResponseWriter, r *http.Request) {
	log.Error("%s: not found %s", this.name, r.RequestURI)

	writeNotFound(w)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/gocql/gocql"
	"github.com/grafana/globalconf"
	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/api/response"
	"github.com/grafana/metrictank/cmd/mt-control-server/controlmodels"
	"github.com/grafana/metrictank/expr/tagquery"
	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/schema/msg"
	"github.com/grafana/metrictank/util"
	log "github.com/sirupsen/logrus"
	"gopkg.in/macaron.v1"
)

var metrictankUrl string
var maxDefsPerMsg int
var indexConcurrency int

func ConfigHandlers() {
	FlagSet := flag.NewFlagSet("handlers", flag.ExitOnError)

	FlagSet.StringVar(&metrictankUrl, "metrictank-url", "http://metrictank:6060", "URL for MT cluster")
	FlagSet.IntVar(&maxDefsPerMsg, "max-defs-per-msg", 1000, "Maximum defs per control message")
	FlagSet.IntVar(&indexConcurrency, "index-concurrency", 5, "Concurrent queries to run against cassandra (per request)")

	globalconf.Register("handlers", FlagSet, flag.ExitOnError)
}

func makeLogId() string {
	logId := "uninit"
	uuid, err := gocql.RandomUUID()
	if err == nil {
		logId = uuid.String()
	}
	return logId
}

func makeMessage(partition int32, op schema.Operation, cm *schema.ControlMsg) (*sarama.ProducerMessage, error) {
	cm.Op = op
	payload, err := msg.WriteIndexControlMsg(cm)
	if err != nil {
		return nil, err
	}

	return &sarama.ProducerMessage{
		Topic:     topic,
		Value:     sarama.ByteEncoder(payload),
		Partition: partition,
	}, nil
}

func appStatus(ctx *macaron.Context) {
	ctx.PlainText(200, []byte("OK"))
}

func tagsDelByQuery(ctx *macaron.Context, request controlmodels.IndexDelByQueryReq) {
	logId := makeLogId()

	log.Infof("Received tagsDelByQuery: logId = %s, req = %v", logId, request)

	// Translate to FindSeries
	findSeries := models.GraphiteTagFindSeries{
		Expr:   request.Expr,
		Limit:  request.Limit,
		To:     request.OlderThan,
		Meta:   true,
		Format: "full-json",
	}
	reqBody, err := json.Marshal(findSeries)
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}
	// send to query address
	resp, err := http.Post(metrictankUrl+"/tags/findSeries", "application/json", bytes.NewBuffer(reqBody))
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}

	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		// Read in body so that the connection can be reused
		io.Copy(ioutil.Discard, resp.Body)
		response.Write(ctx, response.WrapError(fmt.Errorf("Request to cluster failed with status = %d", 502)))
		return
	}

	findSeriesResult := models.GraphiteTagFindSeriesFullResp{}
	err = json.NewDecoder(resp.Body).Decode(&findSeriesResult)
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}

	log.Infof("job %s: Pulled %d series matching filters", logId, len(findSeriesResult.Series))

	retval := controlmodels.IndexDelByQueryResp{}
	retval.Count = len(findSeriesResult.Series)
	retval.Partitions = make(map[int]int)

	for _, s := range findSeriesResult.Series {
		retval.Partitions[int(s.Partition)]++
	}

	if request.Method != "dry-run" {
		op := schema.OpRemove
		if request.Method == "archive" {
			op = schema.OpArchive
		}

		var msgs []*sarama.ProducerMessage

		// Create message per partition
		partitionMsgs := make(map[int32]*schema.ControlMsg)
		for _, s := range findSeriesResult.Series {
			if _, ok := partitionMsgs[s.Partition]; !ok {
				partitionMsgs[s.Partition] = &schema.ControlMsg{}
			}
			cm := partitionMsgs[s.Partition]
			cm.Defs = append(cm.Defs, s)

			if len(cm.Defs) >= maxDefsPerMsg {
				msg, err := makeMessage(s.Partition, op, cm)
				if err == nil {
					msgs = append(msgs, msg)
				}
				partitionMsgs[s.Partition] = &schema.ControlMsg{}
			}
		}

		for partition, cm := range partitionMsgs {
			if len(cm.Defs) == 0 {
				continue
			}
			msg, err := makeMessage(partition, op, cm)
			if err == nil {
				msgs = append(msgs, msg)
			}
		}

		if len(msgs) > 0 {
			go func() {
				pre := time.Now()
				err = producer.producer.SendMessages(msgs)
				log.Infof("job %s: Sending %d batches took %v", logId, len(msgs), time.Since(pre))
			}()
		}
	}

	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}
	response.Write(ctx, response.NewJson(200, retval, ""))
}

func tagsRestore(ctx *macaron.Context, request controlmodels.IndexRestoreReq) {
	// For large archives, this could take quite a long time. Process asynchronously and
	// return response to know the job was kicked off. Provide an identifier to view logs.
	logId := makeLogId()

	log.Infof("Received tagsRestore: logId = %s, req = %v", logId, request)

	// Pre-process expressions
	filter, err := tagquery.NewQueryFromStrings(request.Expr, request.From, request.To)
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}
	filtersByTag := make(map[string][]tagquery.Expression)
	for _, expr := range filter.Expressions {
		filtersByTag[expr.GetKey()] = append(filtersByTag[expr.GetKey()], expr)
	}

	numPartitions := request.NumPartitions
	if numPartitions < 1 {
		numPartitions = producer.numPartitions()
	}

	scanPartition := func(partition int32) {
		q := fmt.Sprintf("SELECT id, orgid, name, interval, unit, mtype, tags, lastupdate FROM %s WHERE partition = ?", cass.Config.ArchiveTable)

		var id, name, unit, mtype string
		var orgId, interval int
		var lastupdate int64
		var tags []string

		var defs []schema.MetricDefinition

		maxLastUpdate := time.Now().Unix()

		session := cass.Session.CurrentSession()
		iter := session.Query(q, partition).Iter()

	ITER:
		for iter.Scan(&id, &orgId, &name, &interval, &unit, &mtype, &tags, &lastupdate) {

			// Quick filter on lastupdate
			if filter.From > 0 && lastupdate < filter.From {
				continue ITER
			}

			if filter.To > 0 && lastupdate > filter.To {
				continue ITER
			}

			// Tags to map
			tagMap := make(map[string]string)
			for _, tag := range tags {
				tagParts := strings.SplitN(tag, "=", 2)
				if len(tagParts) != 2 {
					continue
				}
				tagMap[tagParts[0]] = tagParts[1]
			}

			idTagLookup := func(_ schema.MKey, tag, value string) bool {
				return tagMap[tag] == value
			}

			mkey, err := schema.MKeyFromString(id)
			if err != nil {
				log.Errorf("job=%s: load() could not parse ID %q: %s -> skipping", logId, id, err)
				continue
			}

			for _, expr := range filter.Expressions {
				matchFunc := expr.GetMetricDefinitionFilter(idTagLookup)
				if matchFunc(mkey, name, tags) == tagquery.Fail {
					continue ITER
				}
			}

			if orgId < 0 {
				orgId = int(idx.OrgIdPublic)
			}

			def := schema.MetricDefinition{
				Id:         mkey,
				OrgId:      uint32(orgId),
				Partition:  partition,
				Name:       name,
				Interval:   interval,
				Unit:       unit,
				Mtype:      mtype,
				Tags:       tags,
				LastUpdate: util.MinInt64(maxLastUpdate, lastupdate+request.LastUpdateOffset),
			}

			defs = append(defs, def)

		}
		if err := iter.Close(); err != nil {
			log.Errorf("job=%s: could not close iterator for partition = %d: %s", logId, partition, err.Error())
		}

		log.Infof("job=%s: Restoring %d series from partition %d", logId, len(defs), partition)

		if len(defs) == 0 {
			return
		}

		// Chunk into batches
		var defBatches [][]schema.MetricDefinition
		for i := 0; i < len(defs); i += maxDefsPerMsg {
			end := i + maxDefsPerMsg

			if end > len(defs) {
				end = len(defs)
			}

			defBatches = append(defBatches, defs[i:end])
		}

		for _, batch := range defBatches {
			cm := &schema.ControlMsg{
				Defs: batch,
				Op:   schema.OpRestore,
			}

			msg, err := makeMessage(partition, schema.OpRestore, cm)
			if err != nil {
				log.Warnf("job=%s: Failed to encode ControlMsg from partition = %d, err = %s", logId, partition, err)
				continue
			}

			_, _, err = producer.producer.SendMessage(msg)
			if err != nil {
				log.Warnf("job=%s: Failed to send ControlMsg for partition = %d, err = %s", logId, partition, err)
				continue
			}
		}
	}

	go func() {
		partitions := make(chan int, indexConcurrency)
		var wg sync.WaitGroup
		for i := 0; i < indexConcurrency; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				for partition := range partitions {
					scanPartition(int32(partition))
				}
			}()
		}

		for p := 0; p < numPartitions; p++ {
			partitions <- p
		}

		wg.Wait()
	}()

	retval := controlmodels.IndexRestoreResp{LogId: logId}
	response.Write(ctx, response.NewJson(200, retval, ""))
}
<commit_msg>Log failures as errors<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/gocql/gocql"
	"github.com/grafana/globalconf"
	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/api/response"
	"github.com/grafana/metrictank/cmd/mt-control-server/controlmodels"
	"github.com/grafana/metrictank/expr/tagquery"
	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/schema/msg"
	"github.com/grafana/metrictank/util"
	log "github.com/sirupsen/logrus"
	"gopkg.in/macaron.v1"
)

var metrictankUrl string
var maxDefsPerMsg int
var indexConcurrency int

func ConfigHandlers() {
	FlagSet := flag.NewFlagSet("handlers", flag.ExitOnError)

	FlagSet.StringVar(&metrictankUrl, "metrictank-url", "http://metrictank:6060", "URL for MT cluster")
	FlagSet.IntVar(&maxDefsPerMsg, "max-defs-per-msg", 1000, "Maximum defs per control message")
	FlagSet.IntVar(&indexConcurrency, "index-concurrency", 5, "Concurrent queries to run against cassandra (per request)")

	globalconf.Register("handlers", FlagSet, flag.ExitOnError)
}

func makeLogId() string {
	logId := "uninit"
	uuid, err := gocql.RandomUUID()
	if err == nil {
		logId = uuid.String()
	}
	return logId
}

func makeMessage(partition int32, op schema.Operation, cm *schema.ControlMsg) (*sarama.ProducerMessage, error) {
	cm.Op = op
	payload, err := msg.WriteIndexControlMsg(cm)
	if err != nil {
		return nil, err
	}

	return &sarama.ProducerMessage{
		Topic:     topic,
		Value:     sarama.ByteEncoder(payload),
		Partition: partition,
	}, nil
}

func appStatus(ctx *macaron.Context) {
	ctx.PlainText(200, []byte("OK"))
}

func tagsDelByQuery(ctx *macaron.Context, request controlmodels.IndexDelByQueryReq) {
	logId := makeLogId()

	log.Infof("Received tagsDelByQuery: logId = %s, req = %v", logId, request)

	// Translate to FindSeries
	findSeries := models.GraphiteTagFindSeries{
		Expr:   request.Expr,
		Limit:  request.Limit,
		To:     request.OlderThan,
		Meta:   true,
		Format: "full-json",
	}
	reqBody, err := json.Marshal(findSeries)
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}
	// send to query address
	resp, err := http.Post(metrictankUrl+"/tags/findSeries", "application/json", bytes.NewBuffer(reqBody))
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}

	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		// Read in body so that the connection can be reused
		io.Copy(ioutil.Discard, resp.Body)
		response.Write(ctx, response.WrapError(fmt.Errorf("Request to cluster failed with status = %d", 502)))
		return
	}

	findSeriesResult := models.GraphiteTagFindSeriesFullResp{}
	err = json.NewDecoder(resp.Body).Decode(&findSeriesResult)
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}

	log.Infof("job %s: Pulled %d series matching filters", logId, len(findSeriesResult.Series))

	retval := controlmodels.IndexDelByQueryResp{}
	retval.Count = len(findSeriesResult.Series)
	retval.Partitions = make(map[int]int)

	for _, s := range findSeriesResult.Series {
		retval.Partitions[int(s.Partition)]++
	}

	if request.Method != "dry-run" {
		op := schema.OpRemove
		if request.Method == "archive" {
			op = schema.OpArchive
		}

		var msgs []*sarama.ProducerMessage

		// Create message per partition
		partitionMsgs := make(map[int32]*schema.ControlMsg)
		for _, s := range findSeriesResult.Series {
			if _, ok := partitionMsgs[s.Partition]; !ok {
				partitionMsgs[s.Partition] = &schema.ControlMsg{}
			}
			cm := partitionMsgs[s.Partition]
			cm.Defs = append(cm.Defs, s)

			if len(cm.Defs) >= maxDefsPerMsg {
				msg, err := makeMessage(s.Partition, op, cm)
				if err == nil {
					msgs = append(msgs, msg)
				}
				partitionMsgs[s.Partition] = &schema.ControlMsg{}
			}
		}

		for partition, cm := range partitionMsgs {
			if len(cm.Defs) == 0 {
				continue
			}
			msg, err := makeMessage(partition, op, cm)
			if err == nil {
				msgs = append(msgs, msg)
			}
		}

		if len(msgs) > 0 {
			go func() {
				pre := time.Now()
				err = producer.producer.SendMessages(msgs)
				log.Infof("job %s: Sending %d batches took %v", logId, len(msgs), time.Since(pre))
			}()
		}
	}

	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}
	response.Write(ctx, response.NewJson(200, retval, ""))
}

func tagsRestore(ctx *macaron.Context, request controlmodels.IndexRestoreReq) {
	// For large archives, this could take quite a long time. Process asynchronously and
	// return response to know the job was kicked off. Provide an identifier to view logs.
	logId := makeLogId()

	log.Infof("Received tagsRestore: logId = %s, req = %v", logId, request)

	// Pre-process expressions
	filter, err := tagquery.NewQueryFromStrings(request.Expr, request.From, request.To)
	if err != nil {
		response.Write(ctx, response.WrapError(err))
		return
	}
	filtersByTag := make(map[string][]tagquery.Expression)
	for _, expr := range filter.Expressions {
		filtersByTag[expr.GetKey()] = append(filtersByTag[expr.GetKey()], expr)
	}

	numPartitions := request.NumPartitions
	if numPartitions < 1 {
		numPartitions = producer.numPartitions()
	}

	scanPartition := func(partition int32) {
		q := fmt.Sprintf("SELECT id, orgid, name, interval, unit, mtype, tags, lastupdate FROM %s WHERE partition = ?", cass.Config.ArchiveTable)

		var id, name, unit, mtype string
		var orgId, interval int
		var lastupdate int64
		var tags []string

		var defs []schema.MetricDefinition

		maxLastUpdate := time.Now().Unix()

		session := cass.Session.CurrentSession()
		iter := session.Query(q, partition).Iter()

	ITER:
		for iter.Scan(&id, &orgId, &name, &interval, &unit, &mtype, &tags, &lastupdate) {

			// Quick filter on lastupdate
			if filter.From > 0 && lastupdate < filter.From {
				continue ITER
			}

			if filter.To > 0 && lastupdate > filter.To {
				continue ITER
			}

			// Tags to map
			tagMap := make(map[string]string)
			for _, tag := range tags {
				tagParts := strings.SplitN(tag, "=", 2)
				if len(tagParts) != 2 {
					continue
				}
				tagMap[tagParts[0]] = tagParts[1]
			}

			idTagLookup := func(_ schema.MKey, tag, value string) bool {
				return tagMap[tag] == value
			}

			mkey, err := schema.MKeyFromString(id)
			if err != nil {
				log.Errorf("job=%s: load() could not parse ID %q: %s -> skipping", logId, id, err)
				continue
			}

			for _, expr := range filter.Expressions {
				matchFunc := expr.GetMetricDefinitionFilter(idTagLookup)
				if matchFunc(mkey, name, tags) == tagquery.Fail {
					continue ITER
				}
			}

			if orgId < 0 {
				orgId = int(idx.OrgIdPublic)
			}

			def := schema.MetricDefinition{
				Id:         mkey,
				OrgId:      uint32(orgId),
				Partition:  partition,
				Name:       name,
				Interval:   interval,
				Unit:       unit,
				Mtype:      mtype,
				Tags:       tags,
				LastUpdate: util.MinInt64(maxLastUpdate, lastupdate+request.LastUpdateOffset),
			}

			defs = append(defs, def)

		}
		if err := iter.Close(); err != nil {
			log.Errorf("job=%s: could not close iterator for partition = %d: %s", logId, partition, err.Error())
		}

		log.Infof("job=%s: Restoring %d series from partition %d", logId, len(defs), partition)

		if len(defs) == 0 {
			return
		}

		// Chunk into batches
		var defBatches [][]schema.MetricDefinition
		for i := 0; i < len(defs); i += maxDefsPerMsg {
			end := i + maxDefsPerMsg

			if end > len(defs) {
				end = len(defs)
			}

			defBatches = append(defBatches, defs[i:end])
		}

		for _, batch := range defBatches {
			cm := &schema.ControlMsg{
				Defs: batch,
				Op:   schema.OpRestore,
			}

			msg, err := makeMessage(partition, schema.OpRestore, cm)
			if err != nil {
				log.Errorf("job=%s: Failed to encode ControlMsg from partition = %d, err = %s", logId, partition, err)
				continue
			}

			_, _, err = producer.producer.SendMessage(msg)
			if err != nil {
				log.Errorf("job=%s: Failed to send ControlMsg for partition = %d, err = %s", logId, partition, err)
				continue
			}
		}
	}

	go func() {
		partitions := make(chan int, indexConcurrency)
		var wg sync.WaitGroup
		for i := 0; i < indexConcurrency; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				for partition := range partitions {
					scanPartition(int32(partition))
				}
			}()
		}

		for p := 0; p < numPartitions; p++ {
			partitions <- p
		}

		wg.Wait()
	}()

	retval := controlmodels.IndexRestoreResp{LogId: logId}
	response.Write(ctx, response.NewJson(200, retval, ""))
}
<|endoftext|>


<commit_before>// Sequential WARC file reader library supporting record-at-time compression
// this requires go1.4 due to its usage of compress/gzip.Reader#Multistream
//
// Currently only works on record-at-time compressed .gz files
//
// Example:
//
//     code
//     code
//
// see also:
//     URL to implementation repo
package warc

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"errors"
	"io"
	"strconv"
	"strings"
)

var (
	ErrMalformedRecord = errors.New("malformed record")
	ErrNonWARCRecord   = errors.New("non-WARC/1.0 record")
)

type byteCounter struct {
	r      *bufio.Reader
	offset int64
}

func (bc *byteCounter) Read(p []byte) (n int, err error) {
	n, err = bc.r.Read(p)
	bc.offset += int64(n)
	return
}

func (bc *byteCounter) ReadByte() (c byte, err error) {
	c, err = bc.r.ReadByte()
	if err == nil {
		bc.offset++
	}

	return
}

type Reader struct {
	bc   *byteCounter
	zr   *gzip.Reader
	last int64
}

type NamedField struct {
	Name  string
	Value string
}

type NamedFields []NamedField

type Record struct {
	Fields NamedFields
	Block  []byte
}

func (r *Record) bytes() []byte {
	var b bytes.Buffer
	b.Write([]byte("WARC/1.0\r\n"))
	for _, field := range r.Fields {
		b.Write([]byte(field.Name + ": " + field.Value + "\r\n"))
	}

	b.Write([]byte("\r\n"))
	b.Write([]byte(r.Block))
	return b.Bytes()
}

func (f NamedFields) Value(name string) string {
	for _, el := range f {
		if strings.EqualFold(el.Name, name) {
			return el.Value
		}
	}

	return ""
}

func NewReader(reader io.Reader) *Reader {
	return &Reader{
		bc: &byteCounter{
			r: bufio.NewReader(reader),
		},
	}
}

func NewGZIPReader(reader io.Reader) (r *Reader, err error) {
	r = NewReader(reader)
	r.zr, err = gzip.NewReader(r.bc)
	if err != nil {
		return nil, err
	}

	r.zr.Multistream(false)
	return r, nil
}

func (r *Reader) gzipRecord() ([]byte, error) {
	var rec bytes.Buffer
	_, err := io.Copy(&rec, r.zr)
	if err != nil {
		return nil, err
	}

	if r.last == r.bc.offset {
		return nil, io.EOF
	}

	r.last = r.bc.offset
	r.zr.Reset(r.bc)
	r.zr.Multistream(false)
	return rec.Bytes(), nil
}

func (r *Reader) plainRecord() ([]byte, error) {
	// TODO: Implement
	return nil, errors.New("support for non record-at-time compressed WARC files not yet implemented")
}

func (r *Reader) record() ([]byte, error) {
	if r.zr != nil {
		return r.gzipRecord()
	} else {
		return r.plainRecord()
	}
}

// returns io.EOF when done
func (r *Reader) Next() (*Record, error) {
	var res Record

	rec, err := r.record()
	if err != nil {
		return nil, err
	}

	parts := bytes.SplitN(rec, []byte("\r\n\r\n"), 2)
	if len(parts) != 2 {
		return nil, ErrMalformedRecord
	}

	hdr, block := parts[0], parts[1]
	res.Block = block
	for ix, hdrline := range bytes.Split(hdr, []byte("\r\n")) {
		if ix == 0 {
			if string(hdrline) != "WARC/1.0" {
				return nil, ErrNonWARCRecord
			}

			continue
		}

		l := bytes.SplitN(hdrline, []byte(":"), 2)
		if len(l) != 2 {
			return nil, ErrMalformedRecord
		}

		res.Fields = append(res.Fields, NamedField{
			Name:  string(bytes.Trim(l[0], "\r\n\t ")),
			Value: string(bytes.Trim(l[1], "\r\n\t ")),
		})
	}

	lenStr := res.Fields.Value("content-length")
	if i, err := strconv.Atoi(lenStr); err == nil {
		res.Block = res.Block[:i]
	}

	return &res, nil
}

type Writer struct {
	w  io.Writer
	zw *gzip.Writer
}

func NewWriter(w io.Writer) *Writer {
	return &Writer{w: w, zw: gzip.NewWriter(w)}
}

// Write a record. No validation of mandatory WARC fields is performed.
// The written record will be an independent GZIP stream.
func (w *Writer) WriteRecord(r *Record) error {
	rec := r.bytes()
	if _, err := w.zw.Write(rec); err != nil {
		return err
	}

	err := w.zw.Close()
	w.zw.Reset(w.w)
	return err
}
<commit_msg>Added trailing CRLFs to Record<commit_after>// Sequential WARC file reader library supporting record-at-time compression
// this requires go1.4 due to its usage of compress/gzip.Reader#Multistream
//
// Currently only works on record-at-time compressed .gz files
//
// Example:
//
//     code
//     code
//
// see also:
//     URL to implementation repo
package warc

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"errors"
	"io"
	"strconv"
	"strings"
)

var (
	ErrMalformedRecord = errors.New("malformed record")
	ErrNonWARCRecord   = errors.New("non-WARC/1.0 record")
)

type byteCounter struct {
	r      *bufio.Reader
	offset int64
}

func (bc *byteCounter) Read(p []byte) (n int, err error) {
	n, err = bc.r.Read(p)
	bc.offset += int64(n)
	return
}

func (bc *byteCounter) ReadByte() (c byte, err error) {
	c, err = bc.r.ReadByte()
	if err == nil {
		bc.offset++
	}

	return
}

type Reader struct {
	bc   *byteCounter
	zr   *gzip.Reader
	last int64
}

type NamedField struct {
	Name  string
	Value string
}

type NamedFields []NamedField

type Record struct {
	Fields NamedFields
	Block  []byte
}

func (r *Record) bytes() []byte {
	var b bytes.Buffer
	b.Write([]byte("WARC/1.0\r\n"))
	for _, field := range r.Fields {
		b.Write([]byte(field.Name + ": " + field.Value + "\r\n"))
	}

	b.Write([]byte("\r\n"))
	b.Write([]byte(r.Block))
	b.Write([]byte("\r\n\r\n"))
	return b.Bytes()
}

func (f NamedFields) Value(name string) string {
	for _, el := range f {
		if strings.EqualFold(el.Name, name) {
			return el.Value
		}
	}

	return ""
}

func NewReader(reader io.Reader) *Reader {
	return &Reader{
		bc: &byteCounter{
			r: bufio.NewReader(reader),
		},
	}
}

func NewGZIPReader(reader io.Reader) (r *Reader, err error) {
	r = NewReader(reader)
	r.zr, err = gzip.NewReader(r.bc)
	if err != nil {
		return nil, err
	}

	r.zr.Multistream(false)
	return r, nil
}

func (r *Reader) gzipRecord() ([]byte, error) {
	var rec bytes.Buffer
	_, err := io.Copy(&rec, r.zr)
	if err != nil {
		return nil, err
	}

	if r.last == r.bc.offset {
		return nil, io.EOF
	}

	r.last = r.bc.offset
	r.zr.Reset(r.bc)
	r.zr.Multistream(false)
	return rec.Bytes(), nil
}

func (r *Reader) plainRecord() ([]byte, error) {
	// TODO: Implement
	return nil, errors.New("support for non record-at-time compressed WARC files not yet implemented")
}

func (r *Reader) record() ([]byte, error) {
	if r.zr != nil {
		return r.gzipRecord()
	} else {
		return r.plainRecord()
	}
}

// returns io.EOF when done
func (r *Reader) Next() (*Record, error) {
	var res Record

	rec, err := r.record()
	if err != nil {
		return nil, err
	}

	parts := bytes.SplitN(rec, []byte("\r\n\r\n"), 2)
	if len(parts) != 2 {
		return nil, ErrMalformedRecord
	}

	hdr, block := parts[0], parts[1]
	res.Block = block
	for ix, hdrline := range bytes.Split(hdr, []byte("\r\n")) {
		if ix == 0 {
			if string(hdrline) != "WARC/1.0" {
				return nil, ErrNonWARCRecord
			}

			continue
		}

		l := bytes.SplitN(hdrline, []byte(":"), 2)
		if len(l) != 2 {
			return nil, ErrMalformedRecord
		}

		res.Fields = append(res.Fields, NamedField{
			Name:  string(bytes.Trim(l[0], "\r\n\t ")),
			Value: string(bytes.Trim(l[1], "\r\n\t ")),
		})
	}

	lenStr := res.Fields.Value("content-length")
	if i, err := strconv.Atoi(lenStr); err == nil {
		res.Block = res.Block[:i]
	}

	return &res, nil
}

type Writer struct {
	w  io.Writer
	zw *gzip.Writer
}

func NewWriter(w io.Writer) *Writer {
	return &Writer{w: w, zw: gzip.NewWriter(w)}
}

// Write a record. No validation of mandatory WARC fields is performed.
// The written record will be an independent GZIP stream.
func (w *Writer) WriteRecord(r *Record) error {
	rec := r.bytes()
	if _, err := w.zw.Write(rec); err != nil {
		return err
	}

	err := w.zw.Close()
	w.zw.Reset(w.w)
	return err
}
<|endoftext|>


<commit_before>package jvm

import (
    "fmt"
    . "jvmgo/any"
    "jvmgo/cmdline"
    _ "jvmgo/native"
    "jvmgo/jvm/rtda"
    rtc "jvmgo/jvm/rtda/class"
)

func Startup(cmd *cmdline.Command) {
    classPath := cmd.Options().Classpath()
    classLoader := rtc.NewClassLoader(classPath)
    mainThread := createMainThread(classLoader, cmd.Class(), cmd.Args())

    // todo
    defer func() {
        if r := recover(); r != nil {
            for !mainThread.IsStackEmpty() {
                frame := mainThread.PopFrame()
                fmt.Printf("%v %v\n", frame.Method().Class(), frame.Method())
            }

            err, ok := r.(error)
            if !ok {
                err = fmt.Errorf("%v", r)
                panic(err.Error())
            } else {
                panic(err.Error())
            }
        }
    }()

    loop(mainThread)
}

func createMainThread(classLoader Any, className string, args []string) (*rtda.Thread) {
    mainThread := rtda.NewThread(128, nil)
    bootMethod := rtc.NewBootstrapMethod([]byte{0xff, 0xb1}, classLoader)
    bootFrame := mainThread.NewFrame(bootMethod)
    mainThread.PushFrame(bootFrame)

    stack := bootFrame.OperandStack()
    stack.Push(args)
    stack.Push(className)
    stack.Push(classLoader)
    //stack.PushInt(0)

    return mainThread
}
<commit_msg>reorder imports<commit_after>package jvm

import (
    "fmt"
    . "jvmgo/any"
    "jvmgo/cmdline"
    "jvmgo/jvm/rtda"
    rtc "jvmgo/jvm/rtda/class"
    _ "jvmgo/native"
)

func Startup(cmd *cmdline.Command) {
    classPath := cmd.Options().Classpath()
    classLoader := rtc.NewClassLoader(classPath)
    mainThread := createMainThread(classLoader, cmd.Class(), cmd.Args())

    // todo
    defer func() {
        if r := recover(); r != nil {
            for !mainThread.IsStackEmpty() {
                frame := mainThread.PopFrame()
                fmt.Printf("%v %v\n", frame.Method().Class(), frame.Method())
            }

            err, ok := r.(error)
            if !ok {
                err = fmt.Errorf("%v", r)
                panic(err.Error())
            } else {
                panic(err.Error())
            }
        }
    }()

    loop(mainThread)
}

func createMainThread(classLoader Any, className string, args []string) (*rtda.Thread) {
    mainThread := rtda.NewThread(128, nil)
    bootMethod := rtc.NewBootstrapMethod([]byte{0xff, 0xb1}, classLoader)
    bootFrame := mainThread.NewFrame(bootMethod)
    mainThread.PushFrame(bootFrame)

    stack := bootFrame.OperandStack()
    stack.Push(args)
    stack.Push(className)
    stack.Push(classLoader)
    //stack.PushInt(0)

    return mainThread
}
<|endoftext|>


<commit_before>package autonat

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

	host "github.com/libp2p/go-libp2p-host"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

type NATStatus int

const (
	NATStatusUnknown = iota
	NATStatusPublic
	NATStatusPrivate
)

type AutoNAT interface {
	Status() NATStatus
	PublicAddr() (ma.Multiaddr, error)
}

type AutoNATState struct {
	ctx    context.Context
	host   host.Host
	peers  map[peer.ID]struct{}
	status NATStatus
	addr   ma.Multiaddr
	mx     sync.Mutex
}

func NewAutoNAT(ctx context.Context, h host.Host) AutoNAT {
	as := &AutoNATState{
		ctx:    ctx,
		host:   h,
		peers:  make(map[peer.ID]struct{}),
		status: NATStatusUnknown,
	}

	h.Network().Notify(as)
	go as.background()

	return as
}

func (as *AutoNATState) Status() NATStatus {
	return as.status
}

func (as *AutoNATState) PublicAddr() (ma.Multiaddr, error) {
	as.mx.Lock()
	defer as.mx.Unlock()

	if as.status != NATStatusPublic {
		return nil, errors.New("NAT Status is not public")
	}

	return as.addr, nil
}

func (as *AutoNATState) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	time.Sleep(10 * time.Second)
	for {
		as.autodetect()
		select {
		case <-time.After(15 * time.Minute):
		case <-as.ctx.Done():
			return
		}
	}
}

func (as *AutoNATState) autodetect() {
	if len(as.peers) == 0 {
		log.Debugf("skipping NAT auto detection; no autonat peers")
		return
	}

	as.mx.Lock()
	peers := make([]peer.ID, 0, len(as.peers))
	for p := range as.peers {
		if len(as.host.Network().ConnsToPeer(p)) > 0 {
			peers = append(peers, p)
		}
	}

	if len(peers) == 0 {
		// we don't have any open connections, try any autonat peer that we know about
		for p := range as.peers {
			peers = append(peers, p)
		}
	}

	as.mx.Unlock()

	shufflePeers(peers)

	for _, p := range peers {
		cli := NewAutoNATClient(as.host, p)
		ctx, cancel := context.WithTimeout(as.ctx, 60*time.Second)
		a, err := cli.Dial(ctx)
		cancel()

		switch {
		case err == nil:
			log.Debugf("NAT status is public; address through %s: %s", p.Pretty(), a.String())
			as.mx.Lock()
			as.addr = a
			as.status = NATStatusPublic
			as.mx.Unlock()
			return

		case IsDialError(err):
			log.Debugf("NAT status is private; dial error through %s: %s", p.Pretty(), err.Error())
			as.mx.Lock()
			as.status = NATStatusPrivate
			as.mx.Unlock()
			return

		default:
			log.Debugf("Error dialing through %s: %s", p.Pretty(), err.Error())
		}
	}

	as.mx.Lock()
	as.status = NATStatusUnknown
	as.mx.Unlock()
}

func shufflePeers(peers []peer.ID) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}
<commit_msg>typed NATStatus constants<commit_after>package autonat

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

	host "github.com/libp2p/go-libp2p-host"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

type NATStatus int

const (
	NATStatusUnknown NATStatus = iota
	NATStatusPublic
	NATStatusPrivate
)

type AutoNAT interface {
	Status() NATStatus
	PublicAddr() (ma.Multiaddr, error)
}

type AutoNATState struct {
	ctx    context.Context
	host   host.Host
	peers  map[peer.ID]struct{}
	status NATStatus
	addr   ma.Multiaddr
	mx     sync.Mutex
}

func NewAutoNAT(ctx context.Context, h host.Host) AutoNAT {
	as := &AutoNATState{
		ctx:    ctx,
		host:   h,
		peers:  make(map[peer.ID]struct{}),
		status: NATStatusUnknown,
	}

	h.Network().Notify(as)
	go as.background()

	return as
}

func (as *AutoNATState) Status() NATStatus {
	return as.status
}

func (as *AutoNATState) PublicAddr() (ma.Multiaddr, error) {
	as.mx.Lock()
	defer as.mx.Unlock()

	if as.status != NATStatusPublic {
		return nil, errors.New("NAT Status is not public")
	}

	return as.addr, nil
}

func (as *AutoNATState) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	time.Sleep(10 * time.Second)
	for {
		as.autodetect()
		select {
		case <-time.After(15 * time.Minute):
		case <-as.ctx.Done():
			return
		}
	}
}

func (as *AutoNATState) autodetect() {
	if len(as.peers) == 0 {
		log.Debugf("skipping NAT auto detection; no autonat peers")
		return
	}

	as.mx.Lock()
	peers := make([]peer.ID, 0, len(as.peers))
	for p := range as.peers {
		if len(as.host.Network().ConnsToPeer(p)) > 0 {
			peers = append(peers, p)
		}
	}

	if len(peers) == 0 {
		// we don't have any open connections, try any autonat peer that we know about
		for p := range as.peers {
			peers = append(peers, p)
		}
	}

	as.mx.Unlock()

	shufflePeers(peers)

	for _, p := range peers {
		cli := NewAutoNATClient(as.host, p)
		ctx, cancel := context.WithTimeout(as.ctx, 60*time.Second)
		a, err := cli.Dial(ctx)
		cancel()

		switch {
		case err == nil:
			log.Debugf("NAT status is public; address through %s: %s", p.Pretty(), a.String())
			as.mx.Lock()
			as.addr = a
			as.status = NATStatusPublic
			as.mx.Unlock()
			return

		case IsDialError(err):
			log.Debugf("NAT status is private; dial error through %s: %s", p.Pretty(), err.Error())
			as.mx.Lock()
			as.status = NATStatusPrivate
			as.mx.Unlock()
			return

		default:
			log.Debugf("Error dialing through %s: %s", p.Pretty(), err.Error())
		}
	}

	as.mx.Lock()
	as.status = NATStatusUnknown
	as.mx.Unlock()
}

func shufflePeers(peers []peer.ID) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}
<|endoftext|>


<commit_before>package relay

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

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus

	cachedAddrs       []ma.Multiaddr
	cachedAddrsExpiry time.Time
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			ar.mx.Lock()
			ar.cachedAddrs = nil
			ar.mx.Unlock()
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
	if ar.numRelays() >= DesiredRelays {
		return false
	}

	update := false
	for retry := 0; retry < 5; retry++ {
		if retry > 0 {
			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
			case <-ctx.Done():
				return update
			}
		}

		update = ar.findRelaysOnce(ctx) || update
		if ar.numRelays() > 0 {
			return update
		}
	}
	return update
}

func (ar *AutoRelay) findRelaysOnce(ctx context.Context) bool {
	pis, err := ar.discoverRelays(ctx)
	if err != nil {
		log.Debugf("error discovering relays: %s", err)
		return false
	}
	log.Debugf("discovered %d relays", len(pis))
	pis = ar.selectRelays(ctx, pis)
	log.Debugf("selected %d relays", len(pis))

	update := false
	for _, pi := range pis {
		update = ar.tryRelay(ctx, pi) || update
		if ar.numRelays() >= DesiredRelays {
			break
		}
	}
	return update
}

func (ar *AutoRelay) numRelays() int {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	return len(ar.relays)
}

// usingRelay returns if we're currently using the given relay.
func (ar *AutoRelay) usingRelay(p peer.ID) bool {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	_, ok := ar.relays[p]
	return ok
}

// addRelay adds the given relay to our set of relays.
// returns true when we add a new relay
func (ar *AutoRelay) tryRelay(ctx context.Context, pi pstore.PeerInfo) bool {
	if ar.usingRelay(pi.ID) {
		return false
	}

	if !ar.connect(ctx, pi) {
		return false
	}

	ar.mx.Lock()
	defer ar.mx.Unlock()

	// make sure we're still connected.
	if ar.host.Network().Connectedness(pi.ID) != inet.Connected {
		return false
	}
	ar.relays[pi.ID] = struct{}{}

	return true
}

func (ar *AutoRelay) connect(ctx context.Context, pi pstore.PeerInfo) bool {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	if len(pi.Addrs) == 0 {
		var err error
		pi, err = ar.router.FindPeer(ctx, pi.ID)
		if err != nil {
			log.Debugf("error finding relay peer %s: %s", pi.ID, err.Error())
			return false
		}
	}

	err := ar.host.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
		return false
	}

	// tag the connection as very important
	ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)
	return true
}

func (ar *AutoRelay) discoverRelays(ctx context.Context) ([]pstore.PeerInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	return discovery.FindPeers(ctx, ar.discover, RelayRendezvous, 1000)
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	shuffleRelays(pis)
	return pis
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.status != autonat.NATStatusPrivate {
		return addrs
	}

	if ar.cachedAddrs != nil && time.Now().Before(ar.cachedAddrsExpiry) {
		return ar.cachedAddrs
	}

	relays := make([]peer.ID, 0, len(ar.relays))
	for p := range ar.relays {
		relays = append(relays, p)
	}

	raddrs := make([]ma.Multiaddr, 0, 4*len(relays)+2)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for _, p := range relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	ar.cachedAddrs = raddrs
	ar.cachedAddrsExpiry = time.Now().Add(30 * time.Second)

	return raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<commit_msg>avoid intermediate allocation in relayAddrs<commit_after>package relay

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

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus

	cachedAddrs       []ma.Multiaddr
	cachedAddrsExpiry time.Time
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			ar.mx.Lock()
			ar.cachedAddrs = nil
			ar.mx.Unlock()
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
	if ar.numRelays() >= DesiredRelays {
		return false
	}

	update := false
	for retry := 0; retry < 5; retry++ {
		if retry > 0 {
			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
			case <-ctx.Done():
				return update
			}
		}

		update = ar.findRelaysOnce(ctx) || update
		if ar.numRelays() > 0 {
			return update
		}
	}
	return update
}

func (ar *AutoRelay) findRelaysOnce(ctx context.Context) bool {
	pis, err := ar.discoverRelays(ctx)
	if err != nil {
		log.Debugf("error discovering relays: %s", err)
		return false
	}
	log.Debugf("discovered %d relays", len(pis))
	pis = ar.selectRelays(ctx, pis)
	log.Debugf("selected %d relays", len(pis))

	update := false
	for _, pi := range pis {
		update = ar.tryRelay(ctx, pi) || update
		if ar.numRelays() >= DesiredRelays {
			break
		}
	}
	return update
}

func (ar *AutoRelay) numRelays() int {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	return len(ar.relays)
}

// usingRelay returns if we're currently using the given relay.
func (ar *AutoRelay) usingRelay(p peer.ID) bool {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	_, ok := ar.relays[p]
	return ok
}

// addRelay adds the given relay to our set of relays.
// returns true when we add a new relay
func (ar *AutoRelay) tryRelay(ctx context.Context, pi pstore.PeerInfo) bool {
	if ar.usingRelay(pi.ID) {
		return false
	}

	if !ar.connect(ctx, pi) {
		return false
	}

	ar.mx.Lock()
	defer ar.mx.Unlock()

	// make sure we're still connected.
	if ar.host.Network().Connectedness(pi.ID) != inet.Connected {
		return false
	}
	ar.relays[pi.ID] = struct{}{}

	return true
}

func (ar *AutoRelay) connect(ctx context.Context, pi pstore.PeerInfo) bool {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	if len(pi.Addrs) == 0 {
		var err error
		pi, err = ar.router.FindPeer(ctx, pi.ID)
		if err != nil {
			log.Debugf("error finding relay peer %s: %s", pi.ID, err.Error())
			return false
		}
	}

	err := ar.host.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
		return false
	}

	// tag the connection as very important
	ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)
	return true
}

func (ar *AutoRelay) discoverRelays(ctx context.Context) ([]pstore.PeerInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	return discovery.FindPeers(ctx, ar.discover, RelayRendezvous, 1000)
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	shuffleRelays(pis)
	return pis
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.status != autonat.NATStatusPrivate {
		return addrs
	}

	if ar.cachedAddrs != nil && time.Now().Before(ar.cachedAddrsExpiry) {
		return ar.cachedAddrs
	}

	raddrs := make([]ma.Multiaddr, 0, 4*len(ar.relays)+4)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for p := range ar.relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	ar.cachedAddrs = raddrs
	ar.cachedAddrsExpiry = time.Now().Add(30 * time.Second)

	return raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package zip

import (
	"bytes"
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
)

// TODO(adg): a more sophisticated test suite

type WriteTest struct {
	Name   string
	Data   []byte
	Method uint16
	Mode   os.FileMode
}

var writeTests = []WriteTest{
	{
		Name:   "foo",
		Data:   []byte("Rabbits, guinea pigs, gophers, marsupial rats, and quolls."),
		Method: Store,
		Mode:   0666,
	},
	{
		Name:   "bar",
		Data:   nil, // large data set in the test
		Method: Deflate,
		Mode:   0644,
	},
	{
		Name:   "setuid",
		Data:   []byte("setuid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetuid,
	},
	{
		Name:   "setgid",
		Data:   []byte("setgid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetgid,
	},
	{
		Name:   "symlink",
		Data:   []byte("../link/target"),
		Method: Deflate,
		Mode:   0755 | os.ModeSymlink,
	},
}

func TestWriter(t *testing.T) {
	largeData := make([]byte, 1<<17)
	for i := range largeData {
		largeData[i] = byte(rand.Int())
	}
	writeTests[1].Data = largeData
	defer func() {
		writeTests[1].Data = nil
	}()

	// write a zip file
	buf := new(bytes.Buffer)
	w := NewWriter(buf)

	for _, wt := range writeTests {
		testCreate(t, w, &wt)
	}

	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	// read it back
	r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
	if err != nil {
		t.Fatal(err)
	}
	for i, wt := range writeTests {
		testReadFile(t, r.File[i], &wt)
	}
}

func testCreate(t *testing.T, w *Writer, wt *WriteTest) {
	header := &FileHeader{
		Name:   wt.Name,
		Method: wt.Method,
	}
	if wt.Mode != 0 {
		header.SetMode(wt.Mode)
	}
	f, err := w.CreateHeader(header)
	if err != nil {
		t.Fatal(err)
	}
	_, err = f.Write(wt.Data)
	if err != nil {
		t.Fatal(err)
	}
}

func testReadFile(t *testing.T, f *File, wt *WriteTest) {
	if f.Name != wt.Name {
		t.Fatalf("File name: got %q, want %q", f.Name, wt.Name)
	}
	testFileMode(t, wt.Name, f, wt.Mode)
	rc, err := f.Open()
	if err != nil {
		t.Fatal("opening:", err)
	}
	b, err := ioutil.ReadAll(rc)
	if err != nil {
		t.Fatal("reading:", err)
	}
	err = rc.Close()
	if err != nil {
		t.Fatal("closing:", err)
	}
	if !bytes.Equal(b, wt.Data) {
		t.Errorf("File contents %q, want %q", b, wt.Data)
	}
}
<commit_msg>archive/zip: add flate writing benchmark<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package zip

import (
	"bytes"
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
)

// TODO(adg): a more sophisticated test suite

type WriteTest struct {
	Name   string
	Data   []byte
	Method uint16
	Mode   os.FileMode
}

var writeTests = []WriteTest{
	{
		Name:   "foo",
		Data:   []byte("Rabbits, guinea pigs, gophers, marsupial rats, and quolls."),
		Method: Store,
		Mode:   0666,
	},
	{
		Name:   "bar",
		Data:   nil, // large data set in the test
		Method: Deflate,
		Mode:   0644,
	},
	{
		Name:   "setuid",
		Data:   []byte("setuid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetuid,
	},
	{
		Name:   "setgid",
		Data:   []byte("setgid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetgid,
	},
	{
		Name:   "symlink",
		Data:   []byte("../link/target"),
		Method: Deflate,
		Mode:   0755 | os.ModeSymlink,
	},
}

func TestWriter(t *testing.T) {
	largeData := make([]byte, 1<<17)
	for i := range largeData {
		largeData[i] = byte(rand.Int())
	}
	writeTests[1].Data = largeData
	defer func() {
		writeTests[1].Data = nil
	}()

	// write a zip file
	buf := new(bytes.Buffer)
	w := NewWriter(buf)

	for _, wt := range writeTests {
		testCreate(t, w, &wt)
	}

	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	// read it back
	r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
	if err != nil {
		t.Fatal(err)
	}
	for i, wt := range writeTests {
		testReadFile(t, r.File[i], &wt)
	}
}

func testCreate(t *testing.T, w *Writer, wt *WriteTest) {
	header := &FileHeader{
		Name:   wt.Name,
		Method: wt.Method,
	}
	if wt.Mode != 0 {
		header.SetMode(wt.Mode)
	}
	f, err := w.CreateHeader(header)
	if err != nil {
		t.Fatal(err)
	}
	_, err = f.Write(wt.Data)
	if err != nil {
		t.Fatal(err)
	}
}

func testReadFile(t *testing.T, f *File, wt *WriteTest) {
	if f.Name != wt.Name {
		t.Fatalf("File name: got %q, want %q", f.Name, wt.Name)
	}
	testFileMode(t, wt.Name, f, wt.Mode)
	rc, err := f.Open()
	if err != nil {
		t.Fatal("opening:", err)
	}
	b, err := ioutil.ReadAll(rc)
	if err != nil {
		t.Fatal("reading:", err)
	}
	err = rc.Close()
	if err != nil {
		t.Fatal("closing:", err)
	}
	if !bytes.Equal(b, wt.Data) {
		t.Errorf("File contents %q, want %q", b, wt.Data)
	}
}

func BenchmarkCompressedZipGarbage(b *testing.B) {
	b.ReportAllocs()
	var buf bytes.Buffer
	bigBuf := bytes.Repeat([]byte("a"), 1<<20)
	for i := 0; i < b.N; i++ {
		buf.Reset()
		zw := NewWriter(&buf)
		for j := 0; j < 3; j++ {
			w, _ := zw.CreateHeader(&FileHeader{
				Name:   "foo",
				Method: Deflate,
			})
			w.Write(bigBuf)
		}
		zw.Close()
	}
}
<|endoftext|>