<commit_before>package torrent

import (
	"bufio"
	"bytes"
	"container/list"
	"encoding"
	"errors"
	"expvar"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/anacrolix/torrent/bencode"
	"github.com/anacrolix/torrent/internal/pieceordering"
	pp "github.com/anacrolix/torrent/peer_protocol"
)

var optimizedCancels = expvar.NewInt("optimizedCancels")

type peerSource byte

const (
	peerSourceIncoming = 'I'
	peerSourceDHT      = 'H'
	peerSourcePEX      = 'X'
)

// Maintains the state of a connection with a peer.
type connection struct {
	conn      net.Conn
	rw        io.ReadWriter // The real slim shady
	encrypted bool
	Discovery peerSource
	uTP       bool
	closing   chan struct{}
	mu        sync.Mutex // Only for closing.
	post      chan pp.Message
	writeCh   chan []byte

	// The connection's preferred order to download pieces. The index is the
	// piece, the value is its priority.
	piecePriorities []int
	// The piece request order based on piece priorities.
	pieceRequestOrder *pieceordering.Instance

	UnwantedChunksReceived int
	UsefulChunksReceived   int
	chunksSent             int

	lastMessageReceived     time.Time
	completedHandshake      time.Time
	lastUsefulChunkReceived time.Time
	lastChunkSent           time.Time

	// Stuff controlled by the local peer.
	Interested       bool
	Choked           bool
	Requests         map[request]struct{}
	requestsLowWater int
	// Indexed by metadata piece, set to true if posted and pending a
	// response.
	metadataRequests []bool

	// Stuff controlled by the remote peer.
	PeerID             [20]byte
	PeerInterested     bool
	PeerChoked         bool
	PeerRequests       map[request]struct{}
	PeerExtensionBytes peerExtensionBytes
	// Whether the peer has the given piece. nil if they've not sent any
	// related messages yet.
	PeerPieces []bool
	peerHasAll bool

	PeerMaxRequests  int // Maximum pending requests the peer allows.
	PeerExtensionIDs map[string]byte
	PeerClientName   string
}

func newConnection() (c *connection) {
	c = &connection{
		Choked:          true,
		PeerChoked:      true,
		PeerMaxRequests: 250,

		closing: make(chan struct{}),
		writeCh: make(chan []byte),
		post:    make(chan pp.Message),
	}
	return
}

func (cn *connection) remoteAddr() net.Addr {
	return cn.conn.RemoteAddr()
}

func (cn *connection) localAddr() net.Addr {
	return cn.conn.LocalAddr()
}

// Adjust piece position in the request order for this connection based on the
// given piece priority.
func (cn *connection) pendPiece(piece int, priority piecePriority) {
	if priority == PiecePriorityNone {
		cn.pieceRequestOrder.DeletePiece(piece)
		return
	}
	pp := cn.piecePriorities[piece]
	// Priority regions not to scale. Within each region, piece is randomized
	// according to connection.

	// <-request first -- last->
	// [ Now         ]
	//  [ Next       ]
	//   [ Readahead ]
	//                [ Normal ]
	key := func() int {
		switch priority {
		case PiecePriorityNow:
			return -3*len(cn.piecePriorities) + 3*pp
		case PiecePriorityNext:
			return -2*len(cn.piecePriorities) + 2*pp
		case PiecePriorityReadahead:
			return -len(cn.piecePriorities) + pp
		case PiecePriorityNormal:
			return pp
		default:
			panic(priority)
		}
	}()
	cn.pieceRequestOrder.SetPiece(piece, key)
}

func (cn *connection) supportsExtension(ext string) bool {
	_, ok := cn.PeerExtensionIDs[ext]
	return ok
}

func (cn *connection) completedString(t *torrent) string {
	if cn.PeerPieces == nil && !cn.peerHasAll {
		return "?"
	}
	return fmt.Sprintf("%d/%d", func() int {
		if cn.peerHasAll {
			if t.haveInfo() {
				return t.numPieces()
			}
			return -1
		}
		ret := 0
		for _, b := range cn.PeerPieces {
			if b {
				ret++
			}
		}
		return ret
	}(), func() int {
		if cn.peerHasAll || cn.PeerPieces == nil {
			if t.haveInfo() {
				return t.numPieces()
			}
			return -1
		}
		return len(cn.PeerPieces)
	}())
}

// Correct the PeerPieces slice length. Return false if the existing slice is
// invalid, such as by receiving badly sized BITFIELD, or invalid HAVE
// messages.
func (cn *connection) setNumPieces(num int) error {
	if cn.peerHasAll {
		return nil
	}
	if cn.PeerPieces == nil {
		return nil
	}
	if len(cn.PeerPieces) == num {
	} else if len(cn.PeerPieces) < num {
		cn.PeerPieces = append(cn.PeerPieces, make([]bool, num-len(cn.PeerPieces))...)
	} else if len(cn.PeerPieces) <= (num+7)/8*8 {
		for _, have := range cn.PeerPieces[num:] {
			if have {
				return errors.New("peer has invalid piece")
			}
		}
		cn.PeerPieces = cn.PeerPieces[:num]
	} else {
		return fmt.Errorf("peer bitfield is excessively long: expected %d, have %d", num, len(cn.PeerPieces))
	}
	if len(cn.PeerPieces) != num {
		panic("wat")
	}
	return nil
}

func eventAgeString(t time.Time) string {
	if t.IsZero() {
		return "never"
	}
	return fmt.Sprintf("%.2fs ago", time.Now().Sub(t).Seconds())
}

func (cn *connection) connectionFlags() (ret string) {
	c := func(b byte) {
		ret += string([]byte{b})
	}
	if cn.encrypted {
		c('E')
	}
	if cn.Discovery != 0 {
		c(byte(cn.Discovery))
	}
	if cn.uTP {
		c('T')
	}
	return
}

// Inspired by https://trac.transmissionbt.com/wiki/PeerStatusText
func (cn *connection) statusFlags() (ret string) {
	c := func(b byte) {
		ret += string([]byte{b})
	}
	if cn.Interested {
		c('i')
	}
	if cn.Choked {
		c('c')
	}
	c('-')
	ret += cn.connectionFlags()
	c('-')
	if cn.PeerInterested {
		c('i')
	}
	if cn.PeerChoked {
		c('c')
	}
	return
}

func (cn *connection) String() string {
	var buf bytes.Buffer
	cn.WriteStatus(&buf, nil)
	return buf.String()
}

func (cn *connection) WriteStatus(w io.Writer, t *torrent) {
	// \t isn't preserved in <pre> blocks?
	fmt.Fprintf(w, "%+q: %s-%s\n", cn.PeerID, cn.localAddr(), cn.remoteAddr())
	fmt.Fprintf(w, "    last msg: %s, connected: %s, last useful chunk: %s\n",
		eventAgeString(cn.lastMessageReceived),
		eventAgeString(cn.completedHandshake),
		eventAgeString(cn.lastUsefulChunkReceived))
	fmt.Fprintf(w,
		"    %s completed, good chunks: %d/%d-%d reqq: %d-%d, flags: %s\n",
		cn.completedString(t),
		cn.UsefulChunksReceived,
		cn.UnwantedChunksReceived+cn.UsefulChunksReceived,
		cn.chunksSent,
		len(cn.Requests),
		len(cn.PeerRequests),
		cn.statusFlags(),
	)
}

func (c *connection) Close() {
	c.mu.Lock()
	defer c.mu.Unlock()
	select {
	case <-c.closing:
		return
	default:
	}
	close(c.closing)
	// TODO: This call blocks sometimes, why?
	go c.conn.Close()
}

func (c *connection) PeerHasPiece(piece int) bool {
	if c.peerHasAll {
		return true
	}
	if piece >= len(c.PeerPieces) {
		return false
	}
	return c.PeerPieces[piece]
}

func (c *connection) Post(msg pp.Message) {
	select {
	case c.post <- msg:
	case <-c.closing:
	}
}

func (c *connection) RequestPending(r request) bool {
	_, ok := c.Requests[r]
	return ok
}

func (c *connection) requestMetadataPiece(index int) {
	eID := c.PeerExtensionIDs["ut_metadata"]
	if eID == 0 {
		return
	}
	if index < len(c.metadataRequests) && c.metadataRequests[index] {
		return
	}
	c.Post(pp.Message{
		Type:       pp.Extended,
		ExtendedID: eID,
		ExtendedPayload: func() []byte {
			b, err := bencode.Marshal(map[string]int{
				"msg_type": pp.RequestMetadataExtensionMsgType,
				"piece":    index,
			})
			if err != nil {
				panic(err)
			}
			return b
		}(),
	})
	for index >= len(c.metadataRequests) {
		c.metadataRequests = append(c.metadataRequests, false)
	}
	c.metadataRequests[index] = true
}

func (c *connection) requestedMetadataPiece(index int) bool {
	return index < len(c.metadataRequests) && c.metadataRequests[index]
}

// Returns true if more requests can be sent.
func (c *connection) Request(chunk request) bool {
	if len(c.Requests) >= c.PeerMaxRequests {
		return false
	}
	if !c.PeerHasPiece(int(chunk.Index)) {
		return true
	}
	if c.RequestPending(chunk) {
		return true
	}
	c.SetInterested(true)
	if c.PeerChoked {
		return false
	}
	if c.Requests == nil {
		c.Requests = make(map[request]struct{}, c.PeerMaxRequests)
	}
	c.Requests[chunk] = struct{}{}
	c.requestsLowWater = len(c.Requests) / 2
	c.Post(pp.Message{
		Type:   pp.Request,
		Index:  chunk.Index,
		Begin:  chunk.Begin,
		Length: chunk.Length,
	})
	return true
}

// Returns true if an unsatisfied request was canceled.
func (c *connection) Cancel(r request) bool {
	if c.Requests == nil {
		return false
	}
	if _, ok := c.Requests[r]; !ok {
		return false
	}
	delete(c.Requests, r)
	c.Post(pp.Message{
		Type:   pp.Cancel,
		Index:  r.Index,
		Begin:  r.Begin,
		Length: r.Length,
	})
	return true
}

// Returns true if an unsatisfied request was canceled.
func (c *connection) PeerCancel(r request) bool {
	if c.PeerRequests == nil {
		return false
	}
	if _, ok := c.PeerRequests[r]; !ok {
		return false
	}
	delete(c.PeerRequests, r)
	return true
}

func (c *connection) Choke() {
	if c.Choked {
		return
	}
	c.Post(pp.Message{
		Type: pp.Choke,
	})
	c.PeerRequests = nil
	c.Choked = true
}

func (c *connection) Unchoke() {
	if !c.Choked {
		return
	}
	c.Post(pp.Message{
		Type: pp.Unchoke,
	})
	c.Choked = false
}

func (c *connection) SetInterested(interested bool) {
	if c.Interested == interested {
		return
	}
	c.Post(pp.Message{
		Type: func() pp.MessageType {
			if interested {
				return pp.Interested
			} else {
				return pp.NotInterested
			}
		}(),
	})
	c.Interested = interested
}

// Writes buffers to the socket from the write channel.
func (conn *connection) writer() {
	// Reduce write syscalls.
	buf := bufio.NewWriterSize(conn.rw, 0x8000) // 32 KiB
	// Receives when buf is not empty.
	notEmpty := make(chan struct{}, 1)
	for {
		if buf.Buffered() != 0 {
			// Make sure it's receivable.
			select {
			case notEmpty <- struct{}{}:
			default:
			}
		}
		select {
		case b, ok := <-conn.writeCh:
			if !ok {
				return
			}
			_, err := buf.Write(b)
			if err != nil {
				conn.Close()
				return
			}
		case <-conn.closing:
			return
		case <-notEmpty:
			err := buf.Flush()
			if err != nil {
				return
			}
		}
	}
}

func (conn *connection) writeOptimizer(keepAliveDelay time.Duration) {
	defer close(conn.writeCh) // Responsible for notifying downstream routines.
	pending := list.New()     // Message queue.
	var nextWrite []byte      // Set to nil if we need to need to marshal the next message.
	timer := time.NewTimer(keepAliveDelay)
	defer timer.Stop()
	lastWrite := time.Now()
	for {
		write := conn.writeCh // Set to nil if there's nothing to write.
		if pending.Len() == 0 {
			write = nil
		} else if nextWrite == nil {
			var err error
			nextWrite, err = pending.Front().Value.(encoding.BinaryMarshaler).MarshalBinary()
			if err != nil {
				panic(err)
			}
		}
	event:
		select {
		case <-timer.C:
			if pending.Len() != 0 {
				break
			}
			keepAliveTime := lastWrite.Add(keepAliveDelay)
			if time.Now().Before(keepAliveTime) {
				timer.Reset(keepAliveTime.Sub(time.Now()))
				break
			}
			pending.PushBack(pp.Message{Keepalive: true})
		case msg, ok := <-conn.post:
			if !ok {
				return
			}
			if msg.Type == pp.Cancel {
				for e := pending.Back(); e != nil; e = e.Prev() {
					elemMsg := e.Value.(pp.Message)
					if elemMsg.Type == pp.Request && msg.Index == elemMsg.Index && msg.Begin == elemMsg.Begin && msg.Length == elemMsg.Length {
						pending.Remove(e)
						optimizedCancels.Add(1)
						break event
					}
				}
			}
			pending.PushBack(msg)
		case write <- nextWrite:
			pending.Remove(pending.Front())
			nextWrite = nil
			lastWrite = time.Now()
			if pending.Len() == 0 {
				timer.Reset(keepAliveDelay)
			}
		case <-conn.closing:
			return
		}
	}
}
<commit_msg>Track connection writes and flushes<commit_after>package torrent

import (
	"bufio"
	"bytes"
	"container/list"
	"encoding"
	"errors"
	"expvar"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/anacrolix/torrent/bencode"
	"github.com/anacrolix/torrent/internal/pieceordering"
	pp "github.com/anacrolix/torrent/peer_protocol"
)

var optimizedCancels = expvar.NewInt("optimizedCancels")

type peerSource byte

const (
	peerSourceIncoming = 'I'
	peerSourceDHT      = 'H'
	peerSourcePEX      = 'X'
)

// Maintains the state of a connection with a peer.
type connection struct {
	conn      net.Conn
	rw        io.ReadWriter // The real slim shady
	encrypted bool
	Discovery peerSource
	uTP       bool
	closing   chan struct{}
	mu        sync.Mutex // Only for closing.
	post      chan pp.Message
	writeCh   chan []byte

	// The connection's preferred order to download pieces. The index is the
	// piece, the value is its priority.
	piecePriorities []int
	// The piece request order based on piece priorities.
	pieceRequestOrder *pieceordering.Instance

	UnwantedChunksReceived int
	UsefulChunksReceived   int
	chunksSent             int

	lastMessageReceived     time.Time
	completedHandshake      time.Time
	lastUsefulChunkReceived time.Time
	lastChunkSent           time.Time

	// Stuff controlled by the local peer.
	Interested       bool
	Choked           bool
	Requests         map[request]struct{}
	requestsLowWater int
	// Indexed by metadata piece, set to true if posted and pending a
	// response.
	metadataRequests []bool

	// Stuff controlled by the remote peer.
	PeerID             [20]byte
	PeerInterested     bool
	PeerChoked         bool
	PeerRequests       map[request]struct{}
	PeerExtensionBytes peerExtensionBytes
	// Whether the peer has the given piece. nil if they've not sent any
	// related messages yet.
	PeerPieces []bool
	peerHasAll bool

	PeerMaxRequests  int // Maximum pending requests the peer allows.
	PeerExtensionIDs map[string]byte
	PeerClientName   string
}

func newConnection() (c *connection) {
	c = &connection{
		Choked:          true,
		PeerChoked:      true,
		PeerMaxRequests: 250,

		closing: make(chan struct{}),
		writeCh: make(chan []byte),
		post:    make(chan pp.Message),
	}
	return
}

func (cn *connection) remoteAddr() net.Addr {
	return cn.conn.RemoteAddr()
}

func (cn *connection) localAddr() net.Addr {
	return cn.conn.LocalAddr()
}

// Adjust piece position in the request order for this connection based on the
// given piece priority.
func (cn *connection) pendPiece(piece int, priority piecePriority) {
	if priority == PiecePriorityNone {
		cn.pieceRequestOrder.DeletePiece(piece)
		return
	}
	pp := cn.piecePriorities[piece]
	// Priority regions not to scale. Within each region, piece is randomized
	// according to connection.

	// <-request first -- last->
	// [ Now         ]
	//  [ Next       ]
	//   [ Readahead ]
	//                [ Normal ]
	key := func() int {
		switch priority {
		case PiecePriorityNow:
			return -3*len(cn.piecePriorities) + 3*pp
		case PiecePriorityNext:
			return -2*len(cn.piecePriorities) + 2*pp
		case PiecePriorityReadahead:
			return -len(cn.piecePriorities) + pp
		case PiecePriorityNormal:
			return pp
		default:
			panic(priority)
		}
	}()
	cn.pieceRequestOrder.SetPiece(piece, key)
}

func (cn *connection) supportsExtension(ext string) bool {
	_, ok := cn.PeerExtensionIDs[ext]
	return ok
}

func (cn *connection) completedString(t *torrent) string {
	if cn.PeerPieces == nil && !cn.peerHasAll {
		return "?"
	}
	return fmt.Sprintf("%d/%d", func() int {
		if cn.peerHasAll {
			if t.haveInfo() {
				return t.numPieces()
			}
			return -1
		}
		ret := 0
		for _, b := range cn.PeerPieces {
			if b {
				ret++
			}
		}
		return ret
	}(), func() int {
		if cn.peerHasAll || cn.PeerPieces == nil {
			if t.haveInfo() {
				return t.numPieces()
			}
			return -1
		}
		return len(cn.PeerPieces)
	}())
}

// Correct the PeerPieces slice length. Return false if the existing slice is
// invalid, such as by receiving badly sized BITFIELD, or invalid HAVE
// messages.
func (cn *connection) setNumPieces(num int) error {
	if cn.peerHasAll {
		return nil
	}
	if cn.PeerPieces == nil {
		return nil
	}
	if len(cn.PeerPieces) == num {
	} else if len(cn.PeerPieces) < num {
		cn.PeerPieces = append(cn.PeerPieces, make([]bool, num-len(cn.PeerPieces))...)
	} else if len(cn.PeerPieces) <= (num+7)/8*8 {
		for _, have := range cn.PeerPieces[num:] {
			if have {
				return errors.New("peer has invalid piece")
			}
		}
		cn.PeerPieces = cn.PeerPieces[:num]
	} else {
		return fmt.Errorf("peer bitfield is excessively long: expected %d, have %d", num, len(cn.PeerPieces))
	}
	if len(cn.PeerPieces) != num {
		panic("wat")
	}
	return nil
}

func eventAgeString(t time.Time) string {
	if t.IsZero() {
		return "never"
	}
	return fmt.Sprintf("%.2fs ago", time.Now().Sub(t).Seconds())
}

func (cn *connection) connectionFlags() (ret string) {
	c := func(b byte) {
		ret += string([]byte{b})
	}
	if cn.encrypted {
		c('E')
	}
	if cn.Discovery != 0 {
		c(byte(cn.Discovery))
	}
	if cn.uTP {
		c('T')
	}
	return
}

// Inspired by https://trac.transmissionbt.com/wiki/PeerStatusText
func (cn *connection) statusFlags() (ret string) {
	c := func(b byte) {
		ret += string([]byte{b})
	}
	if cn.Interested {
		c('i')
	}
	if cn.Choked {
		c('c')
	}
	c('-')
	ret += cn.connectionFlags()
	c('-')
	if cn.PeerInterested {
		c('i')
	}
	if cn.PeerChoked {
		c('c')
	}
	return
}

func (cn *connection) String() string {
	var buf bytes.Buffer
	cn.WriteStatus(&buf, nil)
	return buf.String()
}

func (cn *connection) WriteStatus(w io.Writer, t *torrent) {
	// \t isn't preserved in <pre> blocks?
	fmt.Fprintf(w, "%+q: %s-%s\n", cn.PeerID, cn.localAddr(), cn.remoteAddr())
	fmt.Fprintf(w, "    last msg: %s, connected: %s, last useful chunk: %s\n",
		eventAgeString(cn.lastMessageReceived),
		eventAgeString(cn.completedHandshake),
		eventAgeString(cn.lastUsefulChunkReceived))
	fmt.Fprintf(w,
		"    %s completed, good chunks: %d/%d-%d reqq: %d-%d, flags: %s\n",
		cn.completedString(t),
		cn.UsefulChunksReceived,
		cn.UnwantedChunksReceived+cn.UsefulChunksReceived,
		cn.chunksSent,
		len(cn.Requests),
		len(cn.PeerRequests),
		cn.statusFlags(),
	)
}

func (c *connection) Close() {
	c.mu.Lock()
	defer c.mu.Unlock()
	select {
	case <-c.closing:
		return
	default:
	}
	close(c.closing)
	// TODO: This call blocks sometimes, why?
	go c.conn.Close()
}

func (c *connection) PeerHasPiece(piece int) bool {
	if c.peerHasAll {
		return true
	}
	if piece >= len(c.PeerPieces) {
		return false
	}
	return c.PeerPieces[piece]
}

func (c *connection) Post(msg pp.Message) {
	select {
	case c.post <- msg:
	case <-c.closing:
	}
}

func (c *connection) RequestPending(r request) bool {
	_, ok := c.Requests[r]
	return ok
}

func (c *connection) requestMetadataPiece(index int) {
	eID := c.PeerExtensionIDs["ut_metadata"]
	if eID == 0 {
		return
	}
	if index < len(c.metadataRequests) && c.metadataRequests[index] {
		return
	}
	c.Post(pp.Message{
		Type:       pp.Extended,
		ExtendedID: eID,
		ExtendedPayload: func() []byte {
			b, err := bencode.Marshal(map[string]int{
				"msg_type": pp.RequestMetadataExtensionMsgType,
				"piece":    index,
			})
			if err != nil {
				panic(err)
			}
			return b
		}(),
	})
	for index >= len(c.metadataRequests) {
		c.metadataRequests = append(c.metadataRequests, false)
	}
	c.metadataRequests[index] = true
}

func (c *connection) requestedMetadataPiece(index int) bool {
	return index < len(c.metadataRequests) && c.metadataRequests[index]
}

// Returns true if more requests can be sent.
func (c *connection) Request(chunk request) bool {
	if len(c.Requests) >= c.PeerMaxRequests {
		return false
	}
	if !c.PeerHasPiece(int(chunk.Index)) {
		return true
	}
	if c.RequestPending(chunk) {
		return true
	}
	c.SetInterested(true)
	if c.PeerChoked {
		return false
	}
	if c.Requests == nil {
		c.Requests = make(map[request]struct{}, c.PeerMaxRequests)
	}
	c.Requests[chunk] = struct{}{}
	c.requestsLowWater = len(c.Requests) / 2
	c.Post(pp.Message{
		Type:   pp.Request,
		Index:  chunk.Index,
		Begin:  chunk.Begin,
		Length: chunk.Length,
	})
	return true
}

// Returns true if an unsatisfied request was canceled.
func (c *connection) Cancel(r request) bool {
	if c.Requests == nil {
		return false
	}
	if _, ok := c.Requests[r]; !ok {
		return false
	}
	delete(c.Requests, r)
	c.Post(pp.Message{
		Type:   pp.Cancel,
		Index:  r.Index,
		Begin:  r.Begin,
		Length: r.Length,
	})
	return true
}

// Returns true if an unsatisfied request was canceled.
func (c *connection) PeerCancel(r request) bool {
	if c.PeerRequests == nil {
		return false
	}
	if _, ok := c.PeerRequests[r]; !ok {
		return false
	}
	delete(c.PeerRequests, r)
	return true
}

func (c *connection) Choke() {
	if c.Choked {
		return
	}
	c.Post(pp.Message{
		Type: pp.Choke,
	})
	c.PeerRequests = nil
	c.Choked = true
}

func (c *connection) Unchoke() {
	if !c.Choked {
		return
	}
	c.Post(pp.Message{
		Type: pp.Unchoke,
	})
	c.Choked = false
}

func (c *connection) SetInterested(interested bool) {
	if c.Interested == interested {
		return
	}
	c.Post(pp.Message{
		Type: func() pp.MessageType {
			if interested {
				return pp.Interested
			} else {
				return pp.NotInterested
			}
		}(),
	})
	c.Interested = interested
}

var connectionWriterFlush = expvar.NewInt("connectionWriterFlush")
var connectionWriterWrite = expvar.NewInt("connectionWriterWrite")

// Writes buffers to the socket from the write channel.
func (conn *connection) writer() {
	// Reduce write syscalls.
	buf := bufio.NewWriterSize(conn.rw, 0x8000) // 32 KiB
	// Receives when buf is not empty.
	notEmpty := make(chan struct{}, 1)
	for {
		if buf.Buffered() != 0 {
			// Make sure it's receivable.
			select {
			case notEmpty <- struct{}{}:
			default:
			}
		}
		select {
		case b, ok := <-conn.writeCh:
			if !ok {
				return
			}
			connectionWriterWrite.Add(1)
			_, err := buf.Write(b)
			if err != nil {
				conn.Close()
				return
			}
		case <-conn.closing:
			return
		case <-notEmpty:
			connectionWriterFlush.Add(1)
			err := buf.Flush()
			if err != nil {
				return
			}
		}
	}
}

func (conn *connection) writeOptimizer(keepAliveDelay time.Duration) {
	defer close(conn.writeCh) // Responsible for notifying downstream routines.
	pending := list.New()     // Message queue.
	var nextWrite []byte      // Set to nil if we need to need to marshal the next message.
	timer := time.NewTimer(keepAliveDelay)
	defer timer.Stop()
	lastWrite := time.Now()
	for {
		write := conn.writeCh // Set to nil if there's nothing to write.
		if pending.Len() == 0 {
			write = nil
		} else if nextWrite == nil {
			var err error
			nextWrite, err = pending.Front().Value.(encoding.BinaryMarshaler).MarshalBinary()
			if err != nil {
				panic(err)
			}
		}
	event:
		select {
		case <-timer.C:
			if pending.Len() != 0 {
				break
			}
			keepAliveTime := lastWrite.Add(keepAliveDelay)
			if time.Now().Before(keepAliveTime) {
				timer.Reset(keepAliveTime.Sub(time.Now()))
				break
			}
			pending.PushBack(pp.Message{Keepalive: true})
		case msg, ok := <-conn.post:
			if !ok {
				return
			}
			if msg.Type == pp.Cancel {
				for e := pending.Back(); e != nil; e = e.Prev() {
					elemMsg := e.Value.(pp.Message)
					if elemMsg.Type == pp.Request && msg.Index == elemMsg.Index && msg.Begin == elemMsg.Begin && msg.Length == elemMsg.Length {
						pending.Remove(e)
						optimizedCancels.Add(1)
						break event
					}
				}
			}
			pending.PushBack(msg)
		case write <- nextWrite:
			pending.Remove(pending.Front())
			nextWrite = nil
			lastWrite = time.Now()
			if pending.Len() == 0 {
				timer.Reset(keepAliveDelay)
			}
		case <-conn.closing:
			return
		}
	}
}
<|endoftext|>


<commit_before>package apns

import (
	"crypto/tls"
	"encoding/binary"
	"log"
	"time"
)

//ResponseQueueSize indicates how many APNS responses may be buffered.
var ResponseQueueSize = 10000

//SentBufferSize is the maximum number of sent notifications which may be buffered.
var SentBufferSize = 10000

var maxBackoff = 20 * time.Second

//Connection represents a single connection to APNS.
type Connection struct {
	Client
	conn   *tls.Conn
	queue  chan PushNotification
	errors chan *BadPushNotification
}

//Response is a reply from APNS - see apns.ApplePushResponses.
type Response struct {
	Status     uint8
	Identifier uint32
}

func newResponse() *Response {
	return new(Response)
}

//BadPushNotification represents a notification which APNS didn't like.
type BadPushNotification struct {
	PushNotification
	Status uint8
}

//Enqueue adds a push notification to the end of the "sending" queue.
func (conn *Connection) Enqueue(pn *PushNotification) {
	go func(pn *PushNotification) {
		conn.queue <- *pn
	}(pn)
}

//Errors gives you a channel of the push notifications Apple rejected.
func (conn *Connection) Errors() (errors <-chan *BadPushNotification) {
	return conn.errors
}

//Start initiates a connection to APNS and asnchronously sends notifications which have been queued.
func (conn *Connection) Start() error {
	//Connect to APNS. The reason this is here as well as in sender is that this probably catches any unavoidable errors in a synchronous fashion, while in sender it can reconnect after temporary errors (which should work most of the time.)
	err := conn.connect()
	if err != nil {
		return err
	}
	//Start sender goroutine
	sent := make(chan PushNotification)
	go conn.sender(conn.queue, sent)
	//Start reader goroutine
	responses := make(chan *Response, ResponseQueueSize)
	go conn.reader(responses)
	//Start limbo goroutine
	return nil
}

//Stop gracefully closes the connection - it waits for the sending queue to clear, and then shuts down.
func (conn *Connection) Stop() {
	//We can't just close the main queue channel, because retries might still need to be sent there.
	//
}

func (conn *Connection) sender(queue <-chan PushNotification, sent chan PushNotification) {
	defer conn.conn.Close()
	var backoff = time.Duration(100)
	for {
		pn, ok := <-conn.queue
		if !ok {
			//That means the Connection is stopped
			//close sent?
			return
		} else {
			//If not connected, connect
			if conn.conn == nil {
				for {
					err := conn.connect()
					if err != nil {
						//Exponential backoff up to a limit
						log.Println("APNS: Error connecting to server: ", err)
						backoff = backoff * 2
						if backoff > maxBackoff {
							backoff = maxBackoff
						}
						time.Sleep(backoff)
					} else {
						backoff = 100
						break
					}
				}
			}
			//Then send the push notification
			//TODO(draaglom): Do buffering as per the APNS docs
		}
	}
}

func (conn *Connection) reader(responses chan<- *Response) {
	buffer := make([]byte, 6)
	for {
		_, err := conn.conn.Read(buffer)
		if err != nil {
			log.Println("APNS: Error reading from connection: ", err)
			conn.conn.Close()
			return
		}
		resp := newResponse()
		resp.Identifier = binary.BigEndian.Uint32(buffer[2:6])
		resp.Status = uint8(buffer[1])
		responses <- resp
	}
}

func (conn *Connection) limbo(sent <-chan PushNotification, responses chan Response, errors chan PushNotification, queue chan PushNotification) {
	limbo := make(chan PushNotification, SentBufferSize)
	ticker := time.NewTicker(1 * time.Second)
	timeNextNotification := true
	for {
		select {
		case pn := <-sent:
			//Drop it into the array
			limbo <- pn
			if timeNextNotification {
				//Is there a cleaner way of doing this?
				go func(pn PushNotification) {
					<-time.After(TimeoutSeconds * time.Second)
					successResp := newResponse()
					successResp.Identifier = pn.Identifier
					responses <- *successResp
				}(pn)
				timeNextNotification = false
			}
		case resp, ok := <-responses:
			if !ok {
				//If the responses channel is closed,
				//that means we're shutting down the connection.
			}
			switch {
			case resp.Status == 0:
				//Status 0 is a "success" response generated by a timeout in the library.
				for pn := range limbo {
					//Drop all the notifications until we get to the timed-out one.
					//(and leave the others in limbo)
					if pn.Identifier == resp.Identifier {
						break
					}
				}
			default:
				hit := false
				for pn := range limbo {
					switch {
					case pn.Identifier != resp.Identifier && !hit:
						//We haven't seen the identified notification yet
						//so these are all successful (drop silently)
					case pn.Identifier == resp.Identifier:
						hit = true
						if resp.Status != 10 {
							//It was an error, we should report this on the error channel
						}
					case pn.Identifier != resp.Identifier && hit:
						//We've already seen the identified notification,
						//so these should be requeued
						conn.Enqueue(&pn)
					}
				}
			}
		case <-ticker.C:
			timeNextNotification = true
		}
	}
}

func (conn *Connection) connect() error {
	if conn.conn != nil {
		conn.conn.Close()
	}

	var cert tls.Certificate
	var err error
	if len(conn.CertificateBase64) == 0 && len(conn.KeyBase64) == 0 {
		// The user did not specify raw block contents, so check the filesystem.
		cert, err = tls.LoadX509KeyPair(conn.CertificateFile, conn.KeyFile)
	} else {
		// The user provided the raw block contents, so use that.
		cert, err = tls.X509KeyPair([]byte(conn.CertificateBase64), []byte(conn.KeyBase64))
	}

	if err != nil {
		return err
	}

	conf := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	tlsConn, err := tls.Dial("tcp", conn.Gateway, conf)
	if err != nil {
		return err
	}
	err = tlsConn.Handshake()
	if err != nil {
		_ = tlsConn.Close()
		return err
	}
	conn.conn = tlsConn
	return nil
}
<commit_msg>Report a bad notification on the error channel<commit_after>package apns

import (
	"crypto/tls"
	"encoding/binary"
	"log"
	"time"
)

//ResponseQueueSize indicates how many APNS responses may be buffered.
var ResponseQueueSize = 10000

//SentBufferSize is the maximum number of sent notifications which may be buffered.
var SentBufferSize = 10000

var maxBackoff = 20 * time.Second

//Connection represents a single connection to APNS.
type Connection struct {
	Client
	conn   *tls.Conn
	queue  chan PushNotification
	errors chan *BadPushNotification
}

//Response is a reply from APNS - see apns.ApplePushResponses.
type Response struct {
	Status     uint8
	Identifier uint32
}

func newResponse() *Response {
	return new(Response)
}

//BadPushNotification represents a notification which APNS didn't like.
type BadPushNotification struct {
	PushNotification
	Status uint8
}

//Enqueue adds a push notification to the end of the "sending" queue.
func (conn *Connection) Enqueue(pn *PushNotification) {
	go func(pn *PushNotification) {
		conn.queue <- *pn
	}(pn)
}

//Errors gives you a channel of the push notifications Apple rejected.
func (conn *Connection) Errors() (errors <-chan *BadPushNotification) {
	return conn.errors
}

//Start initiates a connection to APNS and asnchronously sends notifications which have been queued.
func (conn *Connection) Start() error {
	//Connect to APNS. The reason this is here as well as in sender is that this probably catches any unavoidable errors in a synchronous fashion, while in sender it can reconnect after temporary errors (which should work most of the time.)
	err := conn.connect()
	if err != nil {
		return err
	}
	//Start sender goroutine
	sent := make(chan PushNotification)
	go conn.sender(conn.queue, sent)
	//Start reader goroutine
	responses := make(chan *Response, ResponseQueueSize)
	go conn.reader(responses)
	//Start limbo goroutine
	return nil
}

//Stop gracefully closes the connection - it waits for the sending queue to clear, and then shuts down.
func (conn *Connection) Stop() {
	//We can't just close the main queue channel, because retries might still need to be sent there.
	//
}

func (conn *Connection) sender(queue <-chan PushNotification, sent chan PushNotification) {
	defer conn.conn.Close()
	var backoff = time.Duration(100)
	for {
		pn, ok := <-conn.queue
		if !ok {
			//That means the Connection is stopped
			//close sent?
			return
		} else {
			//If not connected, connect
			if conn.conn == nil {
				for {
					err := conn.connect()
					if err != nil {
						//Exponential backoff up to a limit
						log.Println("APNS: Error connecting to server: ", err)
						backoff = backoff * 2
						if backoff > maxBackoff {
							backoff = maxBackoff
						}
						time.Sleep(backoff)
					} else {
						backoff = 100
						break
					}
				}
			}
			//Then send the push notification
			//TODO(draaglom): Do buffering as per the APNS docs
		}
	}
}

func (conn *Connection) reader(responses chan<- *Response) {
	buffer := make([]byte, 6)
	for {
		_, err := conn.conn.Read(buffer)
		if err != nil {
			log.Println("APNS: Error reading from connection: ", err)
			conn.conn.Close()
			return
		}
		resp := newResponse()
		resp.Identifier = binary.BigEndian.Uint32(buffer[2:6])
		resp.Status = uint8(buffer[1])
		responses <- resp
	}
}

func (conn *Connection) limbo(sent <-chan PushNotification, responses chan Response, errors chan BadPushNotification, queue chan PushNotification) {
	limbo := make(chan PushNotification, SentBufferSize)
	ticker := time.NewTicker(1 * time.Second)
	timeNextNotification := true
	for {
		select {
		case pn := <-sent:
			//Drop it into the array
			limbo <- pn
			if timeNextNotification {
				//Is there a cleaner way of doing this?
				go func(pn PushNotification) {
					<-time.After(TimeoutSeconds * time.Second)
					successResp := newResponse()
					successResp.Identifier = pn.Identifier
					responses <- *successResp
				}(pn)
				timeNextNotification = false
			}
		case resp, ok := <-responses:
			if !ok {
				//If the responses channel is closed,
				//that means we're shutting down the connection.
			}
			switch {
			case resp.Status == 0:
				//Status 0 is a "success" response generated by a timeout in the library.
				for pn := range limbo {
					//Drop all the notifications until we get to the timed-out one.
					//(and leave the others in limbo)
					if pn.Identifier == resp.Identifier {
						break
					}
				}
			default:
				hit := false
				for pn := range limbo {
					switch {
					case pn.Identifier != resp.Identifier && !hit:
						//We haven't seen the identified notification yet
						//so these are all successful (drop silently)
					case pn.Identifier == resp.Identifier:
						hit = true
						if resp.Status != 10 {
							//It was an error, we should report this on the error channel
							bad := BadPushNotification{PushNotification: pn, Status: resp.Status}
							go func(bad BadPushNotification) {
								errors <- bad
							}(bad)
						}
					case pn.Identifier != resp.Identifier && hit:
						//We've already seen the identified notification,
						//so these should be requeued
						conn.Enqueue(&pn)
					}
				}
			}
		case <-ticker.C:
			timeNextNotification = true
		}
	}
}

func (conn *Connection) connect() error {
	if conn.conn != nil {
		conn.conn.Close()
	}

	var cert tls.Certificate
	var err error
	if len(conn.CertificateBase64) == 0 && len(conn.KeyBase64) == 0 {
		// The user did not specify raw block contents, so check the filesystem.
		cert, err = tls.LoadX509KeyPair(conn.CertificateFile, conn.KeyFile)
	} else {
		// The user provided the raw block contents, so use that.
		cert, err = tls.X509KeyPair([]byte(conn.CertificateBase64), []byte(conn.KeyBase64))
	}

	if err != nil {
		return err
	}

	conf := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	tlsConn, err := tls.Dial("tcp", conn.Gateway, conf)
	if err != nil {
		return err
	}
	err = tlsConn.Handshake()
	if err != nil {
		_ = tlsConn.Close()
		return err
	}
	conn.conn = tlsConn
	return nil
}
<|endoftext|>


<commit_before>/*
Open Source Initiative OSI - The MIT License (MIT):Licensing

The MIT License (MIT)
Copyright (c) 2013 Ralph Caraveo (deckarep@gmail.com)

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

// Package mapset implements a simple and generic set collection.
// Items stored within it are unordered and unique
// It supports typical set operations: membership testing, intersection, union, difference and symmetric difference

package mapset

import "fmt"
import "strings"

type Set struct {
	set map[interface{}]_placeHolder
}

type _placeHolder struct{}

func NewSet() *Set {
	return &Set{make(map[interface{}]_placeHolder)}
}

func (set *Set) Add(i interface{}) bool {
	_, found := set.set[i]
	set.set[i] = _placeHolder{}
	return !found //False if it existed already
}

func (set *Set) Contains(i interface{}) bool {
	if _, found := set.set[i]; found {
		return found //true if it existed already
	}
	return false
}

func (set *Set) IsSubset(other *Set) bool {
	for key, _ := range set.set {
		if !other.Contains(key) {
			return false
		}
	}
	return true
}

func (set *Set) IsSuperset(other *Set) bool {
	for key, _ := range other.set {
		if !set.Contains(key) {
			return false
		}
	}
	return true
}

func (set *Set) Union(other *Set) *Set {
	if set != nil && other != nil {
		unionedSet := NewSet()

		for key, _ := range set.set {
			unionedSet.Add(key)
		}
		for key, _ := range other.set {
			unionedSet.Add(key)
		}
		return unionedSet
	}
	return nil
}

func (set *Set) Intersect(other *Set) *Set {
	if set != nil && other != nil {
		intersectedSet := NewSet()
		var smallerSet *Set = nil
		var largerSet *Set = nil

		//figure out the smaller of the two sets and loop on that one as an optimization.
		if set.Size() < other.Size() {
			smallerSet = set
			largerSet = other
		} else {
			smallerSet = other
			largerSet = set
		}

		for key, _ := range smallerSet.set {
			if largerSet.Contains(key) {
				intersectedSet.Add(key)
			}
		}
		return intersectedSet
	}
	return nil
}

func (set *Set) Difference(other *Set) *Set {
	if set != nil && other != nil {
		differencedSet := NewSet()

		for key, _ := range set.set {
			if !other.Contains(key) {
				differencedSet.Add(key)
			}
		}

		return differencedSet
	}
	return nil
}

func (set *Set) SymmetricDifference(other *Set) *Set {
	if set != nil && other != nil {
		aDiff := set.Difference(other)
		bDiff := other.Difference(set)

		symDifferencedSet := aDiff.Union(bDiff)

		return symDifferencedSet
	}
	return nil
}

func (set *Set) Clear() {
	set.set = make(map[interface{}]_placeHolder)
}

func (set *Set) Remove(i interface{}) {
	delete(set.set, i)
}

// Size returns the how many items are currently in the set.
func (set *Set) Size() int {
	return len(set.set)
}

// Equal determines if two sets are equal to each other.
// If they both are the same size and have the same items they are considered equal.
// Order of items is not relevent for sets to be equal.
func (set *Set) Equal(other *Set) bool {
	if set != nil && other != nil {
		if !(set.Size() == other.Size()) {
			return false
		} else {
			for key, _ := range set.set {
				if !other.Contains(key) {
					return false
				}
			}
			return true
		}
	}
	return false
}

func (set *Set) String() string {
	items := make([]string, 0, len(set.set))

	for key := range set.set {
		items = append(items, fmt.Sprintf("%v", key))
	}
	return fmt.Sprintf("Set{%s}", strings.Join(items, ", "))
}
<commit_msg>Adding one-liner docs to all methods to help flesh-out the documentation.<commit_after>/*
Open Source Initiative OSI - The MIT License (MIT):Licensing

The MIT License (MIT)
Copyright (c) 2013 Ralph Caraveo (deckarep@gmail.com)

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

// Package mapset implements a simple and generic set collection.
// Items stored within it are unordered and unique
// It supports typical set operations: membership testing, intersection, union, difference and symmetric difference

package mapset

import "fmt"
import "strings"

// The primary type that represents a set
type Set struct {
	set map[interface{}]_placeHolder
}

type _placeHolder struct{}

// Creates and returns a pointer to an empty set.
func NewSet() *Set {
	return &Set{make(map[interface{}]_placeHolder)}
}

// Adds an item to the current set if it doesn't already exist in the set.
func (set *Set) Add(i interface{}) bool {
	_, found := set.set[i]
	set.set[i] = _placeHolder{}
	return !found //False if it existed already
}

// Determines if a given item is already in the set.
func (set *Set) Contains(i interface{}) bool {
	if _, found := set.set[i]; found {
		return found //true if it existed already
	}
	return false
}

// Determines if every item in the other set is in this set.
func (set *Set) IsSubset(other *Set) bool {
	for key, _ := range set.set {
		if !other.Contains(key) {
			return false
		}
	}
	return true
}

// Determines if every item of this set is in the other set.
func (set *Set) IsSuperset(other *Set) bool {
	for key, _ := range other.set {
		if !set.Contains(key) {
			return false
		}
	}
	return true
}

// Returns a new set with all items in both sets.
func (set *Set) Union(other *Set) *Set {
	if set != nil && other != nil {
		unionedSet := NewSet()

		for key, _ := range set.set {
			unionedSet.Add(key)
		}
		for key, _ := range other.set {
			unionedSet.Add(key)
		}
		return unionedSet
	}
	return nil
}

// Returns a new set with items that exist only in both sets.
func (set *Set) Intersect(other *Set) *Set {
	if set != nil && other != nil {
		intersectedSet := NewSet()
		var smallerSet *Set = nil
		var largerSet *Set = nil

		//figure out the smaller of the two sets and loop on that one as an optimization.
		if set.Size() < other.Size() {
			smallerSet = set
			largerSet = other
		} else {
			smallerSet = other
			largerSet = set
		}

		for key, _ := range smallerSet.set {
			if largerSet.Contains(key) {
				intersectedSet.Add(key)
			}
		}
		return intersectedSet
	}
	return nil
}

// Returns a new set with items in the current set but not in the other set
func (set *Set) Difference(other *Set) *Set {
	if set != nil && other != nil {
		differencedSet := NewSet()

		for key, _ := range set.set {
			if !other.Contains(key) {
				differencedSet.Add(key)
			}
		}

		return differencedSet
	}
	return nil
}

// Returns a new set with items in the current set or the other set but not in both.
func (set *Set) SymmetricDifference(other *Set) *Set {
	if set != nil && other != nil {
		aDiff := set.Difference(other)
		bDiff := other.Difference(set)

		symDifferencedSet := aDiff.Union(bDiff)

		return symDifferencedSet
	}
	return nil
}

// Clears the entire set to be the empty set.
func (set *Set) Clear() {
	set.set = make(map[interface{}]_placeHolder)
}

// Allows the removal of a single item in the set.
func (set *Set) Remove(i interface{}) {
	delete(set.set, i)
}

// Size returns the how many items are currently in the set.
func (set *Set) Size() int {
	return len(set.set)
}

// Equal determines if two sets are equal to each other.
// If they both are the same size and have the same items they are considered equal.
// Order of items is not relevent for sets to be equal.
func (set *Set) Equal(other *Set) bool {
	if set != nil && other != nil {
		if !(set.Size() == other.Size()) {
			return false
		} else {
			for key, _ := range set.set {
				if !other.Contains(key) {
					return false
				}
			}
			return true
		}
	}
	return false
}

// Provides a convenient string representation of the current state of the set.
func (set *Set) String() string {
	items := make([]string, 0, len(set.set))

	for key := range set.set {
		items = append(items, fmt.Sprintf("%v", key))
	}
	return fmt.Sprintf("Set{%s}", strings.Join(items, ", "))
}
<|endoftext|>


<commit_before>// Package set provides both threadsafe and non-threadsafe implementations of
// a generic set data structure. In the threadsafe set, safety encompasses all
// operations on one set. Operations on multiple sets are consistent in that
// the elements of each set used was valid at exactly one point in time
// between the start and the end of the operation.
package set

// SetType denotes which type of set is created. ThreadSafe or NonThreadSafe
type SetType int

const (
	ThreadSafe = iota
	NonThreadSafe
)

func (s SetType) String() string {
	switch s {
	case ThreadSafe:
		return "ThreadSafe"
	case NonThreadSafe:
		return "NonThreadSafe"
	}
	return ""
}

// Interface is describing a Set. Sets are an unordered, unique list of values.
type Interface interface {
	Add(items ...interface{})
	Remove(items ...interface{})
	Pop() interface{}
	Has(items ...interface{}) bool
	Size() int
	Clear()
	IsEmpty() bool
	IsEqual(s Interface) bool
	IsSubset(s Interface) bool
	IsSuperset(s Interface) bool
	Each(func(interface{}) bool)
	String() string
	List() []interface{}
	Copy() Interface
	Merge(s Interface)
	Separate(s Interface)
}

// helpful to not write everywhere struct{}{}
var keyExists = struct{}{}

// New creates and initalizes a new Set interface. Its single parameter
// denotes the type of set to create. Either ThreadSafe or
// NonThreadSafe. The default is ThreadSafe.
func New(settype SetType) Interface {
	if settype == NonThreadSafe {
		return newNonTS()
	}
	return newTS()
}

// Union is the merger of multiple sets. It returns a new set with all the
// elements present in all the sets that are passed.
//
// The dynamic type of the returned set is determined by the first passed set's
// implementation of the New() method.
func Union(set1, set2 Interface, sets ...Interface) Interface {
	u := set1.Copy()
	set2.Each(func(item interface{}) bool {
		u.Add(item)
		return true
	})
	for _, set := range sets {
		set.Each(func(item interface{}) bool {
			u.Add(item)
			return true
		})
	}

	return u
}

// Difference returns a new set which contains items which are in in the first
// set but not in the others. Unlike the Difference() method you can use this
// function separately with multiple sets.
func Difference(set1, set2 Interface, sets ...Interface) Interface {
	s := set1.Copy()
	s.Separate(set2)
	for _, set := range sets {
		s.Separate(set) // seperate is thread safe
	}
	return s
}

// Intersection returns a new set which contains items that only exist in all given sets.
func Intersection(set1, set2 Interface, sets ...Interface) Interface {
	all := Union(set1, set2, sets...)
	result := Union(set1, set2, sets...)
	all.Each(func(item interface{}) bool {
		for _, set := range sets {
			if !set.Has(item) || !set1.Has(item) || !set2.Has(item) {
				result.Remove(item)
			}
		}
		return true
	})
	return result
}

// SymmetricDifference returns a new set which s is the difference of items which are in
// one of either, but not in both.
func SymmetricDifference(s Interface, t Interface) Interface {
	u := Difference(s, t)
	v := Difference(t, s)
	return Union(u, v)
}

// StringSlice is a helper function that returns a slice of strings of s. If
// the set contains mixed types of items only items of type string are returned.
func StringSlice(s Interface) []string {
	slice := make([]string, 0)
	for _, item := range s.List() {
		v, ok := item.(string)
		if !ok {
			continue
		}

		slice = append(slice, v)
	}
	return slice
}

// IntSlice is a helper function that returns a slice of ints of s. If
// the set contains mixed types of items only items of type int are returned.
func IntSlice(s Interface) []int {
	slice := make([]int, 0)
	for _, item := range s.List() {
		v, ok := item.(int)
		if !ok {
			continue
		}

		slice = append(slice, v)
	}
	return slice
}
<commit_msg>fix intersection with two sets<commit_after>// Package set provides both threadsafe and non-threadsafe implementations of
// a generic set data structure. In the threadsafe set, safety encompasses all
// operations on one set. Operations on multiple sets are consistent in that
// the elements of each set used was valid at exactly one point in time
// between the start and the end of the operation.
package set

// SetType denotes which type of set is created. ThreadSafe or NonThreadSafe
type SetType int

const (
	ThreadSafe = iota
	NonThreadSafe
)

func (s SetType) String() string {
	switch s {
	case ThreadSafe:
		return "ThreadSafe"
	case NonThreadSafe:
		return "NonThreadSafe"
	}
	return ""
}

// Interface is describing a Set. Sets are an unordered, unique list of values.
type Interface interface {
	Add(items ...interface{})
	Remove(items ...interface{})
	Pop() interface{}
	Has(items ...interface{}) bool
	Size() int
	Clear()
	IsEmpty() bool
	IsEqual(s Interface) bool
	IsSubset(s Interface) bool
	IsSuperset(s Interface) bool
	Each(func(interface{}) bool)
	String() string
	List() []interface{}
	Copy() Interface
	Merge(s Interface)
	Separate(s Interface)
}

// helpful to not write everywhere struct{}{}
var keyExists = struct{}{}

// New creates and initalizes a new Set interface. Its single parameter
// denotes the type of set to create. Either ThreadSafe or
// NonThreadSafe. The default is ThreadSafe.
func New(settype SetType) Interface {
	if settype == NonThreadSafe {
		return newNonTS()
	}
	return newTS()
}

// Union is the merger of multiple sets. It returns a new set with all the
// elements present in all the sets that are passed.
//
// The dynamic type of the returned set is determined by the first passed set's
// implementation of the New() method.
func Union(set1, set2 Interface, sets ...Interface) Interface {
	u := set1.Copy()
	set2.Each(func(item interface{}) bool {
		u.Add(item)
		return true
	})
	for _, set := range sets {
		set.Each(func(item interface{}) bool {
			u.Add(item)
			return true
		})
	}

	return u
}

// Difference returns a new set which contains items which are in in the first
// set but not in the others. Unlike the Difference() method you can use this
// function separately with multiple sets.
func Difference(set1, set2 Interface, sets ...Interface) Interface {
	s := set1.Copy()
	s.Separate(set2)
	for _, set := range sets {
		s.Separate(set) // seperate is thread safe
	}
	return s
}

// Intersection returns a new set which contains items that only exist in all given sets.
func Intersection(set1, set2 Interface, sets ...Interface) Interface {
	all := Union(set1, set2, sets...)
	result := Union(set1, set2, sets...)

	all.Each(func(item interface{}) bool {
		if !set1.Has(item) || !set2.Has(item) {
			result.Remove(item)
		}

		for _, set := range sets {
			if !set.Has(item) {
				result.Remove(item)
			}
		}
		return true
	})
	return result
}

// SymmetricDifference returns a new set which s is the difference of items which are in
// one of either, but not in both.
func SymmetricDifference(s Interface, t Interface) Interface {
	u := Difference(s, t)
	v := Difference(t, s)
	return Union(u, v)
}

// StringSlice is a helper function that returns a slice of strings of s. If
// the set contains mixed types of items only items of type string are returned.
func StringSlice(s Interface) []string {
	slice := make([]string, 0)
	for _, item := range s.List() {
		v, ok := item.(string)
		if !ok {
			continue
		}

		slice = append(slice, v)
	}
	return slice
}

// IntSlice is a helper function that returns a slice of ints of s. If
// the set contains mixed types of items only items of type int are returned.
func IntSlice(s Interface) []int {
	slice := make([]int, 0)
	for _, item := range s.List() {
		v, ok := item.(int)
		if !ok {
			continue
		}

		slice = append(slice, v)
	}
	return slice
}
<|endoftext|>


<commit_before>package main

type mockI2CBus struct {}

func (bus mockI2CBus) ReadByte(addr byte) (value byte, err error) {
    return 1, nil
}

func (bus mockI2CBus) WriteByte(addr, value byte) error {
    return nil
}

func (bus mockI2CBus) Close() error {
    return nil
}
<commit_msg>fixing tests missing ReadBytes<commit_after>package main

type mockI2CBus struct {}

func (bus mockI2CBus) ReadBytes(addr byte, num int) (value []byte, err error) {
    return []byte{1}, nil
}

func (bus mockI2CBus) ReadByte(addr byte) (value byte, err error) {
    return 1, nil
}

func (bus mockI2CBus) WriteByte(addr, value byte) error {
    return nil
}

func (bus mockI2CBus) Close() error {
    return nil
}
<|endoftext|>


<commit_before>package kite

import (
	"errors"
	"fmt"
	"strings"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kitekey"
)

// Request contains information about the incoming request.
type Request struct {
	Method         string
	Args           *dnode.Partial
	LocalKite      *Kite
	Client         *Client
	Username       string
	Authentication *Authentication
}

// Response is the type of the object that is returned from request handlers
// and the type of only argument that is passed to callback functions.
type Response struct {
	Error  *Error      `json:"error" dnode:"-"`
	Result interface{} `json:"result"`
}

// HandlerFunc is the type of the handlers registered to Kite.
// The returned result must be Marshalable with json package.
type HandlerFunc func(*Request) (result interface{}, err error)

// runMethod is called when a method is received from remote Kite.
func (c *Client) runMethod(method string, handlerFunc HandlerFunc, args *dnode.Partial) {
	var (
		// The request that will be constructed from incoming dnode message.
		request *Request

		// First value to the response.
		result interface{}

		// Second value to the response.
		kiteErr *Error

		// Will send the response when called.
		callback dnode.Function
	)

	// Send result if "responseCallback" exists in the request.
	defer func() {
		if callback.Caller == nil {
			return
		}

		// Only argument to the callback.
		response := Response{
			Result: result,
			Error:  kiteErr,
		}

		// Call response callback function.
		if err := callback.Call(response); err != nil {
			c.LocalKite.Log.Error(err.Error())
		}
	}()

	// Recover dnode argument errors. The caller can use functions like
	// MustString(), MustSlice()... without the fear of panic.
	defer c.LocalKite.recoverError(&kiteErr)()

	request, callback = c.newRequest(method, args)

	if !c.LocalKite.Config.DisableAuthentication {
		kiteErr = request.authenticate()
		if kiteErr != nil {
			return
		}
	}

	// Call the handler function.
	var err error
	result, err = handlerFunc(request)

	if err != nil {
		panic(err) // This will be recoverd from kite.recoverError() above.
	}
}

// HandleFunc registers a handler to run when a method call is received from a Kite.
func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.handlers[method] = handler
}

// runCallback is called when a callback method call is received from remote Kite.
func (c *Client) runCallback(callback func(*dnode.Partial), args *dnode.Partial) {
	kiteErr := new(Error)                      // Not used. For recovering the error.
	defer c.LocalKite.recoverError(&kiteErr)() // Do not panic no matter what.

	// Call the callback function.
	callback(args)
}

// newRequest returns a new *Request from the method and arguments passed.
func (c *Client) newRequest(method string, arguments *dnode.Partial) (request *Request, responseCallback dnode.Function) {
	// Parse dnode method arguments: [options]
	var options callOptions
	arguments.One().MustUnmarshal(&options)

	// Notify the handlers registered with Kite.OnFirstRequest().
	if c.RemoteAddr() != "" {
		c.firstRequestHandlersNotified.Do(func() {
			c.Kite = options.Kite
			c.LocalKite.callOnFirstRequestHandlers(c)
		})
	}

	request = &Request{
		Method:         method,
		Args:           options.WithArgs,
		LocalKite:      c.LocalKite,
		Client:         c,
		Authentication: options.Authentication,
	}

	return request, options.ResponseCallback
}

// authenticate tries to authenticate the user by selecting appropriate
// authenticator function.
func (r *Request) authenticate() *Error {
	// Trust the Kite if we have initiated the connection.
	// RemoteAddr() returns "" if this is an outgoing connection.
	if r.Client.RemoteAddr() == "" {
		return nil
	}

	if r.Authentication == nil {
		return &Error{
			Type:    "authenticationError",
			Message: "No authentication information is provided",
		}
	}

	// Select authenticator function.
	f := r.LocalKite.Authenticators[r.Authentication.Type]
	if f == nil {
		return &Error{
			Type:    "authenticationError",
			Message: fmt.Sprintf("Unknown authentication type: %s", r.Authentication.Type),
		}
	}

	// Call authenticator function. It sets the Request.Username field.
	err := f(r)
	if err != nil {
		return &Error{
			Type:    "authenticationError",
			Message: err.Error(),
		}
	}

	// Fix username of the remote Kite if it is invalid.
	// This prevents a Kite to impersonate someone else's Kite.
	r.Client.Kite.Username = r.Username

	return nil
}

// AuthenticateFromToken is the default Authenticator for Kite.
func (k *Kite) AuthenticateFromToken(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, r.LocalKite.RSAKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if audience, ok := token.Claims["aud"].(string); !ok || !strings.HasPrefix(k.Kite().String(), audience) {
		return fmt.Errorf("Invalid audience in token: %s", audience)
	}

	// We don't check for exp and nbf claims here because jwt-go package already checks them.

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}

// AuthenticateFromKiteKey authenticates user from kite key.
func (k *Kite) AuthenticateFromKiteKey(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, kitekey.GetKontrolKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}
<commit_msg>request: do not panic, instead try to normalize<commit_after>package kite

import (
	"errors"
	"fmt"
	"strings"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kitekey"
)

// Request contains information about the incoming request.
type Request struct {
	Method         string
	Args           *dnode.Partial
	LocalKite      *Kite
	Client         *Client
	Username       string
	Authentication *Authentication
}

// Response is the type of the object that is returned from request handlers
// and the type of only argument that is passed to callback functions.
type Response struct {
	Error  *Error      `json:"error" dnode:"-"`
	Result interface{} `json:"result"`
}

// HandlerFunc is the type of the handlers registered to Kite.
// The returned result must be Marshalable with json package.
type HandlerFunc func(*Request) (result interface{}, err error)

// runMethod is called when a method is received from remote Kite.
func (c *Client) runMethod(method string, handlerFunc HandlerFunc, args *dnode.Partial) {
	var (
		// The request that will be constructed from incoming dnode message.
		request *Request

		// First value to the response.
		result interface{}

		// Second value to the response.
		kiteErr *Error

		// Will send the response when called.
		callback dnode.Function
	)

	// Send result if "responseCallback" exists in the request.
	defer func() {
		if callback.Caller == nil {
			return
		}

		// Only argument to the callback.
		response := Response{
			Result: result,
			Error:  kiteErr,
		}

		// Call response callback function.
		if err := callback.Call(response); err != nil {
			c.LocalKite.Log.Error(err.Error())
		}
	}()

	// Recover dnode argument errors. The caller can use functions like
	// MustString(), MustSlice()... without the fear of panic.
	defer c.LocalKite.recoverError(&kiteErr)()

	request, callback = c.newRequest(method, args)

	if !c.LocalKite.Config.DisableAuthentication {
		kiteErr = request.authenticate()
		if kiteErr != nil {
			return
		}
	}

	// Call the handler function.
	var err error
	result, err = handlerFunc(request)
	if err != nil {
		switch t := err.(type) {
		case *Error:
			kiteErr = t
		case *dnode.ArgumentError:
			kiteErr = &Error{
				Type:    "argumentError",
				Message: err.Error(),
			}
		default:
			kiteErr = &Error{
				Type:    "genericError",
				Message: err.Error(),
			}
		}
	}

}

// HandleFunc registers a handler to run when a method call is received from a Kite.
func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.handlers[method] = handler
}

// runCallback is called when a callback method call is received from remote Kite.
func (c *Client) runCallback(callback func(*dnode.Partial), args *dnode.Partial) {
	kiteErr := new(Error)                      // Not used. For recovering the error.
	defer c.LocalKite.recoverError(&kiteErr)() // Do not panic no matter what.

	// Call the callback function.
	callback(args)
}

// newRequest returns a new *Request from the method and arguments passed.
func (c *Client) newRequest(method string, arguments *dnode.Partial) (request *Request, responseCallback dnode.Function) {
	// Parse dnode method arguments: [options]
	var options callOptions
	arguments.One().MustUnmarshal(&options)

	// Notify the handlers registered with Kite.OnFirstRequest().
	if c.RemoteAddr() != "" {
		c.firstRequestHandlersNotified.Do(func() {
			c.Kite = options.Kite
			c.LocalKite.callOnFirstRequestHandlers(c)
		})
	}

	request = &Request{
		Method:         method,
		Args:           options.WithArgs,
		LocalKite:      c.LocalKite,
		Client:         c,
		Authentication: options.Authentication,
	}

	return request, options.ResponseCallback
}

// authenticate tries to authenticate the user by selecting appropriate
// authenticator function.
func (r *Request) authenticate() *Error {
	// Trust the Kite if we have initiated the connection.
	// RemoteAddr() returns "" if this is an outgoing connection.
	if r.Client.RemoteAddr() == "" {
		return nil
	}

	if r.Authentication == nil {
		return &Error{
			Type:    "authenticationError",
			Message: "No authentication information is provided",
		}
	}

	// Select authenticator function.
	f := r.LocalKite.Authenticators[r.Authentication.Type]
	if f == nil {
		return &Error{
			Type:    "authenticationError",
			Message: fmt.Sprintf("Unknown authentication type: %s", r.Authentication.Type),
		}
	}

	// Call authenticator function. It sets the Request.Username field.
	err := f(r)
	if err != nil {
		return &Error{
			Type:    "authenticationError",
			Message: err.Error(),
		}
	}

	// Fix username of the remote Kite if it is invalid.
	// This prevents a Kite to impersonate someone else's Kite.
	r.Client.Kite.Username = r.Username

	return nil
}

// AuthenticateFromToken is the default Authenticator for Kite.
func (k *Kite) AuthenticateFromToken(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, r.LocalKite.RSAKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if audience, ok := token.Claims["aud"].(string); !ok || !strings.HasPrefix(k.Kite().String(), audience) {
		return fmt.Errorf("Invalid audience in token: %s", audience)
	}

	// We don't check for exp and nbf claims here because jwt-go package already checks them.

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}

// AuthenticateFromKiteKey authenticates user from kite key.
func (k *Kite) AuthenticateFromKiteKey(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, kitekey.GetKontrolKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}
<|endoftext|>


<commit_before>package gen

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"

	"github.com/robertkrimen/otto"
	"github.com/spf13/afero"
)

const (
	msgModduleNotFound = "module %s not found"
)

type require struct {
	cache map[string]otto.Value
	paths []string
	fs    afero.Fs
}

type moduleInfo struct {
	id string
}

func (r *require) load(call otto.FunctionCall) otto.Value {
	id, err := call.Argument(0).ToString()
	if err != nil {
		Panic(err)
	}
	newID, err := r.resolve(id)
	if err != nil {
		Panic(err)
	}
	return r.loadFromFile(newID, call.Otto)
}

func (r *require) resolve(id string) (string, error) {
	if id == "" {
		return "", errors.New("empty module name")
	}
	if filepath.IsAbs(id) {
		return id, nil
	}
	if strings.HasPrefix(id, ".") {
		_, err := r.fs.Stat(id)
		if err != nil {
			return "", err
		}
	}
	ext := filepath.Ext(id)
	opts := []string{".js", ".json"}
	for i := 0; i < len(r.paths); i++ {
		fullPath := filepath.Join(r.paths[i], id)
		if ext != "" {
			_, err := r.fs.Stat(fullPath)
			if err != nil {
				return "", err
			}
			return fullPath, nil
		}

		for _, e := range opts {
			_, err := r.fs.Stat(fullPath + e)
			if err != nil {
				return "", err
			}
			return fullPath + e, nil
		}
	}
	return "", fmt.Errorf(msgModduleNotFound, id)
}

func (r *require) loadFromFile(path string, vm *otto.Otto) otto.Value {
	f, err := r.fs.Open(path)
	if err != nil {
		Panic(err.Error())
	}
	defer func() { _ = f.Close() }()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		Panic(err.Error())
	}
	if filepath.Ext(path) == ".json" {
		v, err := vm.Call("JSON.parse", nil, string(data))
		if err != nil {
			Panic(err.Error())
		}
		return v
	}
	return r.loadFromSource(string(data), vm)
}

func (r *require) loadFromSource(source string, vm *otto.Otto) otto.Value {
	source = "(function(module) {var require = module.require;var exports = module.exports;\n" + source + "\n})"

	jsModule, _ := vm.Object(`({exports: {}})`)
	jsModule.Set("require", r.load)
	jsExports, _ := jsModule.Get("exports")

	moduleReturn, err := vm.Call(source, jsExports, jsModule)
	if err != nil {
		Panic(err.Error())
	}
	var moduleValue otto.Value
	if !moduleReturn.IsUndefined() {
		moduleValue = moduleReturn
		jsModule.Set("exports", moduleValue)
	} else {
		moduleValue, _ = jsModule.Get("exports")
	}
	return moduleValue
}
<commit_msg>Fix path resolver<commit_after>package gen

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"

	"github.com/robertkrimen/otto"
	"github.com/spf13/afero"
)

const (
	msgModduleNotFound = "module %s not found"
)

type require struct {
	cache map[string]otto.Value
	paths []string
	fs    afero.Fs
}

type moduleInfo struct {
	id string
}

func (r *require) load(call otto.FunctionCall) otto.Value {
	id, err := call.Argument(0).ToString()
	if err != nil {
		Panic(err)
	}
	newID, err := r.resolve(id)
	if err != nil {
		Panic(err.Error())
	}
	return r.loadFromFile(newID, call.Otto)
}

func (r *require) resolve(id string) (string, error) {
	if id == "" {
		return "", errors.New("empty module name")
	}
	if filepath.IsAbs(id) {
		return id, nil
	}
	if strings.HasPrefix(id, ".") {
		_, err := r.fs.Stat(id)
		if err != nil {
			return "", err
		}
	}
	ext := filepath.Ext(id)
	opts := []string{".js", ".json"}
	for i := 0; i < len(r.paths); i++ {
		fullPath := filepath.Join(r.paths[i], id)
		if ext != "" {
			_, err := r.fs.Stat(fullPath)
			if err != nil {
				if i != len(r.paths)-1 {
					continue
				}
				return "", err
			}
			return fullPath, nil
		}

		for _, e := range opts {
			_, err := r.fs.Stat(fullPath + e)
			if err != nil {
				if i != len(r.paths)-1 {
					continue
				}
				return "", err
			}
			return fullPath + e, nil
		}
	}
	return "", fmt.Errorf(msgModduleNotFound, id)
}

func (r *require) loadFromFile(path string, vm *otto.Otto) otto.Value {
	f, err := r.fs.Open(path)
	if err != nil {
		Panic(err.Error())
	}
	defer func() { _ = f.Close() }()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		Panic(err.Error())
	}
	if filepath.Ext(path) == ".json" {
		v, err := vm.Call("JSON.parse", nil, string(data))
		if err != nil {
			Panic(err.Error())
		}
		return v
	}
	return r.loadFromSource(string(data), vm)
}

func (r *require) loadFromSource(source string, vm *otto.Otto) otto.Value {
	source = "(function(module) {var require = module.require;var exports = module.exports;\n" + source + "\n})"

	jsModule, _ := vm.Object(`({exports: {}})`)
	jsModule.Set("require", r.load)
	jsExports, _ := jsModule.Get("exports")

	moduleReturn, err := vm.Call(source, jsExports, jsModule)
	if err != nil {
		Panic(err.Error())
	}
	var moduleValue otto.Value
	if !moduleReturn.IsUndefined() {
		moduleValue = moduleReturn
		jsModule.Set("exports", moduleValue)
	} else {
		moduleValue, _ = jsModule.Get("exports")
	}
	return moduleValue
}
<|endoftext|>


<commit_before>package respond

import (
	"bytes"
	"encoding/json"
	"html/template"
	"net/http"
)

/*
ByteSlice send a byte slice as a response.
*/
func ByteSlice(w http.ResponseWriter, message []byte, statusCode int) {
	w.WriteHeader(statusCode)
	w.Write(message)
}

/*
String sends a string as your response.
*/
func String(w http.ResponseWriter, message string, statusCode int) {
	ByteSlice(w, []byte(message), statusCode)
}

/*
Simple calls String, populating message with the equivalent reason
phrase as returned by http.StatusText, based on the
statusCode.

For example, Status(w, http.StatusBadRequest) is the equivalent
of String(w, "Bad Request", http.StatusBadRequest)

If the statusCode is not recognized an empty string is used.
*/
func Simple(w http.ResponseWriter, statusCode int) {
	String(w, http.StatusText(statusCode), statusCode)
}

/*
JSON sends a json-serializable object as the response and sets
"Content-Type" header to "application/json". If there is an error
the error is returned and the Content-Type is not set, the
http.ResponseWriter object is not written to or modified.
*/
func JSON(w http.ResponseWriter, message interface{}, statusCode int) error {
	jsonAsByteSlice, err := json.Marshal(message)

	if err != nil {
		return err
	}

	w.Header().Add("Content-Type", "application/json")
	ByteSlice(w, jsonAsByteSlice, statusCode)
	return nil
}

/*
HTML sends an html string as the response and sets
"Content-Type" header to "text/html".
*/
func HTML(w http.ResponseWriter, message string, statusCode int) {
	w.Header().Add("Content-Type", "text/html")
	String(w, message, statusCode)
}

/*
HTMLTemplate renders an html template and sends an html string as the response. Will return an error
if the template was not found or otherwise had a problem. In the case of an error,
the http.ResponseWriter object will not be written to or modified and the developer
should take appropriate action, such as respond.Simple(w, http.StatusInternalServerError)
*/
func HTMLTemplate(w http.ResponseWriter, templateName string, data interface{}, statusCode int) error {
	t, err := template.ParseFiles(templateName)
	if err != nil {
		return err
	}

	var b bytes.Buffer
	if err := t.Execute(&b, data); err != nil {
		return err
	}

	HTML(w, b.String(), statusCode)
	return nil
}
<commit_msg>Add some nicer formatting to godoc comments<commit_after>package respond

import (
	"bytes"
	"encoding/json"
	"html/template"
	"net/http"
)

/*
ByteSlice send a byte slice as a response.
*/
func ByteSlice(w http.ResponseWriter, message []byte, statusCode int) {
	w.WriteHeader(statusCode)
	w.Write(message)
}

/*
String sends a string as your response.
*/
func String(w http.ResponseWriter, message string, statusCode int) {
	ByteSlice(w, []byte(message), statusCode)
}

/*
Simple calls String, populating message with the equivalent reason
phrase as returned by http.StatusText, based on the
statusCode.

For example, Status(w, http.StatusBadRequest) is the equivalent
of String(w, "Bad Request", http.StatusBadRequest)

If the statusCode is not recognized by http.StatusText an empty string is used.
*/
func Simple(w http.ResponseWriter, statusCode int) {
	String(w, http.StatusText(statusCode), statusCode)
}

/*
JSON sends a json-serializable object as the response and sets
"Content-Type" header to "application/json".

If there is an error the error is returned and the Content-Type is not set, the
http.ResponseWriter object is not written to or modified.
*/
func JSON(w http.ResponseWriter, message interface{}, statusCode int) error {
	jsonAsByteSlice, err := json.Marshal(message)

	if err != nil {
		return err
	}

	w.Header().Add("Content-Type", "application/json")
	ByteSlice(w, jsonAsByteSlice, statusCode)
	return nil
}

/*
HTML sends an html string as the response and sets
"Content-Type" header to "text/html".
*/
func HTML(w http.ResponseWriter, message string, statusCode int) {
	w.Header().Add("Content-Type", "text/html")
	String(w, message, statusCode)
}

/*
HTMLTemplate renders an html template and sends an html string as the response.

It will return an error if the template was not found or otherwise had a problem.

In the case of an error, the http.ResponseWriter object will not be written to or
modified. The developer should take appropriate action, such as
respond.Simple(w, http.StatusInternalServerError)
*/
func HTMLTemplate(w http.ResponseWriter, templateName string, data interface{}, statusCode int) error {
	t, err := template.ParseFiles(templateName)
	if err != nil {
		return err
	}

	var b bytes.Buffer
	if err := t.Execute(&b, data); err != nil {
		return err
	}

	HTML(w, b.String(), statusCode)
	return nil
}
<|endoftext|>


<commit_before>package worker

import (
	"sphere2cubeGo/cache"
	"log"
	"image"
	"os"
	"image/jpeg"
	"image/color"
	"io"
	"math"
)

const (
	TileUp    = "up"
	TileDown  = "down"
	TileFront = "front"
	TileRight = "right"
	TileBack  = "back"
	TileLeft  = "left"
)

type TileResult struct {
	Tile  Tile
	Image image.Image
}

type Tile struct {
	TileSize int
	TileName string
}

func (tile *Tile) getHalf() float64 {
	return float64(tile.TileSize-1) / 2
}

// Pixel struct
type Pixel struct {
	R int
	G int
	B int
	A int
}

func (pixel *Pixel) pixelToRGBA() color.Color {
	return color.RGBA64{uint16(pixel.R * 257), uint16(pixel.G * 257), uint16(pixel.B * 257), uint16(pixel.A * 257)}
}

// img.At(x, y).RGBA() returns four uint32 values; we want a Pixel
func rgbaToPixel(r uint32, g uint32, b uint32, a uint32) Pixel {
	return Pixel{int(r / 257), int(g / 257), int(b / 257), int(a / 257)}
}

// Get the bi-dimensional pixel array
func getPixels(file io.Reader) ([][]Pixel, error) {
	img, err := jpeg.Decode(file)

	if err != nil {
		return nil, err
	}

	bounds := img.Bounds()
	width, height := bounds.Max.X, bounds.Max.Y

	var pixels [][]Pixel
	for y := 0; y < height; y++ {
		var row []Pixel
		for x := 0; x < width; x++ {
			row = append(row, rgbaToPixel(img.At(x, y).RGBA()))
		}
		pixels = append(pixels, row)
	}

	return pixels, nil
}

func updatePhi(half_size float64, phi float64, major_dir int, minor_dir int, major_m float64, major_p float64, minor_m float64, minor_p float64) float64 {

	if float64(major_dir) < half_size {
		phi = phi + major_m
	} else if float64(major_dir) > half_size {
		phi = phi + major_p
	} else if float64(minor_dir) < half_size {
		phi = minor_m
	} else {
		phi = minor_p
	}

	return phi
}

func phi2Width(width int, phi float64) int {
	x := 0.5 * float64(width) * (phi/math.Pi + 1)

	if x < 1 {
		x += float64(width)
	} else if x > float64(width) {
		x -= float64(width)
	}

	return int(x)
}

func theta2Height(height int, theta float64) int {
	return int(float64(height) * theta / math.Pi)
}

func processCords(tileX int, tileY int, originalImage [][]Pixel, tile Tile, mathCache cache.CacheAngles) Pixel {

	theta := 0.0
	phi := 0.0

	sphereHeight, sphereWidth := len(originalImage), len(originalImage[0])

	if tile.TileName == TileUp {
		theta = mathCache.ZP[tileY][tileX]
		phi = mathCache.PHI[tileX][tileY]
		phi = updatePhi(tile.getHalf(), phi, tileY, tileX, math.Pi, 0, -math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileDown {
		theta = mathCache.ZM[tileY][tileX]
		phi = mathCache.PHI[tileX][tile.TileSize-tileY-1]
		phi = updatePhi(tile.getHalf(), phi, tileY, tileX, 0, math.Pi, -math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileFront {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tileX][tile.TileSize-1] //tile_x, tile_size - 1
		phi = updatePhi(tile.getHalf(), phi, tileY, tileX, 0, 0, -math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileRight {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tile.TileSize-1][tile.TileSize-tileX-1]
		phi = updatePhi(tile.getHalf(), phi, tileX, tileY, 0, math.Pi, math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileBack {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tileX][tile.TileSize-1] + math.Pi
	} else if tile.TileName == TileLeft {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tile.TileSize-1][tile.TileSize-tileX-1]
		phi = updatePhi(tile.getHalf(), phi, tileX, tileY, math.Pi, 0, -math.Pi/2, -math.Pi/2)
	}

	spX := phi2Width(sphereWidth, phi)
	spY := theta2Height(sphereHeight, theta)

	return originalImage[spY][spX]
}

func Worker(tile Tile, mathCache cache.CacheAngles, originalImagePath string, done chan TileResult) {
	log.Printf("Process for tileImage %v --> started", tile.TileName)
	tileImage := image.NewRGBA(image.Rect(0, 0, tile.TileSize, tile.TileSize))
	reader, err := os.Open(originalImagePath)

	if err != nil {
		panic(err)
	}

	defer reader.Close()

	originalPixels, err := getPixels(reader)

	if err != nil {
		panic(err)
	}

	for tileY := 0; tileY < tile.TileSize; tileY++ {
		for tileX := 0; tileX < tile.TileSize; tileX++ {
			pixelToMove := processCords(tileX, tileY, originalPixels, tile, mathCache)
			colorPixel := pixelToMove.pixelToRGBA()
			tileImage.Set(tileX, tileY, colorPixel)
		}
	}

	result := TileResult{Tile: tile, Image: tileImage}

	done <- result

}
<commit_msg>some refactoring<commit_after>package worker

import (
	"sphere2cubeGo/cache"
	"log"
	"image"
	"os"
	"image/jpeg"
	"image/color"
	"io"
	"math"
)

const (
	TileUp    = "up"
	TileDown  = "down"
	TileFront = "front"
	TileRight = "right"
	TileBack  = "back"
	TileLeft  = "left"
)

type TileResult struct {
	Tile  Tile
	Image image.Image
}

type Tile struct {
	TileSize int
	TileName string
}

func (tile *Tile) getHalfSize() float64 {
	return float64(tile.TileSize-1) / 2
}

// Pixel struct
type Pixel struct {
	R int
	G int
	B int
	A int
}

func (pixel *Pixel) pixelToRGBA() color.Color {
	return color.RGBA64{uint16(pixel.R * 257), uint16(pixel.G * 257), uint16(pixel.B * 257), uint16(pixel.A * 257)}
}

// img.At(x, y).RGBA() returns four uint32 values; we want a Pixel
func rgbaToPixel(r uint32, g uint32, b uint32, a uint32) Pixel {
	return Pixel{int(r / 257), int(g / 257), int(b / 257), int(a / 257)}
}

// Get the bi-dimensional pixel array
func getPixels(file io.Reader) ([][]Pixel, error) {
	img, err := jpeg.Decode(file)

	if err != nil {
		return nil, err
	}

	bounds := img.Bounds()
	width, height := bounds.Max.X, bounds.Max.Y

	var pixels [][]Pixel
	for y := 0; y < height; y++ {
		var row []Pixel
		for x := 0; x < width; x++ {
			row = append(row, rgbaToPixel(img.At(x, y).RGBA()))
		}
		pixels = append(pixels, row)
	}

	return pixels, nil
}

func updatePhi(half_size float64, phi float64, major_dir int, minor_dir int, major_m float64, major_p float64, minor_m float64, minor_p float64) float64 {

	if float64(major_dir) < half_size {
		phi = phi + major_m
	} else if float64(major_dir) > half_size {
		phi = phi + major_p
	} else if float64(minor_dir) < half_size {
		phi = minor_m
	} else {
		phi = minor_p
	}

	return phi
}

func phi2Width(width int, phi float64) int {
	x := 0.5 * float64(width) * (phi/math.Pi + 1)

	if x < 1 {
		x += float64(width)
	} else if x > float64(width) {
		x -= float64(width)
	}

	return int(x)
}

func theta2Height(height int, theta float64) int {
	return int(float64(height) * theta / math.Pi)
}

func processCords(tileX int, tileY int, originalImage [][]Pixel, tile Tile, mathCache cache.CacheAngles) Pixel {

	theta := 0.0
	phi := 0.0

	sphereHeight, sphereWidth := len(originalImage), len(originalImage[0])

	if tile.TileName == TileUp {
		theta = mathCache.ZP[tileY][tileX]
		phi = mathCache.PHI[tileX][tileY]
		phi = updatePhi(tile.getHalfSize(), phi, tileY, tileX, math.Pi, 0, -math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileDown {
		theta = mathCache.ZM[tileY][tileX]
		phi = mathCache.PHI[tileX][tile.TileSize-tileY-1]
		phi = updatePhi(tile.getHalfSize(), phi, tileY, tileX, 0, math.Pi, -math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileFront {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tileX][tile.TileSize-1] //tile_x, tile_size - 1
		phi = updatePhi(tile.getHalfSize(), phi, tileY, tileX, 0, 0, -math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileRight {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tile.TileSize-1][tile.TileSize-tileX-1]
		phi = updatePhi(tile.getHalfSize(), phi, tileX, tileY, 0, math.Pi, math.Pi/2, math.Pi/2)
	} else if tile.TileName == TileBack {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tileX][tile.TileSize-1] + math.Pi
	} else if tile.TileName == TileLeft {
		theta = mathCache.XYPM[tile.TileSize-tileY-1][tile.TileSize-tileX-1]
		phi = mathCache.PHI[tile.TileSize-1][tile.TileSize-tileX-1]
		phi = updatePhi(tile.getHalfSize(), phi, tileX, tileY, math.Pi, 0, -math.Pi/2, -math.Pi/2)
	}

	spX := phi2Width(sphereWidth, phi)
	spY := theta2Height(sphereHeight, theta)

	return originalImage[spY][spX]
}

func Worker(tile Tile, mathCache cache.CacheAngles, originalImagePath string, done chan TileResult) {
	log.Printf("Process for tileImage %v --> started", tile.TileName)
	tileImage := image.NewRGBA(image.Rect(0, 0, tile.TileSize, tile.TileSize))
	reader, err := os.Open(originalImagePath)

	if err != nil {
		panic(err)
	}

	defer reader.Close()

	originalPixels, err := getPixels(reader)

	if err != nil {
		panic(err)
	}

	for tileY := 0; tileY < tile.TileSize; tileY++ {
		for tileX := 0; tileX < tile.TileSize; tileX++ {
			pixelToMove := processCords(tileX, tileY, originalPixels, tile, mathCache)
			colorPixel := pixelToMove.pixelToRGBA()
			tileImage.Set(tileX, tileY, colorPixel)
		}
	}

	result := TileResult{Tile: tile, Image: tileImage}

	done <- result

}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"container/list"
	"os"
	es "github.com/jensrantil/rewindd/eventstore"
	zmq "github.com/alecthomas/gozmq"
)

var (
	eventStorePath = flag.String("datadir", "data",
	"directory path where incoming events will be stored. Created"+
	" if non-existent.")
	commandSocketZPath = flag.String("commandsocket",
	"tcp://127.0.0.1:9002", "Command socket. Handles new eventsand"+
	" queries.")
	eventPublishZPath = flag.String("evpubsocket",
	"tcp://127.0.0.1:9003", "ZeroMQ event publishing socket.")
)

// Runs the server that distributes requests to workers.
// Panics on error since it is an essential piece of code required to
// run the application correctly.
func runServer(estore *es.EventStore) {
	context, err := zmq.NewContext()
	if err != nil {
		panic(err)
	}
	defer context.Close()

	commandsock, err := context.NewSocket(zmq.ROUTER)
	if err != nil {
		panic(err)
	}
	defer commandsock.Close()
	err = commandsock.Bind(*commandSocketZPath)
	if err != nil {
		panic(err)
	}

	evpubsock, err := context.NewSocket(zmq.PUB)
	if err != nil {
		panic(err)
	}
	defer evpubsock.Close()
	if binderr := evpubsock.Bind(*commandSocketZPath); binderr != nil {
		panic(binderr)
	}

	loopServer(estore, evpubsock, commandsock)
}

// The result of an asynchronous zmq.Poll call.
type zmqPollResult struct {
	nbrOfChanges int
	err error
}

// Polls a bunch of ZeroMQ sockets and notifies the result through a
// channel. This makes it possible to combine ZeroMQ polling with Go's
// own built-in channels.
func asyncPoll(notifier chan zmqPollResult, items zmq.PollItems) {
	a, b := zmq.Poll(items, -1)
	notifier <- zmqPollResult{a, b}
}

// The core ZeroMQ messaging loop. Handles requests and responses
// asynchronously using the router socket. Every request is delegated to
// a goroutine for maximum concurrency.
//
// `gozmq` does currently not support copy-free messages/frames. This means that
// every message passing through this function needs to be copied
// in-memory. If this becomes a bottleneck in the future, multiple
// router sockets can be hooked to this final router to scale message
// copying to multiple cores.
func loopServer(estore *es.EventStore, evpubsock, frontend zmq.Socket) {
	toPoll := zmq.PollItems{
		zmq.PollItem{Socket: frontend, zmq.Events: zmq.POLLIN},
	}

	pubchan := make(chan es.StoredEvent)
	estore.RegisterPublishedEventsChannel(pubchan)
	go publishAllSavedEvents(pubchan, evpubsock)

	pollchan := make(chan zmqPollResult)
	respchan := make(chan [][]byte)
	go asyncPoll(pollchan, toPoll)
	for {
		select {
		case <- pollchan:
			if toPoll[0].REvents&zmq.POLLIN != 0 {
				msg, _ := toPoll[0].Socket.RecvMultipart(0)
				go handleRequest(respchan, estore, msg)
			}
			go asyncPoll(pollchan, toPoll)
		case frames := <-respchan:
			if err := frontend.SendMultipart(frames, 0); err != nil {
				fmt.Println(err)
			}
		}
	}
}

// Publishes stored events to event listeners.
//
// Pops previously stored messages off a channel and published them to a
// ZeroMQ socket.
func publishAllSavedEvents(toPublish chan es.StoredEvent, evpub zmq.Socket) {
	msg := make([][]byte, 3)
	for {
		event := <-toPublish

		msg[0] = event.Stream
		msg[1] = event.Id
		msg[2] = event.Data

		if err := evpub.SendMultipart(msg, 0); err != nil {
			// TODO: Use logger
			fmt.Println(err)
		}
	}
}

// Handles a single ZeroMQ RES/REQ loop.
//
// The full request message stored in `msg` and the full ZeroMQ response
// is pushed to `respchan`. The function does not return any error
// because it is expected to be called asynchronously as a goroutine.
func handleRequest(respchan chan [][]byte, estore *es.EventStore, msg [][]byte) {

	// TODO: Rename to 'framelist'
	// TODO: Avoid this creation and move it into copyList.
	parts := list.New()
	for _, msgpart := range msg {
		parts.PushBack(msgpart)
	}

	// TODO: Possibly wrap ZeroMQ router frames into a Type before
	// calling this method. That would yield a nicer API without
	// nitty gritty ZeroMQ details.
	resptemplate := list.New()
	emptyFrame := []byte("")
	for true {
		resptemplate.PushBack(parts.Remove(parts.Front()))

		if bytes.Equal(parts.Front().Value.([]byte), emptyFrame) {
			break
		}
	}

	if parts.Len() == 0 {
		errstr := "Incoming command was empty. Ignoring it."
		// TODO: Migrate to logging system
		fmt.Println(errstr)
		response := copyList(resptemplate)
		response.PushBack([]byte("ERROR " + errstr))
		respchan <- listToFrames(response)
		return
	}

	command := string(parts.Front().Value.([]byte))
	switch command {
	case "PUBLISH":
		parts.Remove(parts.Front())
		if parts.Len() != 2 {
			errstr := "Wrong number of frames for PUBLISH."
			// TODO: Migrate to logging system
			fmt.Println(errstr)
			response := copyList(resptemplate)
			response.PushBack([]byte("ERROR " + errstr))
			respchan <- listToFrames(response)
		} else {
			estream := parts.Remove(parts.Front())
			data := parts.Remove(parts.Front())
			newevent := es.UnstoredEvent{
				Stream: estream.([]byte),
				Data: data.([]byte),
			}
			err := estore.Add(newevent)
			if err != nil {
				// TODO: Migrate to logging system
				fmt.Println(err.String())

				response := copyList(resptemplate)
				response.PushBack([]byte("ERROR " + err))
				respchan <- listToFrames(response)
			} else {
				// the event was added
				response := copyList(resptemplate)
				response.PushBack([]byte("PUBLISHED"))
				respchan <- listToFrames(response)
			}
		}
	case "QUERY":
		parts.Remove(parts.Front())
		if parts.Len() != 3 {
			errstr := "Wrong number of frames for QUERY."
			// TODO: Migrate to logging system
			fmt.Println(errstr)
			response := copyList(resptemplate)
			response.PushBack([]byte("ERROR " + errstr))
			respchan <- listToFrames(response)
		} else {
			estreamprefix := parts.Remove(parts.Front())
			fromid := parts.Remove(parts.Front())
			toid := parts.Remove(parts.Front())

			events := make(chan es.StoredEvent)
			req := es.QueryRequest{
				StreamPrefix: estreamprefix.([]byte),
				FromId: fromid.([]byte),
				ToId: toid.([]byte),
			}
			go estore.Query(req, events)
			for eventdata := range(events) {
				response := copyList(resptemplate)
				response.PushBack(eventdata.Stream)
				response.PushBack(eventdata.Id)
				response.PushBack(eventdata.Data)
				// TODO: Prepend the router
				// frames before!
				respchan <- listToFrames(response)
			}
			response := copyList(resptemplate)
			response.PushBack([]byte("END"))
			respchan <- listToFrames(response)
		}
	default:
		// TODO: Move these error strings out as constants of
		//       this package.

		// TODO: Move the chunk of code below into a separate
		// function and reuse for similar piece of code above.
		errstr := "Unknown request type."
		// TODO: Migrate to logging system
		fmt.Println(errstr)
		response := copyList(resptemplate)
		response.PushBack([]byte("ERROR " + errstr))
		respchan <- listToFrames(response)
	}
}

// Convert a doubly linked list of message frames to a slice of message
// frames.
func listToFrames(l *list.List) [][]byte {
	frames := make([][]byte, l.Len())
	i := 0
	for e := l.Front(); e != nil; e = e.Next() {
		frames[i] = e.Value.([]byte)
	}
	return frames
}

// Helper function for copying a doubly linked list.
func copyList(l *list.List) *list.List {
	replica := list.New()
	replica.PushBackList(l)
	return replica
}

// Main method. Will panic if things are so bad that the application
// will not start.
func main() {
	flag.Parse()

	desc, err := es.NewFileSystemDescriptor(*eventStorePath)
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not create event store Desc")
		panic(err)
	}
	estore, err := es.NewEventStore(desc)
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not create event store")
		panic(err)
	}
	defer estore.Close()

	runServer(estore)

	// TODO: Handle SIGINT correctly and smoothly.
}
<commit_msg>Making things compile<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"container/list"
	"os"
	es "github.com/jensrantil/rewindd/eventstore"
	zmq "github.com/alecthomas/gozmq"
)

var (
	eventStorePath = flag.String("datadir", "data",
	"directory path where incoming events will be stored. Created"+
	" if non-existent.")
	commandSocketZPath = flag.String("commandsocket",
	"tcp://127.0.0.1:9002", "Command socket. Handles new eventsand"+
	" queries.")
	eventPublishZPath = flag.String("evpubsocket",
	"tcp://127.0.0.1:9003", "ZeroMQ event publishing socket.")
)

// Runs the server that distributes requests to workers.
// Panics on error since it is an essential piece of code required to
// run the application correctly.
func runServer(estore *es.EventStore) {
	context, err := zmq.NewContext()
	if err != nil {
		panic(err)
	}
	defer context.Close()

	commandsock, err := context.NewSocket(zmq.ROUTER)
	if err != nil {
		panic(err)
	}
	defer commandsock.Close()
	err = commandsock.Bind(*commandSocketZPath)
	if err != nil {
		panic(err)
	}

	evpubsock, err := context.NewSocket(zmq.PUB)
	if err != nil {
		panic(err)
	}
	defer evpubsock.Close()
	if binderr := evpubsock.Bind(*commandSocketZPath); binderr != nil {
		panic(binderr)
	}

	loopServer(estore, evpubsock, commandsock)
}

// The result of an asynchronous zmq.Poll call.
type zmqPollResult struct {
	nbrOfChanges int
	err error
}

// Polls a bunch of ZeroMQ sockets and notifies the result through a
// channel. This makes it possible to combine ZeroMQ polling with Go's
// own built-in channels.
func asyncPoll(notifier chan zmqPollResult, items zmq.PollItems) {
	a, b := zmq.Poll(items, -1)
	notifier <- zmqPollResult{a, b}
}

// The core ZeroMQ messaging loop. Handles requests and responses
// asynchronously using the router socket. Every request is delegated to
// a goroutine for maximum concurrency.
//
// `gozmq` does currently not support copy-free messages/frames. This means that
// every message passing through this function needs to be copied
// in-memory. If this becomes a bottleneck in the future, multiple
// router sockets can be hooked to this final router to scale message
// copying to multiple cores.
func loopServer(estore *es.EventStore, evpubsock, frontend zmq.Socket) {
	toPoll := zmq.PollItems{
		zmq.PollItem{Socket: frontend, zmq.Events: zmq.POLLIN},
	}

	pubchan := make(chan es.StoredEvent)
	estore.RegisterPublishedEventsChannel(pubchan)
	go publishAllSavedEvents(pubchan, evpubsock)

	pollchan := make(chan zmqPollResult)
	respchan := make(chan [][]byte)
	go asyncPoll(pollchan, toPoll)
	for {
		select {
		case <- pollchan:
			if toPoll[0].REvents&zmq.POLLIN != 0 {
				msg, _ := toPoll[0].Socket.RecvMultipart(0)
				go handleRequest(respchan, estore, msg)
			}
			go asyncPoll(pollchan, toPoll)
		case frames := <-respchan:
			if err := frontend.SendMultipart(frames, 0); err != nil {
				fmt.Println(err)
			}
		}
	}
}

// Publishes stored events to event listeners.
//
// Pops previously stored messages off a channel and published them to a
// ZeroMQ socket.
func publishAllSavedEvents(toPublish chan es.StoredEvent, evpub zmq.Socket) {
	msg := make([][]byte, 3)
	for {
		event := <-toPublish

		msg[0] = event.Stream
		msg[1] = event.Id
		msg[2] = event.Data

		if err := evpub.SendMultipart(msg, 0); err != nil {
			// TODO: Use logger
			fmt.Println(err)
		}
	}
}

// Handles a single ZeroMQ RES/REQ loop.
//
// The full request message stored in `msg` and the full ZeroMQ response
// is pushed to `respchan`. The function does not return any error
// because it is expected to be called asynchronously as a goroutine.
func handleRequest(respchan chan [][]byte, estore *es.EventStore, msg [][]byte) {

	// TODO: Rename to 'framelist'
	// TODO: Avoid this creation and move it into copyList.
	parts := list.New()
	for _, msgpart := range msg {
		parts.PushBack(msgpart)
	}

	// TODO: Possibly wrap ZeroMQ router frames into a Type before
	// calling this method. That would yield a nicer API without
	// nitty gritty ZeroMQ details.
	resptemplate := list.New()
	emptyFrame := []byte("")
	for true {
		resptemplate.PushBack(parts.Remove(parts.Front()))

		if bytes.Equal(parts.Front().Value.([]byte), emptyFrame) {
			break
		}
	}

	if parts.Len() == 0 {
		errstr := "Incoming command was empty. Ignoring it."
		// TODO: Migrate to logging system
		fmt.Println(errstr)
		response := copyList(resptemplate)
		response.PushBack([]byte("ERROR " + errstr))
		respchan <- listToFrames(response)
		return
	}

	command := string(parts.Front().Value.([]byte))
	switch command {
	case "PUBLISH":
		parts.Remove(parts.Front())
		if parts.Len() != 2 {
			errstr := "Wrong number of frames for PUBLISH."
			// TODO: Migrate to logging system
			fmt.Println(errstr)
			response := copyList(resptemplate)
			response.PushBack([]byte("ERROR " + errstr))
			respchan <- listToFrames(response)
		} else {
			estream := parts.Remove(parts.Front())
			data := parts.Remove(parts.Front())
			newevent := es.UnstoredEvent{
				Stream: estream.([]byte),
				Data: data.([]byte),
			}
			newId, err := estore.Add(newevent)
			if err != nil {
				// TODO: Migrate to logging system
				sErr := err.Error()
				fmt.Println(sErr)

				response := copyList(resptemplate)
				response.PushBack([]byte("ERROR " + sErr))
				respchan <- listToFrames(response)
			} else {
				// the event was added
				response := copyList(resptemplate)
				response.PushBack([]byte("PUBLISHED"))
				response.PushBack([]byte(newId))
				respchan <- listToFrames(response)
			}
		}
	case "QUERY":
		parts.Remove(parts.Front())
		if parts.Len() != 3 {
			errstr := "Wrong number of frames for QUERY."
			// TODO: Migrate to logging system
			fmt.Println(errstr)
			response := copyList(resptemplate)
			response.PushBack([]byte("ERROR " + errstr))
			respchan <- listToFrames(response)
		} else {
			estreamprefix := parts.Remove(parts.Front())
			fromid := parts.Remove(parts.Front())
			toid := parts.Remove(parts.Front())

			events := make(chan es.StoredEvent)
			req := es.QueryRequest{
				StreamPrefix: estreamprefix.([]byte),
				FromId: fromid.([]byte),
				ToId: toid.([]byte),
			}
			go estore.Query(req, events)
			for eventdata := range(events) {
				response := copyList(resptemplate)
				response.PushBack(eventdata.Stream)
				response.PushBack(eventdata.Id)
				response.PushBack(eventdata.Data)
				// TODO: Prepend the router
				// frames before!
				respchan <- listToFrames(response)
			}
			response := copyList(resptemplate)
			response.PushBack([]byte("END"))
			respchan <- listToFrames(response)
		}
	default:
		// TODO: Move these error strings out as constants of
		//       this package.

		// TODO: Move the chunk of code below into a separate
		// function and reuse for similar piece of code above.
		errstr := "Unknown request type."
		// TODO: Migrate to logging system
		fmt.Println(errstr)
		response := copyList(resptemplate)
		response.PushBack([]byte("ERROR " + errstr))
		respchan <- listToFrames(response)
	}
}

// Convert a doubly linked list of message frames to a slice of message
// frames.
func listToFrames(l *list.List) [][]byte {
	frames := make([][]byte, l.Len())
	i := 0
	for e := l.Front(); e != nil; e = e.Next() {
		frames[i] = e.Value.([]byte)
	}
	return frames
}

// Helper function for copying a doubly linked list.
func copyList(l *list.List) *list.List {
	replica := list.New()
	replica.PushBackList(l)
	return replica
}

// Main method. Will panic if things are so bad that the application
// will not start.
func main() {
	flag.Parse()

	desc, err := es.NewFileSystemDescriptor(*eventStorePath)
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not create event store Desc")
		panic(err)
	}
	estore, err := es.NewEventStore(desc)
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not create event store")
		panic(err)
	}
	defer estore.Close()

	runServer(estore)

	// TODO: Handle SIGINT correctly and smoothly.
}
<|endoftext|>