<commit_before>package turnpike

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

	"github.com/gorilla/websocket"
	logrus "github.com/sirupsen/logrus"
)

// errors.
var (
	ErrWSSendTimeout = errors.New("ws peer send timeout")
	ErrWSIsClosed    = errors.New("ws peer is closed")
)

type websocketPeer struct {
	conn        *websocket.Conn
	serializer  Serializer
	sendMsgs    chan Message
	messages    chan Message
	payloadType int
	mutex       sync.Mutex
	inSending   chan struct{}
	closing     chan struct{}
	*ConnectionConfig
	debugURL string
}

func NewWebsocketPeer(serialization Serialization, url string, tlscfg *tls.Config, dial DialFunc) (Peer, error) {
	switch serialization {
	case JSON:
		return newWebsocketPeer(url, jsonWebsocketProtocol,
			new(JSONSerializer), websocket.TextMessage, tlscfg, dial,
		)
	case MSGPACK:
		return newWebsocketPeer(url, msgpackWebsocketProtocol,
			new(MessagePackSerializer), websocket.BinaryMessage, tlscfg, dial,
		)
	default:
		return nil, fmt.Errorf("Unsupported serialization: %v", serialization)
	}
}

func newWebsocketPeer(url, protocol string, serializer Serializer, payloadType int, tlscfg *tls.Config, dial DialFunc) (Peer, error) {
	dialer := websocket.Dialer{
		Subprotocols:    []string{protocol},
		TLSClientConfig: tlscfg,
		Proxy:           http.ProxyFromEnvironment,
		NetDial:         dial,
	}
	conn, _, err := dialer.Dial(url, nil)
	if err != nil {
		return nil, err
	}
	ep := &websocketPeer{
		conn:             conn,
		sendMsgs:         make(chan Message, 16),
		messages:         make(chan Message, 100),
		serializer:       serializer,
		payloadType:      payloadType,
		closing:          make(chan struct{}),
		ConnectionConfig: &ConnectionConfig{},
		debugURL:         url,
	}
	go ep.run()

	return ep, nil
}

func (ep *websocketPeer) Send(msg Message) error {
	select {
	case ep.sendMsgs <- msg:
		return nil
	case <-time.After(5 * time.Second):
		log.Println(ErrWSSendTimeout.Error())
		ep.Close()
		err := fmt.Errorf("ws peer send timeout (%s)", ep.debugURL)
		log.Println(err.Error())
		return err // ErrWSSendTimeout
	case <-ep.closing:
		err := fmt.Errorf("ws peer is closed (%s)", ep.debugURL)
		log.Println(err.Error())
		return err // ErrWSIsClosed
	}
}

func (ep *websocketPeer) Receive() <-chan Message {
	return ep.messages
}

func (ep *websocketPeer) doClosing() {
	ep.mutex.Lock()
	defer ep.mutex.Unlock()

	select {
	case <-ep.closing:
	default:
		close(ep.closing)
	}
}

func (ep *websocketPeer) isClosed() bool {
	select {
	case <-ep.closing:
		return true
	default:
		return false
	}
}

func (ep *websocketPeer) Close() error {
	if ep.isClosed() {
		return nil
	}
	ep.doClosing()

	if ep.inSending != nil {
		<-ep.inSending
	}

	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "goodbye")
	err := ep.conn.WriteControl(websocket.CloseMessage, closeMsg, time.Now().Add(5*time.Second))
	if err != nil {
		log.Warningf("Error sending close message: %s", err.Error())
	}

	return ep.conn.Close()
}

func (ep *websocketPeer) updateReadDeadline() {
	ep.mutex.Lock()
	defer ep.mutex.Unlock()
	if ep.IdleTimeout > 0 {
		ep.conn.SetReadDeadline(time.Now().Add(ep.IdleTimeout))
	}
}

func (ep *websocketPeer) setReadDead() {
	ep.mutex.Lock()
	defer ep.mutex.Unlock()
	ep.conn.SetReadDeadline(time.Now())
}

func (ep *websocketPeer) run() {
	go ep.sending()

	if ep.MaxMsgSize > 0 {
		ep.conn.SetReadLimit(ep.MaxMsgSize)
	}
	ep.conn.SetPongHandler(func(v string) error {
		log.Println("pong:", v)
		ep.updateReadDeadline()
		return nil
	})
	ep.conn.SetPingHandler(func(v string) error {
		log.Println("ping:", v)
		ep.updateReadDeadline()
		return nil
	})

	for {
		// TODO: use conn.NextMessage() and stream
		// TODO: do something different based on binary/text frames
		ep.updateReadDeadline()
		if msgType, b, err := ep.conn.ReadMessage(); err != nil {
			if ep.isClosed() {
				log.Debugf("peer connection closed")
			} else {
				host, _, _ := net.SplitHostPort(ep.conn.RemoteAddr().String())
				log.WithFields(logrus.Fields{
					"remote_addr": host,
				}).Warning("error reading from peer")
				if !websocket.IsCloseError(err) {
					ep.conn.Close()
				}
			}
			log.Println("Closing channel")
			close(ep.messages)
			break
		} else if msgType == websocket.CloseMessage {
			ep.conn.Close()
			close(ep.messages)
			break
		} else {
			log.Debugf("websocketPeer read msg: %s", string(b))
			msg, err := ep.serializer.Deserialize(b)
			if err != nil {
				log.Errorf("Error deserializing peer message: %s", err.Error())
				// TODO: handle error
			} else {
				ep.messages <- msg
			}
		}
	}
}

func (ep *websocketPeer) sending() {
	ep.inSending = make(chan struct{})
	var ticker *time.Ticker
	if ep.PingTimeout == 0 {
		ticker = time.NewTicker(7 * 24 * time.Hour)
	} else {
		ticker = time.NewTicker(ep.PingTimeout)
	}

	defer func() {
		ep.setReadDead()
		ticker.Stop()
		close(ep.inSending)
	}()

	for {
		select {
		case msg := <-ep.sendMsgs:
			if closed, _ := ep.doSend(msg); closed {
				return
			}
		case <-ticker.C:
			wt := ep.WriteTimeout
			if wt == 0 {
				wt = 10 * time.Second
			}
			if err := ep.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(wt)); err != nil {
				log.Println("error sending ping message:", err)
				return
			}
		case <-ep.closing:
			// sending remaining messages.
			for {
				select {
				case msg := <-ep.sendMsgs:
					if closed, _ := ep.doSend(msg); !closed {
						continue
					}
				default:
				}
				break
			}
			return
		}
		ep.updateReadDeadline()
	}
}

func (ep *websocketPeer) doSend(msg Message) (closed bool, err error) {
	b, err := ep.serializer.Serialize(msg)
	if err != nil {
		log.Printf("error serializing peer message: %s, %+v", err, msg)
		return true, err
	}
	if ep.WriteTimeout > 0 {
		ep.conn.SetWriteDeadline(time.Now().Add(ep.WriteTimeout))
	}
	if err = ep.conn.WriteMessage(ep.payloadType, b); err != nil {
		log.Println("error write message: ", err)
		return true, err
	}
	return
}
<commit_msg>Add URL to websocket peer error messages.<commit_after>package turnpike

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

	"github.com/gorilla/websocket"
	logrus "github.com/sirupsen/logrus"
)

// errors.
var (
	ErrWSSendTimeout = errors.New("ws peer send timeout")
	ErrWSIsClosed    = errors.New("ws peer is closed")
)

type websocketPeer struct {
	conn        *websocket.Conn
	serializer  Serializer
	sendMsgs    chan Message
	messages    chan Message
	payloadType int
	mutex       sync.Mutex
	inSending   chan struct{}
	closing     chan struct{}
	*ConnectionConfig
	debugURL string
}

func NewWebsocketPeer(serialization Serialization, url string, tlscfg *tls.Config, dial DialFunc) (Peer, error) {
	switch serialization {
	case JSON:
		return newWebsocketPeer(url, jsonWebsocketProtocol,
			new(JSONSerializer), websocket.TextMessage, tlscfg, dial,
		)
	case MSGPACK:
		return newWebsocketPeer(url, msgpackWebsocketProtocol,
			new(MessagePackSerializer), websocket.BinaryMessage, tlscfg, dial,
		)
	default:
		return nil, fmt.Errorf("Unsupported serialization: %v", serialization)
	}
}

func newWebsocketPeer(url, protocol string, serializer Serializer, payloadType int, tlscfg *tls.Config, dial DialFunc) (Peer, error) {
	dialer := websocket.Dialer{
		Subprotocols:    []string{protocol},
		TLSClientConfig: tlscfg,
		Proxy:           http.ProxyFromEnvironment,
		NetDial:         dial,
	}
	conn, _, err := dialer.Dial(url, nil)
	if err != nil {
		return nil, err
	}
	ep := &websocketPeer{
		conn:             conn,
		sendMsgs:         make(chan Message, 16),
		messages:         make(chan Message, 100),
		serializer:       serializer,
		payloadType:      payloadType,
		closing:          make(chan struct{}),
		ConnectionConfig: &ConnectionConfig{},
		debugURL:         url,
	}
	go ep.run()

	return ep, nil
}

func (ep *websocketPeer) Send(msg Message) error {
	select {
	case ep.sendMsgs <- msg:
		return nil
	case <-time.After(5 * time.Second):
		err := fmt.Errorf("ws peer send timeout (%s)", ep.debugURL)
		log.Println(err.Error())
		ep.Close()
		return err // ErrWSSendTimeout
	case <-ep.closing:
		err := fmt.Errorf("ws peer is closed (%s)", ep.debugURL)
		log.Println(err.Error())
		return err // ErrWSIsClosed
	}
}

func (ep *websocketPeer) Receive() <-chan Message {
	return ep.messages
}

func (ep *websocketPeer) doClosing() {
	ep.mutex.Lock()
	defer ep.mutex.Unlock()

	select {
	case <-ep.closing:
	default:
		close(ep.closing)
	}
}

func (ep *websocketPeer) isClosed() bool {
	select {
	case <-ep.closing:
		return true
	default:
		return false
	}
}

func (ep *websocketPeer) Close() error {
	if ep.isClosed() {
		return nil
	}
	ep.doClosing()

	if ep.inSending != nil {
		<-ep.inSending
	}

	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "goodbye")
	err := ep.conn.WriteControl(websocket.CloseMessage, closeMsg, time.Now().Add(5*time.Second))
	if err != nil {
		log.Warningf("Error sending close message: %s", err.Error())
	}

	return ep.conn.Close()
}

func (ep *websocketPeer) updateReadDeadline() {
	ep.mutex.Lock()
	defer ep.mutex.Unlock()
	if ep.IdleTimeout > 0 {
		ep.conn.SetReadDeadline(time.Now().Add(ep.IdleTimeout))
	}
}

func (ep *websocketPeer) setReadDead() {
	ep.mutex.Lock()
	defer ep.mutex.Unlock()
	ep.conn.SetReadDeadline(time.Now())
}

func (ep *websocketPeer) run() {
	go ep.sending()

	if ep.MaxMsgSize > 0 {
		ep.conn.SetReadLimit(ep.MaxMsgSize)
	}
	ep.conn.SetPongHandler(func(v string) error {
		log.Println("pong:", v)
		ep.updateReadDeadline()
		return nil
	})
	ep.conn.SetPingHandler(func(v string) error {
		log.Println("ping:", v)
		ep.updateReadDeadline()
		return nil
	})

	for {
		// TODO: use conn.NextMessage() and stream
		// TODO: do something different based on binary/text frames
		ep.updateReadDeadline()
		if msgType, b, err := ep.conn.ReadMessage(); err != nil {
			if ep.isClosed() {
				log.Debugf("peer connection closed")
			} else {
				host, _, _ := net.SplitHostPort(ep.conn.RemoteAddr().String())
				log.WithFields(logrus.Fields{
					"remote_addr": host,
				}).Warning("error reading from peer")
				if !websocket.IsCloseError(err) {
					ep.conn.Close()
				}
			}
			log.Println("Closing channel")
			close(ep.messages)
			break
		} else if msgType == websocket.CloseMessage {
			ep.conn.Close()
			close(ep.messages)
			break
		} else {
			log.Debugf("websocketPeer read msg: %s", string(b))
			msg, err := ep.serializer.Deserialize(b)
			if err != nil {
				log.Errorf("Error deserializing peer message: %s", err.Error())
				// TODO: handle error
			} else {
				ep.messages <- msg
			}
		}
	}
}

func (ep *websocketPeer) sending() {
	ep.inSending = make(chan struct{})
	var ticker *time.Ticker
	if ep.PingTimeout == 0 {
		ticker = time.NewTicker(7 * 24 * time.Hour)
	} else {
		ticker = time.NewTicker(ep.PingTimeout)
	}

	defer func() {
		ep.setReadDead()
		ticker.Stop()
		close(ep.inSending)
	}()

	for {
		select {
		case msg := <-ep.sendMsgs:
			if closed, _ := ep.doSend(msg); closed {
				return
			}
		case <-ticker.C:
			wt := ep.WriteTimeout
			if wt == 0 {
				wt = 10 * time.Second
			}
			if err := ep.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(wt)); err != nil {
				log.Println("error sending ping message:", err)
				return
			}
		case <-ep.closing:
			// sending remaining messages.
			for {
				select {
				case msg := <-ep.sendMsgs:
					if closed, _ := ep.doSend(msg); !closed {
						continue
					}
				default:
				}
				break
			}
			return
		}
		ep.updateReadDeadline()
	}
}

func (ep *websocketPeer) doSend(msg Message) (closed bool, err error) {
	b, err := ep.serializer.Serialize(msg)
	if err != nil {
		log.Printf("error serializing peer message: %s, %+v", err, msg)
		return true, err
	}
	if ep.WriteTimeout > 0 {
		ep.conn.SetWriteDeadline(time.Now().Add(ep.WriteTimeout))
	}
	if err = ep.conn.WriteMessage(ep.payloadType, b); err != nil {
		log.Println("error write message: ", err)
		return true, err
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2009 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 time

// Sleep pauses the current goroutine for at least the duration d.
// A negative or zero duration causes Sleep to return immediately.
func Sleep(d Duration)

// runtimeNano returns the current value of the runtime clock in nanoseconds.
func runtimeNano() int64

// Interface to timers implemented in package runtime.
// Must be in sync with ../runtime/runtime.h:/^struct.Timer$
type runtimeTimer struct {
	i      int
	when   int64
	period int64
	f      func(interface{}, uintptr) // NOTE: must not be closure
	arg    interface{}
	seq    uintptr
}

// when is a helper function for setting the 'when' field of a runtimeTimer.
// It returns what the time will be, in nanoseconds, Duration d in the future.
// If d is negative, it is ignored. If the returned value would be less than
// zero because of an overflow, MaxInt64 is returned.
func when(d Duration) int64 {
	if d <= 0 {
		return runtimeNano()
	}
	t := runtimeNano() + int64(d)
	if t < 0 {
		t = 1<<63 - 1 // math.MaxInt64
	}
	return t
}

func startTimer(*runtimeTimer)
func stopTimer(*runtimeTimer) bool

// The Timer type represents a single event.
// When the Timer expires, the current time will be sent on C,
// unless the Timer was created by AfterFunc.
// A Timer must be created with NewTimer or AfterFunc.
type Timer struct {
	C <-chan Time
	r runtimeTimer
}

// Stop prevents the Timer from firing.
// It returns true if the call stops the timer, false if the timer has already
// expired or been stopped.
// Stop does not close the channel, to prevent a read from the channel succeeding
// incorrectly.
//
// To prevent the timer firing after a call to Stop,
// check the return value and drain the channel. For example:
// 	if !t.Stop() {
// 		<-t.C
// 	}
// This cannot be done concurrent to other receives from the Timer's
// channel.
func (t *Timer) Stop() bool {
	if t.r.f == nil {
		panic("time: Stop called on uninitialized Timer")
	}
	return stopTimer(&t.r)
}

// NewTimer creates a new Timer that will send
// the current time on its channel after at least duration d.
func NewTimer(d Duration) *Timer {
	c := make(chan Time, 1)
	t := &Timer{
		C: c,
		r: runtimeTimer{
			when: when(d),
			f:    sendTime,
			arg:  c,
		},
	}
	startTimer(&t.r)
	return t
}

// Reset changes the timer to expire after duration d.
// It returns true if the timer had been active, false if the timer had
// expired or been stopped.
//
// To reuse an active timer, always call its Stop method first and—if it had
// expired—drain the value from its channel. For example:
// 	if !t.Stop() {
// 		<-t.C
// 	}
// 	t.Reset(d)
// This should not be done concurrent to other receives from the Timer's
// channel.
//
// Note that it is not possible to use Reset's return value correctly, as there
// is a race condition between draining the channel and the new timer expiring.
// Reset should always be used in concert with Stop, as described above.
// The return value exists to preserve compatibility with existing programs.
func (t *Timer) Reset(d Duration) bool {
	if t.r.f == nil {
		panic("time: Reset called on uninitialized Timer")
	}
	w := when(d)
	active := stopTimer(&t.r)
	t.r.when = w
	startTimer(&t.r)
	return active
}

func sendTime(c interface{}, seq uintptr) {
	// Non-blocking send of time on c.
	// Used in NewTimer, it cannot block anyway (buffer).
	// Used in NewTicker, dropping sends on the floor is
	// the desired behavior when the reader gets behind,
	// because the sends are periodic.
	select {
	case c.(chan Time) <- Now():
	default:
	}
}

// After waits for the duration to elapse and then sends the current time
// on the returned channel.
// It is equivalent to NewTimer(d).C.
// The underlying Timer is not recovered by the garbage collector
// until the timer fires. If efficiency is a concern, use NewTimer
// instead and call Timer.Stop if the timer is no longer needed.
func After(d Duration) <-chan Time {
	return NewTimer(d).C
}

// AfterFunc waits for the duration to elapse and then calls f
// in its own goroutine. It returns a Timer that can
// be used to cancel the call using its Stop method.
func AfterFunc(d Duration, f func()) *Timer {
	t := &Timer{
		r: runtimeTimer{
			when: when(d),
			f:    goFunc,
			arg:  f,
		},
	}
	startTimer(&t.r)
	return t
}

func goFunc(arg interface{}, seq uintptr) {
	go arg.(func())()
}
<commit_msg>time: update comment to reflect correct file path.<commit_after>// Copyright 2009 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 time

// Sleep pauses the current goroutine for at least the duration d.
// A negative or zero duration causes Sleep to return immediately.
func Sleep(d Duration)

// runtimeNano returns the current value of the runtime clock in nanoseconds.
func runtimeNano() int64

// Interface to timers implemented in package runtime.
// Must be in sync with ../runtime/time.go:/^type timer
type runtimeTimer struct {
	i      int
	when   int64
	period int64
	f      func(interface{}, uintptr) // NOTE: must not be closure
	arg    interface{}
	seq    uintptr
}

// when is a helper function for setting the 'when' field of a runtimeTimer.
// It returns what the time will be, in nanoseconds, Duration d in the future.
// If d is negative, it is ignored. If the returned value would be less than
// zero because of an overflow, MaxInt64 is returned.
func when(d Duration) int64 {
	if d <= 0 {
		return runtimeNano()
	}
	t := runtimeNano() + int64(d)
	if t < 0 {
		t = 1<<63 - 1 // math.MaxInt64
	}
	return t
}

func startTimer(*runtimeTimer)
func stopTimer(*runtimeTimer) bool

// The Timer type represents a single event.
// When the Timer expires, the current time will be sent on C,
// unless the Timer was created by AfterFunc.
// A Timer must be created with NewTimer or AfterFunc.
type Timer struct {
	C <-chan Time
	r runtimeTimer
}

// Stop prevents the Timer from firing.
// It returns true if the call stops the timer, false if the timer has already
// expired or been stopped.
// Stop does not close the channel, to prevent a read from the channel succeeding
// incorrectly.
//
// To prevent the timer firing after a call to Stop,
// check the return value and drain the channel. For example:
// 	if !t.Stop() {
// 		<-t.C
// 	}
// This cannot be done concurrent to other receives from the Timer's
// channel.
func (t *Timer) Stop() bool {
	if t.r.f == nil {
		panic("time: Stop called on uninitialized Timer")
	}
	return stopTimer(&t.r)
}

// NewTimer creates a new Timer that will send
// the current time on its channel after at least duration d.
func NewTimer(d Duration) *Timer {
	c := make(chan Time, 1)
	t := &Timer{
		C: c,
		r: runtimeTimer{
			when: when(d),
			f:    sendTime,
			arg:  c,
		},
	}
	startTimer(&t.r)
	return t
}

// Reset changes the timer to expire after duration d.
// It returns true if the timer had been active, false if the timer had
// expired or been stopped.
//
// To reuse an active timer, always call its Stop method first and—if it had
// expired—drain the value from its channel. For example:
// 	if !t.Stop() {
// 		<-t.C
// 	}
// 	t.Reset(d)
// This should not be done concurrent to other receives from the Timer's
// channel.
//
// Note that it is not possible to use Reset's return value correctly, as there
// is a race condition between draining the channel and the new timer expiring.
// Reset should always be used in concert with Stop, as described above.
// The return value exists to preserve compatibility with existing programs.
func (t *Timer) Reset(d Duration) bool {
	if t.r.f == nil {
		panic("time: Reset called on uninitialized Timer")
	}
	w := when(d)
	active := stopTimer(&t.r)
	t.r.when = w
	startTimer(&t.r)
	return active
}

func sendTime(c interface{}, seq uintptr) {
	// Non-blocking send of time on c.
	// Used in NewTimer, it cannot block anyway (buffer).
	// Used in NewTicker, dropping sends on the floor is
	// the desired behavior when the reader gets behind,
	// because the sends are periodic.
	select {
	case c.(chan Time) <- Now():
	default:
	}
}

// After waits for the duration to elapse and then sends the current time
// on the returned channel.
// It is equivalent to NewTimer(d).C.
// The underlying Timer is not recovered by the garbage collector
// until the timer fires. If efficiency is a concern, use NewTimer
// instead and call Timer.Stop if the timer is no longer needed.
func After(d Duration) <-chan Time {
	return NewTimer(d).C
}

// AfterFunc waits for the duration to elapse and then calls f
// in its own goroutine. It returns a Timer that can
// be used to cancel the call using its Stop method.
func AfterFunc(d Duration, f func()) *Timer {
	t := &Timer{
		r: runtimeTimer{
			when: when(d),
			f:    goFunc,
			arg:  f,
		},
	}
	startTimer(&t.r)
	return t
}

func goFunc(arg interface{}, seq uintptr) {
	go arg.(func())()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"gopkg.in/kothar/brotli-go.v0/enc"

	"github.com/dustin/go-humanize"

	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/sync"
	"github.com/itchio/wharf/tlc"
	"github.com/itchio/wharf/wire"
)

// TODO: make this customizable
var ignoredDirs = []string{
	".git",
	".cvs",
	".svn",
}

func filterDirs(fileInfo os.FileInfo) bool {
	name := fileInfo.Name()
	for _, dir := range ignoredDirs {
		if strings.HasPrefix(name, dir) {
			return false
		}
	}

	return true
}

func diff(target string, source string, recipe string, brotliQuality int) {
	if *appArgs.verbose {
		Logf("Computing TLC signature of %s", target)
	}

	targetContainer, err := tlc.Walk(target, filterDirs)
	must(err)

	sourceContainer, err := tlc.Walk(source, filterDirs)
	must(err)

	StartProgress()
	targetSignature, err := pwr.ComputeDiffSignature(targetContainer, target, Progress)
	EndProgress()
	must(err)

	// index + weak + strong
	sigBytes := len(targetSignature) * (4 + 16)
	if *appArgs.verbose {
		Logf("Target signature size: %s", humanize.Bytes(uint64(sigBytes)))
	}

	brotliParams := enc.NewBrotliParams()
	brotliParams.SetQuality(brotliQuality)

	rawRecipeWriter, err := os.Create(recipe + ".br")
	must(err)

	recipeWriter := enc.NewBrotliWriter(brotliParams, rawRecipeWriter)

	rawSignatureWriter, err := os.Create(recipe + ".sig.br")
	must(err)

	dctx := &pwr.DiffContext{
		SourceContainer: sourceContainer,
		SourcePath:      source,

		TargetContainer: targetContainer,
		TargetSignature: targetSignature,
	}

	signatureWriter := enc.NewBrotliWriter(brotliParams, rawSignatureWriter)
	swc := wire.NewWriteContext(signatureWriter)

	err = swc.WriteMessage(&pwr.SignatureHeader{})
	must(err)

	sourceSignatureWriter := func(bl sync.BlockHash) error {
		// swc.WriteMessage(&pwr.BlockHash{
		// 	WeakHash:   bl.WeakHash,
		// 	StrongHash: bl.StrongHash,
		// })
		return nil
	}

	StartProgress()
	err = dctx.WriteRecipe(recipeWriter, Progress, brotliParams, sourceSignatureWriter)
	must(err)
	EndProgress()

	err = recipeWriter.Close()
	must(err)

	err = signatureWriter.Close()
	must(err)

	if *diffArgs.verify {
		tmpDir, err := ioutil.TempDir(os.TempDir(), "megadiff")
		must(err)
		defer os.RemoveAll(tmpDir)

		Logf("Verifying recipe by rebuilding source in %s", tmpDir)
		apply(recipe, target, tmpDir)

		tmpInfo, err := tlc.Walk(tmpDir, filterDirs)
		must(err)
		fmt.Printf("tmpInfo: %+v", tmpInfo)
	}
}

func apply(recipe string, target string, output string) {
	recipeReader, err := os.Open(recipe)
	must(err)

	StartProgress()
	err = pwr.ApplyRecipe(recipeReader, target, output, Progress)
	EndProgress()
	must(err)

	if *appArgs.verbose {
		Logf("Rebuilt source in %s", output)
	}
}
<commit_msg>:panda_face:<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"gopkg.in/kothar/brotli-go.v0/dec"
	"gopkg.in/kothar/brotli-go.v0/enc"

	"github.com/dustin/go-humanize"

	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/sync"
	"github.com/itchio/wharf/tlc"
	"github.com/itchio/wharf/wire"
)

// TODO: make this customizable
var ignoredDirs = []string{
	".git",
	".cvs",
	".svn",
}

func filterDirs(fileInfo os.FileInfo) bool {
	name := fileInfo.Name()
	for _, dir := range ignoredDirs {
		if strings.HasPrefix(name, dir) {
			return false
		}
	}

	return true
}

func diff(target string, source string, recipe string, brotliQuality int) {
	if *appArgs.verbose {
		Logf("Computing TLC signature of %s", target)
	}

	targetContainer, err := tlc.Walk(target, filterDirs)
	must(err)

	sourceContainer, err := tlc.Walk(source, filterDirs)
	must(err)

	StartProgress()
	targetSignature, err := pwr.ComputeDiffSignature(targetContainer, target, Progress)
	EndProgress()
	must(err)

	// index + weak + strong
	sigBytes := len(targetSignature) * (4 + 16)
	if *appArgs.verbose {
		Logf("Target signature size: %s", humanize.Bytes(uint64(sigBytes)))
	}

	rawRecipeWriter, err := os.Create(recipe)
	must(err)

	// recipeWriter := enc.NewBrotliWriter(brotliParams, rawRecipeWriter)
	recipeWriter := rawRecipeWriter

	signatureWriter, err := os.Create(recipe + ".sig")
	must(err)

	dctx := &pwr.DiffContext{
		SourceContainer: sourceContainer,
		SourcePath:      source,

		TargetContainer: targetContainer,
		TargetSignature: targetSignature,
	}

	hswc := wire.NewWriteContext(signatureWriter)

	err = hswc.WriteMessage(&pwr.SignatureHeader{})
	must(err)

	brotliParams := enc.NewBrotliParams()
	brotliParams.SetQuality(1)

	signatureCompressedWriter := enc.NewBrotliWriter(brotliParams, signatureWriter)
	swc := wire.NewWriteContext(signatureCompressedWriter)

	sourceSignatureWriter := func(bl sync.BlockHash) error {
		swc.WriteMessage(&pwr.BlockHash{
			WeakHash:   bl.WeakHash,
			StrongHash: bl.StrongHash,
		})
		return nil
	}

	StartProgress()
	err = dctx.WriteRecipe(recipeWriter, Progress, sourceSignatureWriter)
	must(err)
	EndProgress()

	err = recipeWriter.Close()
	must(err)

	err = signatureWriter.Close()
	must(err)

	if *diffArgs.verify {
		tmpDir, err := ioutil.TempDir(os.TempDir(), "megadiff")
		must(err)
		defer os.RemoveAll(tmpDir)

		Logf("Verifying recipe by rebuilding source in %s", tmpDir)
		apply(recipe, target, tmpDir)

		tmpInfo, err := tlc.Walk(tmpDir, filterDirs)
		must(err)
		fmt.Printf("tmpInfo: %+v", tmpInfo)
	}
}

func apply(recipe string, target string, output string) {
	rawRecipeReader, err := os.Open(recipe)
	must(err)

	recipeReader := dec.NewBrotliReader(rawRecipeReader)
	must(err)
	// recipeReader := rawRecipeReader

	StartProgress()
	err = pwr.ApplyRecipe(recipeReader, target, output, Progress)
	EndProgress()
	must(err)

	if *appArgs.verbose {
		Logf("Rebuilt source in %s", output)
	}
}
<|endoftext|>


<commit_before>package natsgw

import "encoding/json"

func (srv *Server) application() {
	srv.subscribe("ari.applications.all", func(_ string, _ []byte, reply Reply) {
		ax, err := srv.upstream.Application.List()
		if err != nil {
			reply(nil, err)
			return
		}

		var apps []string
		for _, a := range ax {
			apps = append(apps, a.ID())
		}

		reply(apps, nil)
	})

	srv.subscribe("ari.applications.data.>", func(subj string, _ []byte, reply Reply) {
		name := subj[len("ari.applications.data."):]
		data, err := srv.upstream.Application.Data(name)
		reply(data, err)
	})

	srv.subscribe("ari.applications.subscribe.>", func(subj string, data []byte, reply Reply) {
		name := subj[len("ari.applications.subscribe."):]

		var eventSource string
		if err := json.Unmarshal(data, &eventSource); err != nil {
			reply(nil, &decodingError{subj, err})
			return
		}

		err := srv.upstream.Application.Subscribe(name, eventSource)
		reply(nil, err)
	})

	srv.subscribe("ari.applications.unsubscribe.>", func(subj string, data []byte, reply Reply) {
		name := subj[len("ari.applications.unsubscribe."):]

		var eventSource string
		if err := json.Unmarshal(data, &eventSource); err != nil {
			reply(nil, &decodingError{subj, err})
			return
		}

		err := srv.upstream.Application.Unsubscribe(name, eventSource)
		reply(nil, err)
	})
}
<commit_msg>natsgw - update application to use '*' wildcards'<commit_after>package natsgw

import "encoding/json"

func (srv *Server) application() {
	srv.subscribe("ari.applications.all", func(_ string, _ []byte, reply Reply) {
		ax, err := srv.upstream.Application.List()
		if err != nil {
			reply(nil, err)
			return
		}

		var apps []string
		for _, a := range ax {
			apps = append(apps, a.ID())
		}

		reply(apps, nil)
	})

	srv.subscribe("ari.applications.data.*", func(subj string, _ []byte, reply Reply) {
		name := subj[len("ari.applications.data."):]
		data, err := srv.upstream.Application.Data(name)
		reply(data, err)
	})

	srv.subscribe("ari.applications.subscribe.*", func(subj string, data []byte, reply Reply) {
		name := subj[len("ari.applications.subscribe."):]

		var eventSource string
		if err := json.Unmarshal(data, &eventSource); err != nil {
			reply(nil, &decodingError{subj, err})
			return
		}

		err := srv.upstream.Application.Subscribe(name, eventSource)
		reply(nil, err)
	})

	srv.subscribe("ari.applications.unsubscribe.*", func(subj string, data []byte, reply Reply) {
		name := subj[len("ari.applications.unsubscribe."):]

		var eventSource string
		if err := json.Unmarshal(data, &eventSource); err != nil {
			reply(nil, &decodingError{subj, err})
			return
		}

		err := srv.upstream.Application.Unsubscribe(name, eventSource)
		reply(nil, err)
	})
}
<|endoftext|>


<commit_before>package server

import (
	"io/ioutil"
	"os"

	. "gopkg.in/check.v1"
)

const (
	host      = "localhost"
	port      = 4001
	snapAfter = 1000
	dbfile    = "rqlite-test"
)

type SingleServerSuite struct{}

var _ = Suite(&SingleServerSuite{})

func (s *SingleServerSuite) Test_SingleServer(c *C) {
	dir, err := ioutil.TempDir("", "rqlite-test-")
	c.Assert(err, IsNil)
	defer os.RemoveAll(dir)

	server := NewServer(dir, dbfile, snapAfter, host, port)
	c.Assert(server, NotNil)
	go func() { server.ListenAndServe("") }()
}
<commit_msg>More unit-testing of the Server<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	. "gopkg.in/check.v1"
)

const (
	host      = "localhost"
	port      = 4001
	snapAfter = 1000
	dbfile    = "rqlite-test"
)

type SingleServerSuite struct{}

var _ = Suite(&SingleServerSuite{})

func getEndpoint(endpoint string) (res *http.Response, err error) {
	url := fmt.Sprintf("http://%s:%d%s", host, port, endpoint)
	return http.Get(url)
}

func isJsonBody(res *http.Response) bool {
	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return false
	}

	var o interface{}
	err = json.Unmarshal(b, &o)
	if err != nil {
		return false
	}

	return true
}

func (s *SingleServerSuite) Test_SingleServer(c *C) {
	dir, err := ioutil.TempDir("", "rqlite-test-")
	c.Assert(err, IsNil)
	defer os.RemoveAll(dir)

	server := NewServer(dir, dbfile, snapAfter, host, port)
	c.Assert(server, NotNil)
	go func() { server.ListenAndServe("") }()

	// Wait to ensure server is up. This is not ideal, and the server should
	// really use a channel to flag it is ready.
	time.Sleep(1 * time.Second)

	// Sanity-check admin API endpoints
	var res *http.Response
	res, err = getEndpoint("/statistics")
	c.Assert(err, IsNil)
	c.Assert(isJsonBody(res), Equals, true)

	res, err = getEndpoint("/diagnostics")
	c.Assert(err, IsNil)
	c.Assert(isJsonBody(res), Equals, true)

	res, err = getEndpoint("/raft")
	c.Assert(err, IsNil)
	c.Assert(isJsonBody(res), Equals, true)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Skaffold 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 integration

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/GoogleContainerTools/skaffold/integration/skaffold"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event/proto"
	"github.com/GoogleContainerTools/skaffold/testutil"
	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/ptypes/empty"
	"google.golang.org/grpc"
)

var (
	connectionRetries = 20
	stateRetries      = 5
	numLogEntries     = 5
	waitTime          = 1 * time.Second
)

func TestEventLogRPC(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	rpcAddr := randomPort()
	teardown := setupSkaffoldWithArgs(t, "--rpc-port", rpcAddr)
	defer teardown()

	// start a grpc client and make sure we can connect properly
	var conn *grpc.ClientConn
	var err error
	var client proto.SkaffoldServiceClient
	attempts := 0
	for {
		conn, err = grpc.Dial(fmt.Sprintf(":%s", rpcAddr), grpc.WithInsecure())
		if err != nil {
			t.Logf("unable to establish skaffold grpc connection: retrying...")
			attempts++
			if attempts == connectionRetries {
				t.Fatalf("error establishing skaffold grpc connection")
			}

			time.Sleep(waitTime)
			continue
		}
		defer conn.Close()

		client = proto.NewSkaffoldServiceClient(conn)
		break
	}

	ctx, ctxCancel := context.WithCancel(context.Background())
	defer ctxCancel()

	var stream proto.SkaffoldService_EventLogClient
	attempts = 0
	for {
		stream, err = client.EventLog(ctx)
		if err == nil {
			break
		}
		if attempts < connectionRetries {
			attempts++
			t.Logf("waiting for connection...")
			time.Sleep(3 * time.Second)
			continue
		}
		t.Fatalf("error retrieving event log: %v\n", err)
	}

	// read a preset number of entries from the event log
	var logEntries []*proto.LogEntry
	entriesReceived := 0
	for {
		entry, err := stream.Recv()
		if err != nil {
			t.Errorf("error receiving entry from stream: %s", err)
		}

		if entry != nil {
			logEntries = append(logEntries, entry)
			entriesReceived++
		}
		if entriesReceived == numLogEntries {
			break
		}
	}
	metaEntries, buildEntries, deployEntries := 0, 0, 0
	for _, entry := range logEntries {
		switch entry.Event.GetEventType().(type) {
		case *proto.Event_MetaEvent:
			metaEntries++
		case *proto.Event_BuildEvent:
			buildEntries++
		case *proto.Event_DeployEvent:
			deployEntries++
		default:
		}
	}
	// make sure we have exactly 1 meta entry, 2 deploy entries and 2 build entries
	testutil.CheckDeepEqual(t, 1, metaEntries)
	testutil.CheckDeepEqual(t, 2, deployEntries)
	testutil.CheckDeepEqual(t, 2, buildEntries)
}

func TestEventLogHTTP(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	httpAddr := randomPort()
	teardown := setupSkaffoldWithArgs(t, "--rpc-http-port", httpAddr)
	defer teardown()
	time.Sleep(500 * time.Millisecond) // give skaffold time to process all events

	httpResponse, err := http.Get(fmt.Sprintf("http://localhost:%s/v1/event_log", httpAddr))
	if err != nil {
		t.Fatalf("error connecting to gRPC REST API: %s", err.Error())
	}
	numEntries := 0
	var logEntries []proto.LogEntry
	for {
		e := make([]byte, 1024)
		l, err := httpResponse.Body.Read(e)
		if err != nil {
			t.Errorf("error reading body from http response: %s", err.Error())
		}
		e = e[0:l] // remove empty bytes from slice

		// sometimes reads can encompass multiple log entries, since Read() doesn't count newlines as EOF.
		readEntries := strings.Split(string(e), "\n")
		for _, entryStr := range readEntries {
			if entryStr == "" {
				continue
			}
			var entry proto.LogEntry
			// the HTTP wrapper sticks the proto messages into a map of "result" -> message.
			// attempting to JSON unmarshal drops necessary proto information, so we just manually
			// strip the string off the response and unmarshal directly to the proto message
			entryStr = strings.Replace(entryStr, "{\"result\":", "", 1)
			entryStr = entryStr[:len(entryStr)-1]
			if err := jsonpb.UnmarshalString(entryStr, &entry); err != nil {
				t.Errorf("error converting http response to proto: %s", err.Error())
			}
			numEntries++
			logEntries = append(logEntries, entry)
		}
		if numEntries >= numLogEntries {
			break
		}
	}

	metaEntries, buildEntries, deployEntries := 0, 0, 0
	for _, entry := range logEntries {
		switch entry.Event.GetEventType().(type) {
		case *proto.Event_MetaEvent:
			metaEntries++
		case *proto.Event_BuildEvent:
			buildEntries++
		case *proto.Event_DeployEvent:
			deployEntries++
		default:
		}
	}
	// make sure we have exactly 1 meta entry, 2 deploy entries and 2 build entries
	testutil.CheckDeepEqual(t, 1, metaEntries)
	testutil.CheckDeepEqual(t, 2, deployEntries)
	testutil.CheckDeepEqual(t, 2, buildEntries)
}

func TestGetStateRPC(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	rpcAddr := randomPort()
	// start a skaffold dev loop on an example
	teardown := setupSkaffoldWithArgs(t, "--rpc-port", rpcAddr)
	defer teardown()

	// start a grpc client and make sure we can connect properly
	var conn *grpc.ClientConn
	var err error
	var client proto.SkaffoldServiceClient
	attempts := 0
	for {
		conn, err = grpc.Dial(fmt.Sprintf(":%s", rpcAddr), grpc.WithInsecure())
		if err != nil {
			t.Logf("unable to establish skaffold grpc connection: retrying...")
			attempts++
			if attempts == connectionRetries {
				t.Fatalf("error establishing skaffold grpc connection")
			}

			time.Sleep(waitTime)
			continue
		}
		defer conn.Close()

		client = proto.NewSkaffoldServiceClient(conn)
		break
	}

	ctx, ctxCancel := context.WithCancel(context.Background())
	defer ctxCancel()

	// try a few times and wait around until we see the build is complete, or fail.
	attempts = 0
outerLoop:
	for {
		grpcState := retrieveRPCState(ctx, t, client)
		if attempts == stateRetries {
			// fail condition has been reached: abbreviated way of
			// only logging the artifacts that didn't reach completion.
			for _, v := range grpcState.BuildState.Artifacts {
				testutil.CheckDeepEqual(t, event.Complete, v)
			}
			testutil.CheckDeepEqual(t, event.Complete, grpcState.DeployState.Status)
			break
		}
		for _, v := range grpcState.BuildState.Artifacts {
			if v != event.Complete {
				time.Sleep(500 * time.Millisecond)
				attempts++
				continue outerLoop
			}
		}
		if grpcState.DeployState.Status != event.Complete {
			time.Sleep(500 * time.Millisecond)
			attempts++
			continue outerLoop
		}
		break
	}
}

func retrieveRPCState(ctx context.Context, t *testing.T, client proto.SkaffoldServiceClient) *proto.State {
	var grpcState *proto.State
	var err error
	attempts := 0
	for {
		grpcState, err = client.GetState(ctx, &empty.Empty{})
		if err == nil {
			break
		}
		if attempts < connectionRetries {
			attempts++
			t.Logf("waiting for connection...")
			time.Sleep(3 * time.Second)
			continue
		}
		t.Fatalf("error retrieving state: %v\n", err)
	}
	return grpcState
}

func retrieveHTTPState(t *testing.T, httpAddr string) proto.State {
	var httpState proto.State

	// retrieve the state via HTTP as well, and verify the result is the same
	httpResponse, err := http.Get(fmt.Sprintf("http://localhost:%s/v1/state", httpAddr))
	if err != nil {
		t.Fatalf("error connecting to gRPC REST API: %s", err.Error())
	}
	b, err := ioutil.ReadAll(httpResponse.Body)
	if err != nil {
		t.Errorf("error reading body from http response: %s", err.Error())
	}
	if err := json.Unmarshal(b, &httpState); err != nil {
		t.Errorf("error converting http response to proto: %s", err.Error())
	}
	return httpState
}

func TestGetStateHTTP(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	httpAddr := randomPort()
	teardown := setupSkaffoldWithArgs(t, "--rpc-http-port", httpAddr)
	defer teardown()
	time.Sleep(3 * time.Second) // give skaffold time to process all events

	attempts := 0
outerLoop:
	for {
		httpState := retrieveHTTPState(t, httpAddr)
		if attempts == stateRetries {
			for _, v := range httpState.BuildState.Artifacts {
				testutil.CheckDeepEqual(t, event.Complete, v)
			}
			testutil.CheckDeepEqual(t, event.Complete, httpState.DeployState.Status)
			break
		}
		for _, v := range httpState.BuildState.Artifacts {
			if v != event.Complete {
				time.Sleep(1 * time.Second)
				attempts++
				continue outerLoop
			}
		}
		if httpState.DeployState.Status != event.Complete {
			time.Sleep(1 * time.Second)
			attempts++
			continue outerLoop
		}
		break
	}
}

func setupSkaffoldWithArgs(t *testing.T, args ...string) func() {
	Run(t, "testdata/dev", "sh", "-c", "echo foo > foo")

	// Run skaffold build first to fail quickly on a build failure
	skaffold.Build().InDir("testdata/dev").RunOrFail(t)

	// start a skaffold dev loop on an example
	ns, _, deleteNs := SetupNamespace(t)

	stop := skaffold.Dev(args...).InDir("testdata/dev").InNs(ns.Name).RunBackground(t)

	return func() {
		stop()
		deleteNs()
		Run(t, "testdata/dev", "rm", "foo")
	}
}

func randomPort() string {
	return fmt.Sprintf("%d", rand.Intn(65535))
}
<commit_msg>refactor for clarity<commit_after>/*
Copyright 2019 The Skaffold 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 integration

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/GoogleContainerTools/skaffold/integration/skaffold"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event/proto"
	"github.com/GoogleContainerTools/skaffold/testutil"
	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/ptypes/empty"
	"google.golang.org/grpc"
)

var (
	connectionRetries = 20
	stateRetries      = 5
	numLogEntries     = 5
	waitTime          = 1 * time.Second
)

func TestEventLogRPC(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	rpcAddr := randomPort()
	teardown := setupSkaffoldWithArgs(t, "--rpc-port", rpcAddr)
	defer teardown()

	// start a grpc client and make sure we can connect properly
	var conn *grpc.ClientConn
	var err error
	var client proto.SkaffoldServiceClient
	attempts := 0
	for {
		conn, err = grpc.Dial(fmt.Sprintf(":%s", rpcAddr), grpc.WithInsecure())
		if err != nil {
			t.Logf("unable to establish skaffold grpc connection: retrying...")
			attempts++
			if attempts == connectionRetries {
				t.Fatalf("error establishing skaffold grpc connection")
			}

			time.Sleep(waitTime)
			continue
		}
		defer conn.Close()

		client = proto.NewSkaffoldServiceClient(conn)
		break
	}

	ctx, ctxCancel := context.WithCancel(context.Background())
	defer ctxCancel()

	var stream proto.SkaffoldService_EventLogClient
	attempts = 0
	for {
		stream, err = client.EventLog(ctx)
		if err == nil {
			break
		}
		if attempts < connectionRetries {
			attempts++
			t.Logf("waiting for connection...")
			time.Sleep(3 * time.Second)
			continue
		}
		t.Fatalf("error retrieving event log: %v\n", err)
	}

	// read a preset number of entries from the event log
	var logEntries []*proto.LogEntry
	entriesReceived := 0
	for {
		entry, err := stream.Recv()
		if err != nil {
			t.Errorf("error receiving entry from stream: %s", err)
		}

		if entry != nil {
			logEntries = append(logEntries, entry)
			entriesReceived++
		}
		if entriesReceived == numLogEntries {
			break
		}
	}
	metaEntries, buildEntries, deployEntries := 0, 0, 0
	for _, entry := range logEntries {
		switch entry.Event.GetEventType().(type) {
		case *proto.Event_MetaEvent:
			metaEntries++
		case *proto.Event_BuildEvent:
			buildEntries++
		case *proto.Event_DeployEvent:
			deployEntries++
		default:
		}
	}
	// make sure we have exactly 1 meta entry, 2 deploy entries and 2 build entries
	testutil.CheckDeepEqual(t, 1, metaEntries)
	testutil.CheckDeepEqual(t, 2, deployEntries)
	testutil.CheckDeepEqual(t, 2, buildEntries)
}

func TestEventLogHTTP(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	httpAddr := randomPort()
	teardown := setupSkaffoldWithArgs(t, "--rpc-http-port", httpAddr)
	defer teardown()
	time.Sleep(500 * time.Millisecond) // give skaffold time to process all events

	httpResponse, err := http.Get(fmt.Sprintf("http://localhost:%s/v1/event_log", httpAddr))
	if err != nil {
		t.Fatalf("error connecting to gRPC REST API: %s", err.Error())
	}
	numEntries := 0
	var logEntries []proto.LogEntry
	for {
		e := make([]byte, 1024)
		l, err := httpResponse.Body.Read(e)
		if err != nil {
			t.Errorf("error reading body from http response: %s", err.Error())
		}
		e = e[0:l] // remove empty bytes from slice

		// sometimes reads can encompass multiple log entries, since Read() doesn't count newlines as EOF.
		readEntries := strings.Split(string(e), "\n")
		for _, entryStr := range readEntries {
			if entryStr == "" {
				continue
			}
			var entry proto.LogEntry
			// the HTTP wrapper sticks the proto messages into a map of "result" -> message.
			// attempting to JSON unmarshal drops necessary proto information, so we just manually
			// strip the string off the response and unmarshal directly to the proto message
			entryStr = strings.Replace(entryStr, "{\"result\":", "", 1)
			entryStr = entryStr[:len(entryStr)-1]
			if err := jsonpb.UnmarshalString(entryStr, &entry); err != nil {
				t.Errorf("error converting http response to proto: %s", err.Error())
			}
			numEntries++
			logEntries = append(logEntries, entry)
		}
		if numEntries >= numLogEntries {
			break
		}
	}

	metaEntries, buildEntries, deployEntries := 0, 0, 0
	for _, entry := range logEntries {
		switch entry.Event.GetEventType().(type) {
		case *proto.Event_MetaEvent:
			metaEntries++
		case *proto.Event_BuildEvent:
			buildEntries++
		case *proto.Event_DeployEvent:
			deployEntries++
		default:
		}
	}
	// make sure we have exactly 1 meta entry, 2 deploy entries and 2 build entries
	testutil.CheckDeepEqual(t, 1, metaEntries)
	testutil.CheckDeepEqual(t, 2, deployEntries)
	testutil.CheckDeepEqual(t, 2, buildEntries)
}

func TestGetStateRPC(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	rpcAddr := randomPort()
	// start a skaffold dev loop on an example
	teardown := setupSkaffoldWithArgs(t, "--rpc-port", rpcAddr)
	defer teardown()

	// start a grpc client and make sure we can connect properly
	var conn *grpc.ClientConn
	var err error
	var client proto.SkaffoldServiceClient
	attempts := 0
	for {
		conn, err = grpc.Dial(fmt.Sprintf(":%s", rpcAddr), grpc.WithInsecure())
		if err != nil {
			t.Logf("unable to establish skaffold grpc connection: retrying...")
			attempts++
			if attempts == connectionRetries {
				t.Fatalf("error establishing skaffold grpc connection")
			}

			time.Sleep(waitTime)
			continue
		}
		defer conn.Close()

		client = proto.NewSkaffoldServiceClient(conn)
		break
	}

	ctx, ctxCancel := context.WithCancel(context.Background())
	defer ctxCancel()

	// try a few times and wait around until we see the build is complete, or fail.
	attempts = 0
	success := false
	var grpcState *proto.State
	for i := 0; i < stateRetries; i++ {
		grpcState = retrieveRPCState(ctx, t, client)
		if checkState(*grpcState) {
			success = true
			break
		}
		time.Sleep(1 * time.Second)
	}
	if !success {
		// max attempts exceeded, log errors
		for _, v := range grpcState.BuildState.Artifacts {
			testutil.CheckDeepEqual(t, event.Complete, v)
		}
		testutil.CheckDeepEqual(t, event.Complete, grpcState.DeployState.Status)
	}
}

func retrieveRPCState(ctx context.Context, t *testing.T, client proto.SkaffoldServiceClient) *proto.State {
	var grpcState *proto.State
	var err error
	attempts := 0
	for {
		grpcState, err = client.GetState(ctx, &empty.Empty{})
		if err == nil {
			break
		}
		if attempts < connectionRetries {
			attempts++
			t.Logf("waiting for connection...")
			time.Sleep(3 * time.Second)
			continue
		}
		t.Fatalf("error retrieving state: %v\n", err)
	}
	return grpcState
}

func retrieveHTTPState(t *testing.T, httpAddr string) proto.State {
	var httpState proto.State

	// retrieve the state via HTTP as well, and verify the result is the same
	httpResponse, err := http.Get(fmt.Sprintf("http://localhost:%s/v1/state", httpAddr))
	if err != nil {
		t.Fatalf("error connecting to gRPC REST API: %s", err.Error())
	}
	b, err := ioutil.ReadAll(httpResponse.Body)
	if err != nil {
		t.Errorf("error reading body from http response: %s", err.Error())
	}
	if err := json.Unmarshal(b, &httpState); err != nil {
		t.Errorf("error converting http response to proto: %s", err.Error())
	}
	return httpState
}

func TestGetStateHTTP(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	httpAddr := randomPort()
	teardown := setupSkaffoldWithArgs(t, "--rpc-http-port", httpAddr)
	defer teardown()
	time.Sleep(3 * time.Second) // give skaffold time to process all events

	success := false
	var httpState proto.State
	for i := 0; i < stateRetries; i++ {
		httpState = retrieveHTTPState(t, httpAddr)
		if checkState(httpState) {
			success = true
			break
		}
		time.Sleep(1 * time.Second)
	}
	if !success {
		// max attempts exceeded, log errors
		for _, v := range httpState.BuildState.Artifacts {
			testutil.CheckDeepEqual(t, event.Complete, v)
		}
		testutil.CheckDeepEqual(t, event.Complete, httpState.DeployState.Status)
	}
}

func setupSkaffoldWithArgs(t *testing.T, args ...string) func() {
	Run(t, "testdata/dev", "sh", "-c", "echo foo > foo")

	// Run skaffold build first to fail quickly on a build failure
	skaffold.Build().InDir("testdata/dev").RunOrFail(t)

	// start a skaffold dev loop on an example
	ns, _, deleteNs := SetupNamespace(t)

	stop := skaffold.Dev(args...).InDir("testdata/dev").InNs(ns.Name).RunBackground(t)

	return func() {
		stop()
		deleteNs()
		Run(t, "testdata/dev", "rm", "foo")
	}
}

func randomPort() string {
	return fmt.Sprintf("%d", rand.Intn(65535))
}

func checkState(state proto.State) bool {
	for _, a := range state.BuildState.Artifacts {
		if a != event.Complete {
			return false
		}
	}
	return state.DeployState.Status == event.Complete
}
<|endoftext|>


<commit_before>package agent

import (
	"context"
	"fmt"
	"time"

	"github.com/coreos/go-systemd/login1"
	"github.com/golang/glog"
	"k8s.io/client-go/1.5/kubernetes"
	v1core "k8s.io/client-go/1.5/kubernetes/typed/core/v1"
	"k8s.io/client-go/1.5/pkg/api"
	v1api "k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/fields"
	"k8s.io/client-go/1.5/pkg/watch"

	"github.com/coreos-inc/container-linux-update-operator/internal/constants"
	"github.com/coreos-inc/container-linux-update-operator/internal/drain"
	"github.com/coreos-inc/container-linux-update-operator/internal/k8sutil"
	"github.com/coreos-inc/container-linux-update-operator/internal/updateengine"
)

type Klocksmith struct {
	node string
	kc   *kubernetes.Clientset
	nc   v1core.NodeInterface
	ue   *updateengine.Client
	lc   *login1.Conn
}

func New(node string) (*Klocksmith, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// set up update_engine client
	ue, err := updateengine.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to update_engine dbus: %v", err)
	}

	// set up login1 client for our eventual reboot
	lc, err := login1.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to logind dbus: %v", err)
	}

	return &Klocksmith{node, kc, nc, ue, lc}, nil
}

// Run runs klocksmithd, reacting to the update_engine reboot signal by
// draining pods on this kubernetes node and rebooting.
//
// TODO: try to be more resilient against transient failures
func (k *Klocksmith) Run() error {
	// set annotations with helpful info about the coreos node
	vi, err := k8sutil.GetVersionInfo()
	if err != nil {
		return fmt.Errorf("failed to get version info: %v", err)
	}

	labels := map[string]string{
		constants.LabelID:      vi.ID,
		constants.LabelGroup:   vi.Group,
		constants.LabelVersion: vi.Version,
	}

	glog.Infof("Setting labels: %#v", labels)
	if err := k8sutil.SetNodeLabels(k.nc, k.node, labels); err != nil {
		return err
	}

	// set coreos.com/update1/reboot-in-progress=false and
	// coreos.com/update1/reboot-needed=false
	anno := map[string]string{
		constants.AnnotationRebootInProgress: constants.False,
		constants.AnnotationRebootNeeded:     constants.False,
	}
	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// we are schedulable now.
	glog.Info("Marking node as schedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, false); err != nil {
		return err
	}

	// block until update_engine says to reboot, and set label.
	glog.Infof("Waiting for reboot signal...")
	if err := k.waitForRebootSignal(); err != nil {
		return err
	}

	// indicate we need a reboot
	anno = map[string]string{
		constants.AnnotationRebootNeeded: constants.True,
	}
	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// block until constants.AnnotationOkToReboot is set
	glog.Infof("Waiting for ok-to-reboot from controller...")
	if err := k.waitForOkToReboot(); err != nil {
		return err
	}

	// set constants.AnnotationRebootInProgress and drain self
	anno = map[string]string{
		constants.AnnotationRebootInProgress: constants.True,
	}
	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// drain self equates to:
	// 1. set Unschedulable
	// 2. delete all pods
	// unlike `kubectl drain`, we do not care about emptyDir or orphan pods
	// ('any pods that are neither mirror pods nor managed by
	// ReplicationController, ReplicaSet, DaemonSet or Job')

	glog.Info("Marking node as unschedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, true); err != nil {
		return err
	}

	glog.Info("Getting pod list for deletion")
	pods, err := k.getPodsForDeletion()
	if err != nil {
		return err
	}

	// delete the pods.
	// TODO: explicitly don't terminate self? we'll probably just be a
	// mirror pod or daemonset anyway..
	glog.Infof("Deleting %d pods", len(pods))
	deleteOptions := api.NewDeleteOptions(30)
	for _, pod := range pods {
		glog.Infof("Terminating pod %q...", pod.Name)
		if err := k.kc.Pods(pod.Namespace).Delete(pod.Name, deleteOptions); err != nil {
			return fmt.Errorf("failed terminating pod %q: %v", pod.Name, err)
		}
	}

	glog.Info("Node drained, rebooting")

	// reboot
	k.lc.Reboot(false)

	// cross fingers
	time.Sleep(24 * 7 * time.Hour)

	return nil
}

// waitForRebootSignal watches update_engine and waits for
// updateengine.UpdateStatusUpdatedNeedReboot
func (k *Klocksmith) waitForRebootSignal() error {
	// first, check current status.
	status, err := k.ue.GetStatus()
	if err != nil {
		return fmt.Errorf("failed to get update_engine status: %v", err)
	}

	// if we're not in UpdateStatusUpdatedNeedReboot, then wait for it.
	if status.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		ctx, cancel := context.WithCancel(context.Background())
		stc := make(chan updateengine.Status)
		go k.ue.RebootNeededSignal(stc, ctx.Done())
		<-stc
		cancel()
	}

	// we are now in UpdateStatusUpdatedNeedReboot state.
	return nil
}

func (k *Klocksmith) waitForOkToReboot() error {
	n, err := k.nc.Get(k.node)
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// hopefully 24 hours is enough time between indicating we need a
	// reboot and the controller telling us to do it
	ev, err := watch.Until(time.Hour*24, watcher, k8sutil.NodeAnnotationCondition(constants.AnnotationOkToReboot, constants.True))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no, ok := ev.Object.(*v1api.Node)
	if !ok {
		panic("event contains a non-*api.Node object")
	}

	if no.Annotations[constants.AnnotationOkToReboot] != constants.True {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) getPodsForDeletion() ([]v1api.Pod, error) {
	pods, err := drain.GetPodsForDeletion(k.kc, k.node)
	if err != nil {
		return nil, fmt.Errorf("failed to get list of pods for deletion: %v", err)
	}

	// XXX: ignoring kube-system is a simple way to avoid eviciting
	// critical components such as kube-scheduler and
	// kube-controller-manager.

	pods = k8sutil.FilterPods(pods, func(p *v1api.Pod) bool {
		if p.Namespace == "kube-system" {
			return false
		}

		return true
	})

	return pods, nil
}
<commit_msg>internal/agent: async watch of update_engine status<commit_after>package agent

import (
	"context"
	"fmt"
	"time"

	"github.com/coreos/go-systemd/login1"
	"github.com/golang/glog"
	"k8s.io/client-go/1.5/kubernetes"
	v1core "k8s.io/client-go/1.5/kubernetes/typed/core/v1"
	"k8s.io/client-go/1.5/pkg/api"
	v1api "k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/fields"
	"k8s.io/client-go/1.5/pkg/util/wait"
	"k8s.io/client-go/1.5/pkg/watch"

	"github.com/coreos-inc/container-linux-update-operator/internal/constants"
	"github.com/coreos-inc/container-linux-update-operator/internal/drain"
	"github.com/coreos-inc/container-linux-update-operator/internal/k8sutil"
	"github.com/coreos-inc/container-linux-update-operator/internal/updateengine"
)

type Klocksmith struct {
	node string
	kc   *kubernetes.Clientset
	nc   v1core.NodeInterface
	ue   *updateengine.Client
	lc   *login1.Conn
}

func New(node string) (*Klocksmith, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// set up update_engine client
	ue, err := updateengine.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to update_engine dbus: %v", err)
	}

	// set up login1 client for our eventual reboot
	lc, err := login1.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to logind dbus: %v", err)
	}

	return &Klocksmith{node, kc, nc, ue, lc}, nil
}

// updateStatusCallback receives Status messages from update engine. If the
// status is UpdateStatusUpdatedNeedReboot, indicate that with a label on our
// node.
func (k *Klocksmith) updateStatusCallback(s updateengine.Status) {
	glog.Info("Updating status")
	// update our status
	anno := map[string]string{
		constants.AnnotationStatus: s.CurrentOperation,
	}

	// indicate we need a reboot
	if s.CurrentOperation == updateengine.UpdateStatusUpdatedNeedReboot {
		glog.Info("Indicating a reboot is needed")
		anno[constants.AnnotationRebootNeeded] = constants.True
	}

	wait.PollUntil(10*time.Second, func() (bool, error) {
		if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
			glog.Errorf("Failed to set annotation %q: %v", constants.AnnotationStatus, err)
			return false, nil
		}

		return true, nil
	}, wait.NeverStop)
}

// setInfoLabels labels our node with helpful info about Container Linux.
func (k *Klocksmith) setInfoLabels() error {
	vi, err := k8sutil.GetVersionInfo()
	if err != nil {
		return fmt.Errorf("failed to get version info: %v", err)
	}

	labels := map[string]string{
		constants.LabelID:      vi.ID,
		constants.LabelGroup:   vi.Group,
		constants.LabelVersion: vi.Version,
	}

	if err := k8sutil.SetNodeLabels(k.nc, k.node, labels); err != nil {
		return err
	}

	return nil
}

// Run runs klocksmithd, reacting to the update_engine reboot signal by
// draining pods on this kubernetes node and rebooting.
//
// TODO: try to be more resilient against transient failures
func (k *Klocksmith) Run() error {
	log.Print("Setting info labels")
	if err := k.setInfoLabels(); err != nil {
		return fmt.Errorf("failed to set node info: %v", err)
	}

	// set coreos.com/update1/reboot-in-progress=false and
	// coreos.com/update1/reboot-needed=false
	anno := map[string]string{
		constants.AnnotationRebootInProgress: constants.False,
		constants.AnnotationRebootNeeded:     constants.False,
	}
	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// we are schedulable now.
	glog.Info("Marking node as schedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, false); err != nil {
		return err
	}

	// watch update engine for status updates
	go k.watchUpdateStatus(k.updateStatusCallback)

	// block until constants.AnnotationOkToReboot is set
	glog.Infof("Waiting for ok-to-reboot from controller...")
	if err := k.waitForOkToReboot(); err != nil {
		return err
	}

	// set constants.AnnotationRebootInProgress and drain self
	anno = map[string]string{
		constants.AnnotationRebootInProgress: constants.True,
	}

	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// drain self equates to:
	// 1. set Unschedulable
	// 2. delete all pods
	// unlike `kubectl drain`, we do not care about emptyDir or orphan pods
	// ('any pods that are neither mirror pods nor managed by
	// ReplicationController, ReplicaSet, DaemonSet or Job')

	glog.Info("Marking node as unschedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, true); err != nil {
		return err
	}

	glog.Info("Getting pod list for deletion")
	pods, err := k.getPodsForDeletion()
	if err != nil {
		return err
	}

	// delete the pods.
	// TODO: explicitly don't terminate self? we'll probably just be a
	// mirror pod or daemonset anyway..
	glog.Infof("Deleting %d pods", len(pods))
	deleteOptions := api.NewDeleteOptions(30)
	for _, pod := range pods {
		glog.Infof("Terminating pod %q...", pod.Name)
		if err := k.kc.Pods(pod.Namespace).Delete(pod.Name, deleteOptions); err != nil {
			return fmt.Errorf("failed terminating pod %q: %v", pod.Name, err)
		}
	}

	glog.Info("Node drained, rebooting")

	// reboot
	k.lc.Reboot(false)

	// cross fingers
	time.Sleep(24 * 7 * time.Hour)

	return nil
}

func (k *Klocksmith) watchUpdateStatus(update func(s updateengine.Status)) {
	glog.Info("Beginning to watch update_engine status")

	oldOperation := ""
	ch := make(chan updateengine.Status, 1)
	stop := make(chan struct{})

	go k.ue.ReceiveStatuses(ch, stop)

	// synthesize first message
	st, err := k.ue.GetStatus()
	if err != nil {
		glog.Fatalf("Failed to get update_engine status: %v", err)
	}

	ch <- st

	for status := range ch {
		if status.CurrentOperation != oldOperation && update != nil {
			update(status)
			oldOperation = status.CurrentOperation
		}
	}
}

// waitForRebootSignal watches update_engine and waits for
// updateengine.UpdateStatusUpdatedNeedReboot
func (k *Klocksmith) waitForRebootSignal() error {
	// first, check current status.
	status, err := k.ue.GetStatus()
	if err != nil {
		return fmt.Errorf("failed to get update_engine status: %v", err)
	}

	// if we're not in UpdateStatusUpdatedNeedReboot, then wait for it.
	if status.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		ctx, cancel := context.WithCancel(context.Background())
		stc := make(chan updateengine.Status)
		go k.ue.RebootNeededSignal(stc, ctx.Done())
		<-stc
		cancel()
	}

	// we are now in UpdateStatusUpdatedNeedReboot state.
	return nil
}

func (k *Klocksmith) waitForOkToReboot() error {
	n, err := k.nc.Get(k.node)
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// hopefully 24 hours is enough time between indicating we need a
	// reboot and the controller telling us to do it
	ev, err := watch.Until(time.Hour*24, watcher, k8sutil.NodeAnnotationCondition(constants.AnnotationOkToReboot, constants.True))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no, ok := ev.Object.(*v1api.Node)
	if !ok {
		panic("event contains a non-*api.Node object")
	}

	if no.Annotations[constants.AnnotationOkToReboot] != constants.True {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) getPodsForDeletion() ([]v1api.Pod, error) {
	pods, err := drain.GetPodsForDeletion(k.kc, k.node)
	if err != nil {
		return nil, fmt.Errorf("failed to get list of pods for deletion: %v", err)
	}

	// XXX: ignoring kube-system is a simple way to avoid eviciting
	// critical components such as kube-scheduler and
	// kube-controller-manager.

	pods = k8sutil.FilterPods(pods, func(p *v1api.Pod) bool {
		if p.Namespace == "kube-system" {
			return false
		}

		return true
	})

	return pods, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Vermeer Light Tools 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 core

import (
	"errors"
	m "github.com/jamiec7919/vermeer/math"
	"reflect"
)

var ErrParamNotFloat = errors.New("Parameter is not float or convertible to float")
var ErrParamNotInt = errors.New("Parameter is not int or convertible to int")
var ErrParamNotString = errors.New("Parameter is not string")
var ErrNotStruct = errors.New("Unmarshal: param is not a struct")

type Params map[string]interface{}

func (np *Params) GetString(key string) (string, error) {
	if v, present := (*np)[key]; present {
		return getString(v)
	}

	return "", nil
}

func (np *Params) GetFloat(key string) (float64, error) {
	if v, present := (*np)[key]; present {
		return getFloat(v)
	}

	return 0.0, nil
}

func (np *Params) GetInt(key string) (int64, error) {
	if v, present := (*np)[key]; present {
		return getInt(v)
	}

	return 0, nil
}

func (np *Params) GetVec3(key string) (m.Vec3, error) {
	if v, present := (*np)[key]; present {
		return getVec3(v)
	}

	return m.Vec3{}, nil
}

func getString(v interface{}) (string, error) {
	switch t := v.(type) {
	case string:
		return t, nil
	default:
		return "", ErrParamNotString
	}

}

func getFloat(v interface{}) (float64, error) {
	switch t := v.(type) {
	case float64:
		return t, nil
	case int64:
		return float64(t), nil
	default:
		return 0, ErrParamNotFloat
	}

}

func getInt(v interface{}) (int64, error) {
	switch t := v.(type) {
	case float64:
		return int64(t), nil
	case int64:
		return t, nil
	default:
		return 0, ErrParamNotInt
	}

}

func getVec3(v interface{}) (m.Vec3, error) {
	switch t := v.(type) {
	case []float64:
		if len(t) == 3 {
			return m.Vec3{float32(t[0]), float32(t[1]), float32(t[2])}, nil
		} else {
			return m.Vec3{}, nil
		}
	case []int64:
		if len(t) == 3 {
			return m.Vec3{float32(t[0]), float32(t[1]), float32(t[2])}, nil
		} else {
			return m.Vec3{}, nil
		}
	default:
		return m.Vec3{}, nil
	}
}

// This is a utility function for simple structs representing nodes
func (np *Params) Unmarshal(v interface{}) error {
	// Check v is struct

	rv := reflect.ValueOf(v)

	relem := rv.Elem()

	if relem.Kind() != reflect.Struct {
		return ErrNotStruct

	}

	for k, param := range *np {
		// lookup k in v and assign as appropriate
		field := relem.FieldByName(k)
		// log.Printf("%v", k)
		if field.IsValid() {
			if field.CanSet() {
				switch field.Kind() {
				case reflect.Int:
					v, _ := getInt(param)
					field.SetInt(v)
				case reflect.Float64:
					fallthrough
				case reflect.Float32:
					v, _ := getFloat(param)
					field.SetFloat(v)
				case reflect.String:
					field.SetString(param.(string))
				case reflect.Slice:
					// el := field.Type().Elem()

					arr := param.([]Params)

					v := reflect.MakeSlice(field.Type(), len(arr), len(arr))

					for i := range arr {
						arr[i].Unmarshal(v.Index(i).Addr().Interface())
						// log.Printf("%v %v", i, arr[i])
					}
					field.Set(v)
					// log.Printf("Filling array of %v %v", el, v)
				default:
					switch field.Type() {
					case reflect.TypeOf(m.Vec3{}):
						v, _ := getVec3(param)
						field.Set(reflect.ValueOf(v))
						//				case reflect.TypeOf(Colour{}):
						//					v, _ := getColour(param)
						//					field.Set(reflect.ValueOf(v))
						//				case reflect.TypeOf((*Texture)(nil)):
						//					v, _ := getTexture(param)
						//					field.Set(reflect.ValueOf(v))
					}
				}
			}
		}
	}

	return nil

}
<commit_msg>Updated Params.Unmarshal to iterate over struct fields instead of Params map<commit_after>// Copyright 2016 The Vermeer Light Tools 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 core

import (
	"errors"
	m "github.com/jamiec7919/vermeer/math"
	"reflect"
)

var ErrParamNotFloat = errors.New("Parameter is not float or convertible to float")
var ErrParamNotInt = errors.New("Parameter is not int or convertible to int")
var ErrParamNotBool = errors.New("Parameter is not bool or convertible to bool")
var ErrParamNotString = errors.New("Parameter is not string")
var ErrNotStruct = errors.New("Unmarshal: param is not a struct")
var ErrNotFound = errors.New("Parameter not found")

type Params map[string]interface{}

func (np *Params) GetString(key string) (string, error) {
	if v, present := (*np)[key]; present {
		return getString(v)
	}

	return "", ErrNotFound
}

func (np *Params) GetBool(key string) (bool, error) {
	if v, present := (*np)[key]; present {
		switch t := v.(type) {
		case float64:
			return t > 0.0, nil
		case int64:
			return t > 0, nil
		default:
			return false, ErrParamNotBool
		}
	}

	return false, ErrNotFound
}

func (np *Params) GetFloat(key string) (float64, error) {
	if v, present := (*np)[key]; present {
		return getFloat(v)
	}

	return 0.0, ErrNotFound
}

func (np *Params) GetInt(key string) (int64, error) {
	if v, present := (*np)[key]; present {
		return getInt(v)
	}

	return 0, ErrNotFound
}

func (np *Params) GetVec3(key string) (m.Vec3, error) {
	if v, present := (*np)[key]; present {
		return getVec3(v)
	}

	return m.Vec3{}, ErrNotFound
}

func getString(v interface{}) (string, error) {
	switch t := v.(type) {
	case string:
		return t, nil
	default:
		return "", ErrParamNotString
	}

}

func getFloat(v interface{}) (float64, error) {
	switch t := v.(type) {
	case float64:
		return t, nil
	case int64:
		return float64(t), nil
	default:
		return 0, ErrParamNotFloat
	}

}

func getInt(v interface{}) (int64, error) {
	switch t := v.(type) {
	case float64:
		return int64(t), nil
	case int64:
		return t, nil
	default:
		return 0, ErrParamNotInt
	}

}

func getVec3(v interface{}) (m.Vec3, error) {
	switch t := v.(type) {
	case []float64:
		if len(t) == 3 {
			return m.Vec3{float32(t[0]), float32(t[1]), float32(t[2])}, nil
		} else {
			return m.Vec3{}, nil
		}
	case []int64:
		if len(t) == 3 {
			return m.Vec3{float32(t[0]), float32(t[1]), float32(t[2])}, nil
		} else {
			return m.Vec3{}, nil
		}
	default:
		return m.Vec3{}, nil
	}
}

// This is a utility function for simple structs representing nodes
func (p *Params) Unmarshal(v interface{}) error {
	// Check v is struct

	rv := reflect.ValueOf(v)

	relem := rv.Elem()

	if relem.Kind() != reflect.Struct {
		return ErrNotStruct

	}

	ty := relem.Type()

	for i := 0; i < relem.NumField(); i++ {
		f := ty.Field(i)

		fieldName := f.Name

		if tag := f.Tag.Get("node"); tag != "" {
			fieldName = tag
		}

		fv := relem.Field(i)

		if fv.CanSet() {
			switch fv.Kind() {
			case reflect.Int:
				v, err := p.GetInt(fieldName)
				if err == nil {
					fv.SetInt(v)
				}
			case reflect.Float64:
				fallthrough
			case reflect.Float32:
				v, err := p.GetFloat(fieldName)
				if err == nil {
					fv.SetFloat(v)
				}
			case reflect.String:
				v, err := p.GetString(fieldName)
				if err == nil {
					fv.SetString(v)
				}
			case reflect.Bool:
				v, err := p.GetBool(fieldName)
				if err == nil {
					fv.SetBool(v)
				}

			default:
				switch fv.Type() {
				case reflect.TypeOf(m.Vec3{}):
					v, err := p.GetVec3(fieldName)
					if err == nil {
						fv.Set(reflect.ValueOf(v))
					}
				}
			}

		}
	}

	return nil
}
<|endoftext|>


<commit_before>package tracker

// http://bittorrent.org/beps/bep_0015.html
// http://xbtt.sourceforge.net/udp_tracker_protocol.html
// http://www.rasterbar.com/products/libtorrent/udp_tracker_protocol.html

import (
	"bytes"
	"encoding/binary"
	"errors"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/backoff"

	"github.com/cenkalti/rain/internal/protocol"
)

const connectionIDMagic = 0x41727101980
const connectionIDInterval = time.Minute

type action int32

// UDP tracker Actions
const (
	connect action = iota
	announce
	scrape
	errorAction
)

type transaction struct {
	request  udpReqeust
	response []byte
	err      error
	done     chan struct{}
}

func newTransaction(req udpReqeust) *transaction {
	req.SetTransactionID(rand.Int31())
	return &transaction{
		request: req,
		done:    make(chan struct{}),
	}
}

func (t *transaction) ID() int32 { return t.request.GetTransactionID() }
func (t *transaction) Done()     { close(t.done) }

type udpTracker struct {
	*trackerBase
	conn          *net.UDPConn
	transactions  map[int32]*transaction
	transactionsM sync.Mutex
	writeC        chan *transaction
}

func newUDPTracker(b *trackerBase) *udpTracker {
	return &udpTracker{
		trackerBase:  b,
		transactions: make(map[int32]*transaction),
		writeC:       make(chan *transaction),
	}
}

func (t *udpTracker) Dial() error {
	serverAddr, err := net.ResolveUDPAddr("udp", t.url.Host)
	if err != nil {
		return err
	}
	t.conn, err = net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		return err
	}
	go t.readLoop()
	go t.writeLoop()
	return nil
}

// Close the tracker connection.
// TODO end all goroutines.
func (t *udpTracker) Close() error {
	return t.conn.Close()
}

// readLoop reads datagrams from connection, finds the transaction and
// sends the bytes to the transaction's response channel.
func (t *udpTracker) readLoop() {
	// Read buffer must be big enough to hold a UDP packet of maximum expected size.
	// Current value is: 320 = 20 + 50*6 (AnnounceResponse with 50 peers)
	buf := make([]byte, 320)
	for {
		n, err := t.conn.Read(buf)
		if err != nil {
			t.log.Error(err)
			if nerr, ok := err.(net.Error); ok && !nerr.Temporary() {
				t.log.Debug("End of tracker read loop")
				return
			}
			continue
		}
		t.log.Debug("Read ", n, " bytes")

		var header udpMessageHeader
		if n < binary.Size(header) {
			t.log.Error("response is too small")
			continue
		}

		err = binary.Read(bytes.NewReader(buf), binary.BigEndian, &header)
		if err != nil {
			t.log.Error(err)
			continue
		}

		t.transactionsM.Lock()
		trx, ok := t.transactions[header.TransactionID]
		delete(t.transactions, header.TransactionID)
		t.transactionsM.Unlock()
		if !ok {
			t.log.Errorln("unexpected transaction_id:", header.TransactionID)
			continue
		}

		// Tracker has sent and error.
		if header.Action == errorAction {
			// The part after the header is the error message.
			trx.err = Error(buf[binary.Size(header):])
			trx.Done()
			continue
		}

		// Copy data into a new slice because buf will be overwritten at next read.
		trx.response = make([]byte, n)
		copy(trx.response, buf)
		trx.Done()
	}
}

// writeLoop receives a request from t.transactionC, sets a random TransactionID
// and sends it to the tracker.
func (t *udpTracker) writeLoop() {
	var connectionID int64
	var connectionIDtime time.Time

	for trx := range t.writeC {
		if time.Since(connectionIDtime) > connectionIDInterval {
			connectionID = t.connect()
			connectionIDtime = time.Now()
		}
		trx.request.SetConnectionID(connectionID)

		t.writeTrx(trx)
	}
}

func (t *udpTracker) writeTrx(trx *transaction) {
	t.log.Debugln("Writing transaction. ID:", trx.ID())
	err := binary.Write(t.conn, binary.BigEndian, trx.request)
	if err != nil {
		t.log.Error(err)
	}
}

// connect sends a connectRequest and returns a ConnectionID given by the tracker.
// On error, it backs off with the algorithm described in BEP15 and retries.
// It does not return until tracker sends a ConnectionID.
func (t *udpTracker) connect() int64 {
	req := new(connectRequest)
	req.SetAction(connect)
	req.SetConnectionID(connectionIDMagic)

	trx := newTransaction(req)

	for {
		data, err := t.retryTransaction(t.writeTrx, trx, nil) // Does not return until transaction is completed.
		if err != nil {
			t.log.Error(err)
			continue
		}

		var response connectResponse
		err = binary.Read(bytes.NewReader(data), binary.BigEndian, &response)
		if err != nil {
			t.log.Error(err)
			continue
		}

		if response.Action != connect {
			t.log.Error("invalid action in connect response")
			continue
		}

		t.log.Debugf("connect Response: %#v\n", response)
		return response.ConnectionID
	}
}

func (t *udpTracker) retryTransaction(f func(*transaction), trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	t.transactionsM.Lock()
	t.transactions[trx.ID()] = trx
	t.transactionsM.Unlock()

	ticker := backoff.NewTicker(UDPBackOff())
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			f(trx)
		case <-trx.done:
			// transaction is deleted in readLoop()
			return trx.response, trx.err
		case <-cancel:
			t.transactionsM.Lock()
			delete(t.transactions, trx.ID())
			t.transactionsM.Unlock()
			return nil, errors.New("transaction cancelled")
		}
	}
}

func (t *udpTracker) sendTransaction(trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	f := func(trx *transaction) { t.writeC <- trx }
	return t.retryTransaction(f, trx, cancel)
}

// Announce announces transfer to t periodically.
func (t *udpTracker) Announce(transfer Transfer, cancel <-chan struct{}, event <-chan Event, responseC chan<- *AnnounceResponse) {
	err := t.Dial()
	if err != nil {
		// TODO retry connecting to tracker
		t.log.Fatal(err)
	}

	// TODO below is same as HTTP tracker
	var nextAnnounce time.Duration
	var retry = *defaultRetryBackoff

	announce := func(e Event) {
		r, err := t.announce(transfer, e)
		if err != nil {
			t.log.Error(err)
			r = &AnnounceResponse{Error: err}
			nextAnnounce = retry.NextBackOff()
		} else {
			retry.Reset()
			nextAnnounce = r.Interval
		}
		select {
		case responseC <- r:
		case <-cancel:
			return
		}
	}

	announce(None)
	for {
		select {
		case <-time.After(nextAnnounce):
			announce(None)
		case e := <-event:
			announce(e)
		case <-cancel:
			return
		}
	}
}

func (t *udpTracker) announce(transfer Transfer, e Event) (*AnnounceResponse, error) {
	request := &announceRequest{
		InfoHash:   transfer.InfoHash(),
		PeerID:     t.peerID,
		Event:      e,
		IP:         0, // Tracker uses sender of this UDP packet.
		Key:        0, // TODO set it
		NumWant:    NumWant,
		Port:       t.port,
		Extensions: 0,
	}
	request.SetAction(announce)

	// TODO update on every try.
	request.update(transfer)

	// t.request may block, that's why we pass cancel as argument.
	trx := newTransaction(request)
	reply, err := t.sendTransaction(trx, nil) // TODO pass cancel instead of nil
	if err != nil {
		return nil, err
	}

	response, peers, err := t.parseAnnounceResponse(reply)
	if err != nil {
		return nil, err
	}
	t.log.Debugf("Announce response: %#v", response)

	return &AnnounceResponse{
		Error:    nil, // TODO handler error
		Interval: time.Duration(response.Interval) * time.Second,
		Leechers: response.Leechers,
		Seeders:  response.Seeders,
		Peers:    peers,
	}, nil
}

func (r *announceRequest) update(t Transfer) {
	r.Downloaded = t.Downloaded()
	r.Uploaded = t.Uploaded()
	r.Left = t.Left()
}

func (t *udpTracker) parseAnnounceResponse(data []byte) (*announceResponse, []Peer, error) {
	response := new(announceResponse)
	if len(data) < binary.Size(response) {
		return nil, nil, errors.New("response is too small")
	}

	reader := bytes.NewReader(data)

	err := binary.Read(reader, binary.BigEndian, response)
	if err != nil {
		return nil, nil, err
	}
	t.log.Debugf("annouceResponse: %#v", response)

	if response.Action != announce {
		return nil, nil, errors.New("invalid action")
	}

	peers, err := t.parsePeers(reader)
	if err != nil {
		return nil, nil, err
	}

	return response, peers, nil
}

type udpBackOff int

func (b *udpBackOff) NextBackOff() time.Duration {
	defer func() { *b++ }()
	if *b > 8 {
		*b = 8
	}
	return time.Duration(15*(2^*b)) * time.Second
}

func (b *udpBackOff) Reset() { *b = 0 }

var UDPBackOff = func() backoff.BackOff { return new(udpBackOff) }

type udpMessage interface {
	GetAction() action
	SetAction(action)
	GetTransactionID() int32
	SetTransactionID(int32)
}

type udpReqeust interface {
	udpMessage
	GetConnectionID() int64
	SetConnectionID(int64)
}

// udpMessageHeader implements udpMessage.
type udpMessageHeader struct {
	Action        action
	TransactionID int32
}

func (h *udpMessageHeader) GetAction() action         { return h.Action }
func (h *udpMessageHeader) SetAction(a action)        { h.Action = a }
func (h *udpMessageHeader) GetTransactionID() int32   { return h.TransactionID }
func (h *udpMessageHeader) SetTransactionID(id int32) { h.TransactionID = id }

// udpRequestHeader implements udpMessage and udpReqeust.
type udpRequestHeader struct {
	ConnectionID int64
	udpMessageHeader
}

func (h *udpRequestHeader) GetConnectionID() int64   { return h.ConnectionID }
func (h *udpRequestHeader) SetConnectionID(id int64) { h.ConnectionID = id }

type connectRequest struct {
	udpRequestHeader
}

type connectResponse struct {
	udpMessageHeader
	ConnectionID int64
}

type announceRequest struct {
	udpRequestHeader
	InfoHash   protocol.InfoHash
	PeerID     protocol.PeerID
	Downloaded int64
	Left       int64
	Uploaded   int64
	Event      Event
	IP         uint32
	Key        uint32
	NumWant    int32
	Port       uint16
	Extensions uint16
}

type announceResponse struct {
	udpMessageHeader
	Interval int32
	Leechers int32
	Seeders  int32
}
<commit_msg>udp dial once<commit_after>package tracker

// http://bittorrent.org/beps/bep_0015.html
// http://xbtt.sourceforge.net/udp_tracker_protocol.html
// http://www.rasterbar.com/products/libtorrent/udp_tracker_protocol.html

import (
	"bytes"
	"encoding/binary"
	"errors"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/backoff"

	"github.com/cenkalti/rain/internal/protocol"
)

const connectionIDMagic = 0x41727101980
const connectionIDInterval = time.Minute

type action int32

// UDP tracker Actions
const (
	connect action = iota
	announce
	scrape
	errorAction
)

type transaction struct {
	request  udpReqeust
	response []byte
	err      error
	done     chan struct{}
}

func newTransaction(req udpReqeust) *transaction {
	req.SetTransactionID(rand.Int31())
	return &transaction{
		request: req,
		done:    make(chan struct{}),
	}
}

func (t *transaction) ID() int32 { return t.request.GetTransactionID() }
func (t *transaction) Done()     { close(t.done) }

type udpTracker struct {
	*trackerBase
	conn          *net.UDPConn
	dialMutex     sync.Mutex
	connected     bool
	transactions  map[int32]*transaction
	transactionsM sync.Mutex
	writeC        chan *transaction
}

func newUDPTracker(b *trackerBase) *udpTracker {
	return &udpTracker{
		trackerBase:  b,
		transactions: make(map[int32]*transaction),
		writeC:       make(chan *transaction),
	}
}

func (t *udpTracker) dial() error {
	serverAddr, err := net.ResolveUDPAddr("udp", t.url.Host)
	if err != nil {
		return err
	}
	t.conn, err = net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		return err
	}
	go t.readLoop()
	go t.writeLoop()
	return nil
}

// Close the tracker connection.
// TODO end all goroutines.
func (t *udpTracker) Close() error {
	return t.conn.Close()
}

// readLoop reads datagrams from connection, finds the transaction and
// sends the bytes to the transaction's response channel.
func (t *udpTracker) readLoop() {
	// Read buffer must be big enough to hold a UDP packet of maximum expected size.
	// Current value is: 320 = 20 + 50*6 (AnnounceResponse with 50 peers)
	buf := make([]byte, 320)
	for {
		n, err := t.conn.Read(buf)
		if err != nil {
			t.log.Error(err)
			if nerr, ok := err.(net.Error); ok && !nerr.Temporary() {
				t.log.Debug("End of tracker read loop")
				return
			}
			continue
		}
		t.log.Debug("Read ", n, " bytes")

		var header udpMessageHeader
		if n < binary.Size(header) {
			t.log.Error("response is too small")
			continue
		}

		err = binary.Read(bytes.NewReader(buf), binary.BigEndian, &header)
		if err != nil {
			t.log.Error(err)
			continue
		}

		t.transactionsM.Lock()
		trx, ok := t.transactions[header.TransactionID]
		delete(t.transactions, header.TransactionID)
		t.transactionsM.Unlock()
		if !ok {
			t.log.Errorln("unexpected transaction_id:", header.TransactionID)
			continue
		}

		// Tracker has sent and error.
		if header.Action == errorAction {
			// The part after the header is the error message.
			trx.err = Error(buf[binary.Size(header):])
			trx.Done()
			continue
		}

		// Copy data into a new slice because buf will be overwritten at next read.
		trx.response = make([]byte, n)
		copy(trx.response, buf)
		trx.Done()
	}
}

// writeLoop receives a request from t.transactionC, sets a random TransactionID
// and sends it to the tracker.
func (t *udpTracker) writeLoop() {
	var connectionID int64
	var connectionIDtime time.Time

	for trx := range t.writeC {
		if time.Since(connectionIDtime) > connectionIDInterval {
			connectionID = t.connect()
			connectionIDtime = time.Now()
		}
		trx.request.SetConnectionID(connectionID)

		t.writeTrx(trx)
	}
}

func (t *udpTracker) writeTrx(trx *transaction) {
	t.log.Debugln("Writing transaction. ID:", trx.ID())
	err := binary.Write(t.conn, binary.BigEndian, trx.request)
	if err != nil {
		t.log.Error(err)
	}
}

// connect sends a connectRequest and returns a ConnectionID given by the tracker.
// On error, it backs off with the algorithm described in BEP15 and retries.
// It does not return until tracker sends a ConnectionID.
func (t *udpTracker) connect() int64 {
	req := new(connectRequest)
	req.SetAction(connect)
	req.SetConnectionID(connectionIDMagic)

	trx := newTransaction(req)

	for {
		data, err := t.retryTransaction(t.writeTrx, trx, nil) // Does not return until transaction is completed.
		if err != nil {
			t.log.Error(err)
			continue
		}

		var response connectResponse
		err = binary.Read(bytes.NewReader(data), binary.BigEndian, &response)
		if err != nil {
			t.log.Error(err)
			continue
		}

		if response.Action != connect {
			t.log.Error("invalid action in connect response")
			continue
		}

		t.log.Debugf("connect Response: %#v\n", response)
		return response.ConnectionID
	}
}

func (t *udpTracker) retryTransaction(f func(*transaction), trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	t.transactionsM.Lock()
	t.transactions[trx.ID()] = trx
	t.transactionsM.Unlock()

	ticker := backoff.NewTicker(UDPBackOff())
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			f(trx)
		case <-trx.done:
			// transaction is deleted in readLoop()
			return trx.response, trx.err
		case <-cancel:
			t.transactionsM.Lock()
			delete(t.transactions, trx.ID())
			t.transactionsM.Unlock()
			return nil, errors.New("transaction cancelled")
		}
	}
}

func (t *udpTracker) sendTransaction(trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	f := func(trx *transaction) { t.writeC <- trx }
	return t.retryTransaction(f, trx, cancel)
}

// Announce announces transfer to t periodically.
func (t *udpTracker) Announce(transfer Transfer, cancel <-chan struct{}, event <-chan Event, responseC chan<- *AnnounceResponse) {
	// TODO below is same as HTTP tracker
	var nextAnnounce time.Duration
	var retry = *defaultRetryBackoff

	announce := func(e Event) {
		r, err := t.announce(transfer, e)
		if err != nil {
			t.log.Error(err)
			r = &AnnounceResponse{Error: err}
			nextAnnounce = retry.NextBackOff()
		} else {
			retry.Reset()
			nextAnnounce = r.Interval
		}
		select {
		case responseC <- r:
		case <-cancel:
			return
		}
	}

	announce(None)
	for {
		select {
		case <-time.After(nextAnnounce):
			announce(None)
		case e := <-event:
			announce(e)
		case <-cancel:
			return
		}
	}
}

func (t *udpTracker) announce(transfer Transfer, e Event) (*AnnounceResponse, error) {
	t.dialMutex.Lock()
	if !t.connected {
		err := t.dial()
		if err != nil {
			t.dialMutex.Unlock()
			return nil, err
		}
		t.connected = true
	}
	t.dialMutex.Unlock()

	request := &announceRequest{
		InfoHash:   transfer.InfoHash(),
		PeerID:     t.peerID,
		Event:      e,
		IP:         0, // Tracker uses sender of this UDP packet.
		Key:        0, // TODO set it
		NumWant:    NumWant,
		Port:       t.port,
		Extensions: 0,
	}
	request.SetAction(announce)

	// TODO update on every try.
	request.update(transfer)

	// t.request may block, that's why we pass cancel as argument.
	trx := newTransaction(request)
	reply, err := t.sendTransaction(trx, nil) // TODO pass cancel instead of nil
	if err != nil {
		return nil, err
	}

	response, peers, err := t.parseAnnounceResponse(reply)
	if err != nil {
		return nil, err
	}
	t.log.Debugf("Announce response: %#v", response)

	return &AnnounceResponse{
		Error:    nil, // TODO handler error
		Interval: time.Duration(response.Interval) * time.Second,
		Leechers: response.Leechers,
		Seeders:  response.Seeders,
		Peers:    peers,
	}, nil
}

func (r *announceRequest) update(t Transfer) {
	r.Downloaded = t.Downloaded()
	r.Uploaded = t.Uploaded()
	r.Left = t.Left()
}

func (t *udpTracker) parseAnnounceResponse(data []byte) (*announceResponse, []Peer, error) {
	response := new(announceResponse)
	if len(data) < binary.Size(response) {
		return nil, nil, errors.New("response is too small")
	}

	reader := bytes.NewReader(data)

	err := binary.Read(reader, binary.BigEndian, response)
	if err != nil {
		return nil, nil, err
	}
	t.log.Debugf("annouceResponse: %#v", response)

	if response.Action != announce {
		return nil, nil, errors.New("invalid action")
	}

	peers, err := t.parsePeers(reader)
	if err != nil {
		return nil, nil, err
	}

	return response, peers, nil
}

type udpBackOff int

func (b *udpBackOff) NextBackOff() time.Duration {
	defer func() { *b++ }()
	if *b > 8 {
		*b = 8
	}
	return time.Duration(15*(2^*b)) * time.Second
}

func (b *udpBackOff) Reset() { *b = 0 }

var UDPBackOff = func() backoff.BackOff { return new(udpBackOff) }

type udpMessage interface {
	GetAction() action
	SetAction(action)
	GetTransactionID() int32
	SetTransactionID(int32)
}

type udpReqeust interface {
	udpMessage
	GetConnectionID() int64
	SetConnectionID(int64)
}

// udpMessageHeader implements udpMessage.
type udpMessageHeader struct {
	Action        action
	TransactionID int32
}

func (h *udpMessageHeader) GetAction() action         { return h.Action }
func (h *udpMessageHeader) SetAction(a action)        { h.Action = a }
func (h *udpMessageHeader) GetTransactionID() int32   { return h.TransactionID }
func (h *udpMessageHeader) SetTransactionID(id int32) { h.TransactionID = id }

// udpRequestHeader implements udpMessage and udpReqeust.
type udpRequestHeader struct {
	ConnectionID int64
	udpMessageHeader
}

func (h *udpRequestHeader) GetConnectionID() int64   { return h.ConnectionID }
func (h *udpRequestHeader) SetConnectionID(id int64) { h.ConnectionID = id }

type connectRequest struct {
	udpRequestHeader
}

type connectResponse struct {
	udpMessageHeader
	ConnectionID int64
}

type announceRequest struct {
	udpRequestHeader
	InfoHash   protocol.InfoHash
	PeerID     protocol.PeerID
	Downloaded int64
	Left       int64
	Uploaded   int64
	Event      Event
	IP         uint32
	Key        uint32
	NumWant    int32
	Port       uint16
	Extensions uint16
}

type announceResponse struct {
	udpMessageHeader
	Interval int32
	Leechers int32
	Seeders  int32
}
<|endoftext|>


<commit_before>package json

import (
	"io"
	"strconv"
	"time"
)

var (
	nullBytes  = [...]byte{'n', 'u', 'l', 'l'}
	trueBytes  = [...]byte{'t', 'r', 'u', 'e'}
	falseBytes = [...]byte{'f', 'a', 'l', 's', 'e'}

	arrayOpen  = [...]byte{'['}
	arrayClose = [...]byte{']'}

	mapOpen  = [...]byte{'{'}
	mapClose = [...]byte{'}'}

	comma  = [...]byte{','}
	column = [...]byte{':'}
)

// Emitter implements a JSON emitter that satisfies the objconv.Emitter
// interface.
type Emitter struct {
	b [128]byte
	s []byte
	w io.Writer
}

func NewEmitter(w io.Writer) *Emitter {
	return &Emitter{w: w}
}

func (e *Emitter) Reset(w io.Writer) {
	e.w = w
}

func (e *Emitter) EmitNil() (err error) {
	_, err = e.w.Write(nullBytes[:])
	return
}

func (e *Emitter) EmitBool(v bool) (err error) {
	if v {
		_, err = e.w.Write(trueBytes[:])
	} else {
		_, err = e.w.Write(falseBytes[:])
	}
	return
}

func (e *Emitter) EmitInt(v int64) (err error) {
	_, err = e.w.Write(strconv.AppendInt(e.b[:0], v, 10))
	return
}

func (e *Emitter) EmitUint(v uint64) (err error) {
	_, err = e.w.Write(strconv.AppendUint(e.b[:0], v, 10))
	return
}

func (e *Emitter) EmitFloat(v float64) (err error) {
	_, err = e.w.Write(strconv.AppendFloat(e.b[:0], v, 'g', -1, 64))
	return
}

func (e *Emitter) EmitString(v string) (err error) {
	i := 0
	j := 0
	n := len(v)

	if e.s == nil {
		e.s = e.b[:0]
	}

	s := e.s[:0]
	s = append(s, '"')

	for j != n {
		b := v[j]
		j++

		switch b {
		case '"', '\\', '/':
			// b = b

		case '\b':
			b = 'b'

		case '\f':
			b = 'f'

		case '\n':
			b = 'n'

		case '\r':
			b = 'r'

		case '\t':
			b = 't'

		default:
			continue
		}

		s = append(s, v[i:j-1]...)
		s = append(s, '\\', b)
		i = j
	}

	s = append(s, v[i:j]...)
	s = append(s, '"')
	e.s = s[:0] // in case the buffer was reallocated

	_, err = e.w.Write(s)
	return
}

func (e *Emitter) EmitBytes(v []byte) error {
	return e.EmitString(string(v))
}

func (e *Emitter) EmitTime(v time.Time) error {
	return e.EmitString(v.Format(time.RFC3339Nano))
}

func (e *Emitter) EmitDuration(v time.Duration) error {
	return e.EmitString(v.String())
}

func (e *Emitter) EmitError(v error) error {
	return e.EmitString(v.Error())
}

func (e *Emitter) EmitArrayBegin(_ int) (err error) {
	_, err = e.w.Write(arrayOpen[:])
	return
}

func (e *Emitter) EmitArrayEnd() (err error) {
	_, err = e.w.Write(arrayClose[:])
	return
}

func (e *Emitter) EmitArrayNext() (err error) {
	_, err = e.w.Write(comma[:])
	return
}

func (e *Emitter) EmitMapBegin(_ int) (err error) {
	_, err = e.w.Write(mapOpen[:])
	return
}

func (e *Emitter) EmitMapEnd() (err error) {
	_, err = e.w.Write(mapClose[:])
	return
}

func (e *Emitter) EmitMapValue() (err error) {
	_, err = e.w.Write(column[:])
	return
}

func (e *Emitter) EmitMapNext() (err error) {
	_, err = e.w.Write(comma[:])
	return
}
<commit_msg>optimize json time encoding<commit_after>package json

import (
	"io"
	"strconv"
	"time"
)

var (
	nullBytes  = [...]byte{'n', 'u', 'l', 'l'}
	trueBytes  = [...]byte{'t', 'r', 'u', 'e'}
	falseBytes = [...]byte{'f', 'a', 'l', 's', 'e'}

	arrayOpen  = [...]byte{'['}
	arrayClose = [...]byte{']'}

	mapOpen  = [...]byte{'{'}
	mapClose = [...]byte{'}'}

	comma  = [...]byte{','}
	column = [...]byte{':'}
)

// Emitter implements a JSON emitter that satisfies the objconv.Emitter
// interface.
type Emitter struct {
	b [128]byte
	s []byte
	w io.Writer
}

func NewEmitter(w io.Writer) *Emitter {
	return &Emitter{w: w}
}

func (e *Emitter) Reset(w io.Writer) {
	e.w = w
}

func (e *Emitter) EmitNil() (err error) {
	_, err = e.w.Write(nullBytes[:])
	return
}

func (e *Emitter) EmitBool(v bool) (err error) {
	if v {
		_, err = e.w.Write(trueBytes[:])
	} else {
		_, err = e.w.Write(falseBytes[:])
	}
	return
}

func (e *Emitter) EmitInt(v int64) (err error) {
	_, err = e.w.Write(strconv.AppendInt(e.b[:0], v, 10))
	return
}

func (e *Emitter) EmitUint(v uint64) (err error) {
	_, err = e.w.Write(strconv.AppendUint(e.b[:0], v, 10))
	return
}

func (e *Emitter) EmitFloat(v float64) (err error) {
	_, err = e.w.Write(strconv.AppendFloat(e.b[:0], v, 'g', -1, 64))
	return
}

func (e *Emitter) EmitString(v string) (err error) {
	i := 0
	j := 0
	n := len(v)

	if e.s == nil {
		e.s = e.b[:0]
	}

	s := e.s[:0]
	s = append(s, '"')

	for j != n {
		b := v[j]
		j++

		switch b {
		case '"', '\\', '/':
			// b = b

		case '\b':
			b = 'b'

		case '\f':
			b = 'f'

		case '\n':
			b = 'n'

		case '\r':
			b = 'r'

		case '\t':
			b = 't'

		default:
			continue
		}

		s = append(s, v[i:j-1]...)
		s = append(s, '\\', b)
		i = j
	}

	s = append(s, v[i:j]...)
	s = append(s, '"')
	e.s = s[:0] // in case the buffer was reallocated

	_, err = e.w.Write(s)
	return
}

func (e *Emitter) EmitBytes(v []byte) error {
	return e.EmitString(string(v))
}

func (e *Emitter) EmitTime(v time.Time) error {
	return e.EmitString(string(v.AppendFormat(e.b[:0], time.RFC3339Nano)))
}

func (e *Emitter) EmitDuration(v time.Duration) error {
	return e.EmitString(v.String())
}

func (e *Emitter) EmitError(v error) error {
	return e.EmitString(v.Error())
}

func (e *Emitter) EmitArrayBegin(_ int) (err error) {
	_, err = e.w.Write(arrayOpen[:])
	return
}

func (e *Emitter) EmitArrayEnd() (err error) {
	_, err = e.w.Write(arrayClose[:])
	return
}

func (e *Emitter) EmitArrayNext() (err error) {
	_, err = e.w.Write(comma[:])
	return
}

func (e *Emitter) EmitMapBegin(_ int) (err error) {
	_, err = e.w.Write(mapOpen[:])
	return
}

func (e *Emitter) EmitMapEnd() (err error) {
	_, err = e.w.Write(mapClose[:])
	return
}

func (e *Emitter) EmitMapValue() (err error) {
	_, err = e.w.Write(column[:])
	return
}

func (e *Emitter) EmitMapNext() (err error) {
	_, err = e.w.Write(comma[:])
	return
}
<|endoftext|>