<commit_before>package resp

import "errors"

var (
	ErrNoData            = errors.New("resp: no data")
	ErrInvalidPrefix     = errors.New("resp: invalid prefix")
	ErrMissingCRLF       = errors.New("resp: missing CRLF")
	ErrInvalidInteger    = errors.New("resp: invalid integer character")
	ErrInvalidBulkString = errors.New("resp: invalid bulk string")
	ErrInvalidArray      = errors.New("resp: invalid array")
)

type Integer int64

type Error []byte

type SimpleString []byte

type BulkString []byte

type Array []interface{}

func DecodeRequest(b []byte) ([][]byte, error) {
	return nil, nil
}

func Decode(b []byte) ([]interface{}, error) {
	return nil, nil
}

func decodeValue(b []byte) (val interface{}, n int, err error) {
	if len(b) == 0 {
		return nil, 0, ErrNoData
	}

	switch b[0] {
	case '+':
		// Simple string
		val, n, err = decodeSimpleString(b[1:])
	case '-':
		// Error
		val, n, err = decodeError(b[1:])
	case ':':
		// Integer
		val, n, err = decodeInteger(b[1:])
	case '$':
		// Bulk string
		val, n, err = decodeBulkString(b[1:])
	case '*':
		// Array
		val, n, err = decodeArray(b[1:])
	default:
		err = ErrInvalidPrefix
	}

	// n+1 for the prefix consumed by this func
	return val, n + 1, err
}

func decodeArray(b []byte) (Array, int, error) {
	// First comes the number of elements in the array
	cnt, n, err := decodeInteger(b)
	if err != nil {
		return nil, n, err
	}
	switch {
	case cnt == -1:
		// Nil array
		return Array(nil), n, nil

	case cnt == 0:
		// Empty, but allocated, array
		return Array{}, n, nil

	case cnt < 0:
		// Invalid length
		return nil, n, ErrInvalidArray

	default:
		// Allocate the array
		ar := make(Array, cnt)

		// Decode each value
		for i := 0; i < int(cnt); i++ {
			val, nn, err := decodeValue(b[n:])
			n += nn
			if err != nil {
				return nil, n, err
			}
			ar[i] = val
		}
		return ar, n, nil
	}
}

func decodeBulkString(b []byte) (BulkString, int, error) {
	// First comes the length of the bulk string, an integer
	cnt, n, err := decodeInteger(b)
	if err != nil {
		return nil, n, err
	}
	switch {
	case cnt == -1:
		// Special case to represent a nil bulk string
		return BulkString(nil), n, nil

	case cnt < -1:
		return nil, n, ErrInvalidBulkString

	case len(b) < int(cnt)+n+2:
		return nil, n, ErrInvalidBulkString

	default:
		// Then the string is cnt long, and bytes read is cnt+n+2 (for ending CRLF)
		return BulkString(b[n : int(cnt)+n]), int(cnt) + n + 2, nil
	}
}

// decodeInteger decodes the byte slice as a singed 64bit integer. The
// ':' prefix is assumed to be already consumed.
func decodeInteger(b []byte) (val Integer, n int, err error) {
	var cr bool
	var sign Integer = 1

loop:
	for i := 0; i < len(b); i++ {
		ch := b[i]
		n++

		switch ch {
		case '\r':
			cr = true
			break loop

		case '\n':
			break loop

		case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
			val = val*10 + Integer(ch-'0')

		case '-':
			if i == 0 {
				sign = -1
				continue
			}
			fallthrough
		default:
			return 0, n, ErrInvalidInteger
		}
	}

	if !cr {
		return 0, n, ErrMissingCRLF
	}
	// Presume next byte was \n
	return sign * val, n + 1, nil
}

// decodeSimpleString decodes the byte slice as a SimpleString. The
// '+' prefix is assumed to be already consumed.
func decodeSimpleString(b []byte) (SimpleString, int, error) {
	end := -1
	n := 0
	for i := 0; i < len(b); i++ {
		ch := b[i]
		n++
		if ch == '\r' {
			// Simple strings cannot contain \r nor \n, so at the first \r we know
			// the string is over.
			end = i
			break
		}
	}
	if end == -1 {
		return nil, n, ErrMissingCRLF
	}
	// Presume next byte was \n
	return SimpleString(b[:end]), n + 1, nil
}

// decodeError decodes the byte slice as an Error. The '-' prefix
// is assumed to be already consumed.
func decodeError(b []byte) (Error, int, error) {
	val, n, err := decodeSimpleString(b)
	return Error(val), n, err
}
<commit_msg>document code<commit_after>// Package resp implements an efficient decoder for the Redis Serialization Protocol (RESP).
//
// See http://redis.io/topics/protocol for the reference.
package resp

import "errors"

var (
	// ErrNoData is returned if an empty slice was passed to a Decode function.
	ErrNoData = errors.New("resp: no data")

	// ErrInvalidPrefix is returned if the data contains an unrecognized prefix.
	ErrInvalidPrefix = errors.New("resp: invalid prefix")

	// ErrMissingCRLF is returned if a \r\n is missing in the data slice.
	ErrMissingCRLF = errors.New("resp: missing CRLF")

	// ErrInvalidInteger is returned if an invalid character is found while parsing an integer.
	ErrInvalidInteger = errors.New("resp: invalid integer character")

	// ErrInvalidBulkString is returned if the bulk string data cannot be decoded.
	ErrInvalidBulkString = errors.New("resp: invalid bulk string")

	// ErrInvalidArray is returned if the array data cannot be decoded.
	ErrInvalidArray = errors.New("resp: invalid array")

	// ErrNotAnArray is returned if the DecodeRequest function is called and
	// the decoded value is not an array.
	ErrNotAnArray = errors.New("resp: expected an array type")
)

// Integer represents a signed, 64-bit integer as defined by the RESP.
type Integer int64

// Error represents an error string as defined by the RESP. It cannot
// contain \r or \n characters.
type Error []byte

// SimpleString represents a simple string as defined by the RESP. It
// cannot contain \r or \n characters.
type SimpleString []byte

// BulkString represents a binary-safe string as defined by the RESP.
type BulkString []byte

// Array represents an array of values, as defined by the RESP.
type Array []interface{}

// DecodeRequest decodes the provided byte slice and returns the array
// representing the request. If the encoded value is not an array, it
// returns ErrNotAnArray.
func DecodeRequest(b []byte) (Array, error) {
	val, _, err := decodeValue(b)
	if err != nil {
		return nil, err
	}
	if ar, ok := val.(Array); !ok {
		return nil, ErrNotAnArray
	} else {
		return ar, nil
	}
}

// Decode decodes the provided byte slice and returns the parsed value.
func Decode(b []byte) ([]interface{}, error) {
	return nil, nil
}

// decodeValue parses the byte slice and decodes the value based on its
// prefix, as defined by the RESP protocol.
func decodeValue(b []byte) (val interface{}, n int, err error) {
	if len(b) == 0 {
		return nil, 0, ErrNoData
	}

	switch b[0] {
	case '+':
		// Simple string
		val, n, err = decodeSimpleString(b[1:])
	case '-':
		// Error
		val, n, err = decodeError(b[1:])
	case ':':
		// Integer
		val, n, err = decodeInteger(b[1:])
	case '$':
		// Bulk string
		val, n, err = decodeBulkString(b[1:])
	case '*':
		// Array
		val, n, err = decodeArray(b[1:])
	default:
		err = ErrInvalidPrefix
	}

	// n+1 for the prefix consumed by this func
	return val, n + 1, err
}

// decodeArray decodes the byte slice as an array. It assumes the
// '*' prefix is already consumed.
func decodeArray(b []byte) (Array, int, error) {
	// First comes the number of elements in the array
	cnt, n, err := decodeInteger(b)
	if err != nil {
		return nil, n, err
	}
	switch {
	case cnt == -1:
		// Nil array
		return nil, n, nil

	case cnt == 0:
		// Empty, but allocated, array
		return Array{}, n, nil

	case cnt < 0:
		// Invalid length
		return nil, n, ErrInvalidArray

	default:
		// Allocate the array
		ar := make(Array, cnt)

		// Decode each value
		for i := 0; i < int(cnt); i++ {
			val, nn, err := decodeValue(b[n:])
			n += nn
			if err != nil {
				return nil, n, err
			}
			ar[i] = val
		}
		return ar, n, nil
	}
}

// decodeBulkString decodes the byte slice as a binary-safe string. The
// '$' prefix is assumed to be already consumed.
func decodeBulkString(b []byte) (BulkString, int, error) {
	// First comes the length of the bulk string, an integer
	cnt, n, err := decodeInteger(b)
	if err != nil {
		return nil, n, err
	}
	switch {
	case cnt == -1:
		// Special case to represent a nil bulk string
		return nil, n, nil

	case cnt < -1:
		return nil, n, ErrInvalidBulkString

	case len(b) < int(cnt)+n+2:
		return nil, n, ErrInvalidBulkString

	default:
		// Then the string is cnt long, and bytes read is cnt+n+2 (for ending CRLF)
		return BulkString(b[n : int(cnt)+n]), int(cnt) + n + 2, nil
	}
}

// decodeInteger decodes the byte slice as a singed 64bit integer. The
// ':' prefix is assumed to be already consumed.
func decodeInteger(b []byte) (val Integer, n int, err error) {
	var cr bool
	var sign Integer = 1

loop:
	for i := 0; i < len(b); i++ {
		ch := b[i]
		n++

		switch ch {
		case '\r':
			cr = true
			break loop

		case '\n':
			break loop

		case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
			val = val*10 + Integer(ch-'0')

		case '-':
			if i == 0 {
				sign = -1
				continue
			}
			fallthrough
		default:
			return 0, n, ErrInvalidInteger
		}
	}

	if !cr {
		return 0, n, ErrMissingCRLF
	}
	// Presume next byte was \n
	return sign * val, n + 1, nil
}

// decodeSimpleString decodes the byte slice as a SimpleString. The
// '+' prefix is assumed to be already consumed.
func decodeSimpleString(b []byte) (SimpleString, int, error) {
	end := -1
	n := 0
	for i := 0; i < len(b); i++ {
		ch := b[i]
		n++
		if ch == '\r' {
			// Simple strings cannot contain \r nor \n, so at the first \r we know
			// the string is over.
			end = i
			break
		}
	}
	if end == -1 {
		return nil, n, ErrMissingCRLF
	}
	// Presume next byte was \n
	return SimpleString(b[:end]), n + 1, nil
}

// decodeError decodes the byte slice as an Error. The '-' prefix
// is assumed to be already consumed.
func decodeError(b []byte) (Error, int, error) {
	val, n, err := decodeSimpleString(b)
	return Error(val), n, err
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"sort"
	"strings"

	"github.com/lxc/lxd/shared"
)

// CompareConfigs compares two config maps and returns an error if they differ.
func CompareConfigs(config1, config2 map[string]string, exclude []string) error {
	if exclude == nil {
		exclude = []string{}
	}

	delta := []string{}
	for key, value := range config1 {
		if shared.StringInSlice(key, exclude) {
			continue
		}
		if config2[key] != value {
			delta = append(delta, key)
		}
	}
	for key, value := range config2 {
		if shared.StringInSlice(key, exclude) {
			continue
		}
		if config1[key] != value {
			present := false
			for i := range delta {
				if delta[i] == key {
					present = true
				}
				break
			}
			if !present {
				delta = append(delta, key)
			}
		}
	}
	sort.Strings(delta)
	if len(delta) > 0 {
		return fmt.Errorf("different values for keys: %s", strings.Join(delta, ", "))
	}

	return nil
}

// CopyConfig creates a new map with a copy of the given config.
func CopyConfig(config map[string]string) map[string]string {
	copy := make(map[string]string, len(config))
	for key, value := range config {
		copy[key] = value
	}

	return copy
}

// SplitNTrimSpace returns result of strings.SplitN() and then strings.TrimSpace() on each element.
// Accepts nilIfEmpty argument which if true, will return nil slice if s is empty (after trimming space).
func SplitNTrimSpace(s string, sep string, n int, nilIfEmpty bool) []string {
	if nilIfEmpty && strings.TrimSpace(s) == "" {
		return nil
	}

	parts := strings.SplitN(s, sep, n)

	for i, v := range parts {
		parts[i] = strings.TrimSpace(v)
	}

	return parts
}
<commit_msg>lxd/util: Remove SplitNTrimSpace<commit_after>package util

import (
	"fmt"
	"sort"
	"strings"

	"github.com/lxc/lxd/shared"
)

// CompareConfigs compares two config maps and returns an error if they differ.
func CompareConfigs(config1, config2 map[string]string, exclude []string) error {
	if exclude == nil {
		exclude = []string{}
	}

	delta := []string{}
	for key, value := range config1 {
		if shared.StringInSlice(key, exclude) {
			continue
		}
		if config2[key] != value {
			delta = append(delta, key)
		}
	}
	for key, value := range config2 {
		if shared.StringInSlice(key, exclude) {
			continue
		}
		if config1[key] != value {
			present := false
			for i := range delta {
				if delta[i] == key {
					present = true
				}
				break
			}
			if !present {
				delta = append(delta, key)
			}
		}
	}
	sort.Strings(delta)
	if len(delta) > 0 {
		return fmt.Errorf("different values for keys: %s", strings.Join(delta, ", "))
	}

	return nil
}

// CopyConfig creates a new map with a copy of the given config.
func CopyConfig(config map[string]string) map[string]string {
	copy := make(map[string]string, len(config))
	for key, value := range config {
		copy[key] = value
	}

	return copy
}
<|endoftext|>


<commit_before>package nsq

import (
	"fmt"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

// Producer is a high-level type to publish to NSQ.
//
// A Producer instance is 1:1 with a destination `nsqd`
// and will lazily connect to that instance (and re-connect)
// when Publish commands are executed.
type Producer struct {
	id     int64
	addr   string
	conn   *Conn
	config Config

	logger *log.Logger
	logLvl LogLevel

	responseChan  chan []byte
	errorChan     chan []byte
	ioErrorChan   chan error
	heartbeatChan chan int
	closeChan     chan int

	transactionChan chan *ProducerTransaction
	transactions    []*ProducerTransaction
	state           int32

	concurrentProducers int32
	stopFlag            int32
	exitChan            chan int
	wg                  sync.WaitGroup
	guard               sync.Mutex
}

// ProducerTransaction is returned by the async publish methods
// to retrieve metadata about the command after the
// response is received.
type ProducerTransaction struct {
	cmd      *Command
	doneChan chan *ProducerTransaction
	Error    error         // the error (or nil) of the publish command
	Args     []interface{} // the slice of variadic arguments passed to PublishAsync or MultiPublishAsync
}

func (t *ProducerTransaction) finish() {
	if t.doneChan != nil {
		t.doneChan <- t
	}
}

// NewProducer returns an instance of Producer for the specified address
//
// The only valid way to create a Config is via NewConfig, using a struct literal will panic.
// After Config is passed into NewProducer the values are no longer mutable (they are copied).
func NewProducer(addr string, config *Config) (*Producer, error) {
	config.assertInitialized()
	err := config.Validate()
	if err != nil {
		return nil, err
	}

	p := &Producer{
		id: atomic.AddInt64(&instCount, 1),

		addr:   addr,
		config: *config,

		logger: log.New(os.Stderr, "", log.Flags()),
		logLvl: LogLevelInfo,

		transactionChan: make(chan *ProducerTransaction),
		exitChan:        make(chan int),
		responseChan:    make(chan []byte),
		errorChan:       make(chan []byte),
		ioErrorChan:     make(chan error),
		heartbeatChan:   make(chan int),
		closeChan:       make(chan int),
	}
	return p, nil
}

// SetLogger assigns the logger to use as well as a level
func (w *Producer) SetLogger(logger *log.Logger, lvl LogLevel) {
	w.logger = logger
	w.logLvl = lvl
}

// String returns the address of the Producer
func (w *Producer) String() string {
	return w.addr
}

// Stop initiates a graceful stop of the Producer (permanent)
//
// NOTE: this blocks until completion
func (w *Producer) Stop() {
	w.guard.Lock()
	if !atomic.CompareAndSwapInt32(&w.stopFlag, 0, 1) {
		w.guard.Unlock()
		return
	}
	w.log(LogLevelInfo, "stopping")
	close(w.exitChan)
	w.close()
	w.guard.Unlock()
	w.wg.Wait()
}

// PublishAsync publishes a message body to the specified topic
// but does not wait for the response from `nsqd`.
//
// When the Producer eventually receives the response from `nsqd`,
// the supplied `doneChan` (if specified)
// will receive a `ProducerTransaction` instance with the supplied variadic arguments
// and the response error if present
func (w *Producer) PublishAsync(topic string, body []byte, doneChan chan *ProducerTransaction,
	args ...interface{}) error {
	return w.sendCommandAsync(Publish(topic, body), doneChan, args)
}

// MultiPublishAsync publishes a slice of message bodies to the specified topic
// but does not wait for the response from `nsqd`.
//
// When the Producer eventually receives the response from `nsqd`,
// the supplied `doneChan` (if specified)
// will receive a `ProducerTransaction` instance with the supplied variadic arguments
// and the response error if present
func (w *Producer) MultiPublishAsync(topic string, body [][]byte, doneChan chan *ProducerTransaction,
	args ...interface{}) error {
	cmd, err := MultiPublish(topic, body)
	if err != nil {
		return err
	}
	return w.sendCommandAsync(cmd, doneChan, args)
}

// Publish synchronously publishes a message body to the specified topic, returning
// the an error if publish failed
func (w *Producer) Publish(topic string, body []byte) error {
	return w.sendCommand(Publish(topic, body))
}

// MultiPublish synchronously publishes a slice of message bodies to the specified topic, returning
// the an error if publish failed
func (w *Producer) MultiPublish(topic string, body [][]byte) error {
	cmd, err := MultiPublish(topic, body)
	if err != nil {
		return err
	}
	return w.sendCommand(cmd)
}

func (w *Producer) sendCommand(cmd *Command) error {
	doneChan := make(chan *ProducerTransaction)
	err := w.sendCommandAsync(cmd, doneChan, nil)
	if err != nil {
		close(doneChan)
		return err
	}
	t := <-doneChan
	return t.Error
}

func (w *Producer) sendCommandAsync(cmd *Command, doneChan chan *ProducerTransaction,
	args []interface{}) error {
	// keep track of how many outstanding producers we're dealing with
	// in order to later ensure that we clean them all up...
	atomic.AddInt32(&w.concurrentProducers, 1)
	defer atomic.AddInt32(&w.concurrentProducers, -1)

	if atomic.LoadInt32(&w.state) != StateConnected {
		err := w.connect()
		if err != nil {
			return err
		}
	}

	t := &ProducerTransaction{
		cmd:      cmd,
		doneChan: doneChan,
		Args:     args,
	}

	select {
	case w.transactionChan <- t:
	case <-w.exitChan:
		return ErrStopped
	}

	return nil
}

func (w *Producer) connect() error {
	w.guard.Lock()
	defer w.guard.Unlock()

	if atomic.LoadInt32(&w.stopFlag) == 1 {
		return ErrStopped
	}

	if !atomic.CompareAndSwapInt32(&w.state, StateInit, StateConnected) {
		return ErrNotConnected
	}

	w.log(LogLevelInfo, "(%s) connecting to nsqd", w.addr)

	conn := NewConn(w.addr, &w.config, &producerConnDelegate{w})
	conn.SetLogger(w.logger, w.logLvl, fmt.Sprintf("%3d (%%s)", w.id))

	_, err := conn.Connect()
	if err != nil {
		conn.Close()
		w.log(LogLevelError, "(%s) error connecting to nsqd - %s", w.addr, err)
		atomic.StoreInt32(&w.state, StateInit)
		return err
	}
	w.conn = conn

	w.wg.Add(1)
	go w.router()

	return nil
}

func (w *Producer) close() {
	if !atomic.CompareAndSwapInt32(&w.state, StateConnected, StateDisconnected) {
		return
	}
	w.conn.Close()
	go func() {
		// we need to handle this in a goroutine so we don't
		// block the caller from making progress
		w.wg.Wait()
		atomic.StoreInt32(&w.state, StateInit)
	}()
}

func (w *Producer) router() {
	for {
		select {
		case t := <-w.transactionChan:
			w.transactions = append(w.transactions, t)
			err := w.conn.WriteCommand(t.cmd)
			if err != nil {
				w.log(LogLevelError, "(%s) sending command - %s", w.conn.String(), err)
				w.close()
			}
		case data := <-w.responseChan:
			w.popTransaction(FrameTypeResponse, data)
		case data := <-w.errorChan:
			w.popTransaction(FrameTypeError, data)
		case <-w.heartbeatChan:
		case <-w.ioErrorChan:
			w.close()
		case <-w.closeChan:
			goto exit
		case <-w.exitChan:
			goto exit
		}
	}

exit:
	w.transactionCleanup()
	w.wg.Done()
	w.log(LogLevelInfo, "exiting router")
}

func (w *Producer) popTransaction(frameType int32, data []byte) {
	t := w.transactions[0]
	w.transactions = w.transactions[1:]
	if frameType == FrameTypeError {
		t.Error = ErrProtocol{string(data)}
	}
	t.finish()
}

func (w *Producer) transactionCleanup() {
	// clean up transactions we can easily account for
	for _, t := range w.transactions {
		t.Error = ErrNotConnected
		t.finish()
	}
	w.transactions = w.transactions[:0]

	// spin and free up any writes that might have raced
	// with the cleanup process (blocked on writing
	// to transactionChan)
	for {
		select {
		case t := <-w.transactionChan:
			t.Error = ErrNotConnected
			t.finish()
		default:
			// keep spinning until there are 0 concurrent producers
			if atomic.LoadInt32(&w.concurrentProducers) == 0 {
				return
			}
			// give the runtime a chance to schedule other racing goroutines
			time.Sleep(5 * time.Millisecond)
			continue
		}
	}
}

func (w *Producer) log(lvl LogLevel, line string, args ...interface{}) {
	var prefix string

	if w.logger == nil {
		return
	}

	if w.logLvl > lvl {
		return
	}

	switch lvl {
	case LogLevelDebug:
		prefix = "DBG"
	case LogLevelInfo:
		prefix = "INF"
	case LogLevelWarning:
		prefix = "WRN"
	case LogLevelError:
		prefix = "ERR"
	}

	w.logger.Printf("%-4s %3d %s", prefix, w.id, fmt.Sprintf(line, args...))
}

func (w *Producer) onConnResponse(c *Conn, data []byte) { w.responseChan <- data }
func (w *Producer) onConnError(c *Conn, data []byte)    { w.errorChan <- data }
func (w *Producer) onConnHeartbeat(c *Conn)             { w.heartbeatChan <- 1 }
func (w *Producer) onConnIOError(c *Conn, err error)    { w.ioErrorChan <- err }
func (w *Producer) onConnClose(c *Conn)                 { w.closeChan <- 1 }
<commit_msg>producer: fix write IOError deadlock<commit_after>package nsq

import (
	"fmt"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

// Producer is a high-level type to publish to NSQ.
//
// A Producer instance is 1:1 with a destination `nsqd`
// and will lazily connect to that instance (and re-connect)
// when Publish commands are executed.
type Producer struct {
	id     int64
	addr   string
	conn   *Conn
	config Config

	logger *log.Logger
	logLvl LogLevel

	responseChan  chan []byte
	errorChan     chan []byte
	closeChan     chan int

	transactionChan chan *ProducerTransaction
	transactions    []*ProducerTransaction
	state           int32

	concurrentProducers int32
	stopFlag            int32
	exitChan            chan int
	wg                  sync.WaitGroup
	guard               sync.Mutex
}

// ProducerTransaction is returned by the async publish methods
// to retrieve metadata about the command after the
// response is received.
type ProducerTransaction struct {
	cmd      *Command
	doneChan chan *ProducerTransaction
	Error    error         // the error (or nil) of the publish command
	Args     []interface{} // the slice of variadic arguments passed to PublishAsync or MultiPublishAsync
}

func (t *ProducerTransaction) finish() {
	if t.doneChan != nil {
		t.doneChan <- t
	}
}

// NewProducer returns an instance of Producer for the specified address
//
// The only valid way to create a Config is via NewConfig, using a struct literal will panic.
// After Config is passed into NewProducer the values are no longer mutable (they are copied).
func NewProducer(addr string, config *Config) (*Producer, error) {
	config.assertInitialized()
	err := config.Validate()
	if err != nil {
		return nil, err
	}

	p := &Producer{
		id: atomic.AddInt64(&instCount, 1),

		addr:   addr,
		config: *config,

		logger: log.New(os.Stderr, "", log.Flags()),
		logLvl: LogLevelInfo,

		transactionChan: make(chan *ProducerTransaction),
		exitChan:        make(chan int),
		responseChan:    make(chan []byte),
		errorChan:       make(chan []byte),
		closeChan:       make(chan int),
	}
	return p, nil
}

// SetLogger assigns the logger to use as well as a level
func (w *Producer) SetLogger(logger *log.Logger, lvl LogLevel) {
	w.logger = logger
	w.logLvl = lvl
}

// String returns the address of the Producer
func (w *Producer) String() string {
	return w.addr
}

// Stop initiates a graceful stop of the Producer (permanent)
//
// NOTE: this blocks until completion
func (w *Producer) Stop() {
	w.guard.Lock()
	if !atomic.CompareAndSwapInt32(&w.stopFlag, 0, 1) {
		w.guard.Unlock()
		return
	}
	w.log(LogLevelInfo, "stopping")
	close(w.exitChan)
	w.close()
	w.guard.Unlock()
	w.wg.Wait()
}

// PublishAsync publishes a message body to the specified topic
// but does not wait for the response from `nsqd`.
//
// When the Producer eventually receives the response from `nsqd`,
// the supplied `doneChan` (if specified)
// will receive a `ProducerTransaction` instance with the supplied variadic arguments
// and the response error if present
func (w *Producer) PublishAsync(topic string, body []byte, doneChan chan *ProducerTransaction,
	args ...interface{}) error {
	return w.sendCommandAsync(Publish(topic, body), doneChan, args)
}

// MultiPublishAsync publishes a slice of message bodies to the specified topic
// but does not wait for the response from `nsqd`.
//
// When the Producer eventually receives the response from `nsqd`,
// the supplied `doneChan` (if specified)
// will receive a `ProducerTransaction` instance with the supplied variadic arguments
// and the response error if present
func (w *Producer) MultiPublishAsync(topic string, body [][]byte, doneChan chan *ProducerTransaction,
	args ...interface{}) error {
	cmd, err := MultiPublish(topic, body)
	if err != nil {
		return err
	}
	return w.sendCommandAsync(cmd, doneChan, args)
}

// Publish synchronously publishes a message body to the specified topic, returning
// the an error if publish failed
func (w *Producer) Publish(topic string, body []byte) error {
	return w.sendCommand(Publish(topic, body))
}

// MultiPublish synchronously publishes a slice of message bodies to the specified topic, returning
// the an error if publish failed
func (w *Producer) MultiPublish(topic string, body [][]byte) error {
	cmd, err := MultiPublish(topic, body)
	if err != nil {
		return err
	}
	return w.sendCommand(cmd)
}

func (w *Producer) sendCommand(cmd *Command) error {
	doneChan := make(chan *ProducerTransaction)
	err := w.sendCommandAsync(cmd, doneChan, nil)
	if err != nil {
		close(doneChan)
		return err
	}
	t := <-doneChan
	return t.Error
}

func (w *Producer) sendCommandAsync(cmd *Command, doneChan chan *ProducerTransaction,
	args []interface{}) error {
	// keep track of how many outstanding producers we're dealing with
	// in order to later ensure that we clean them all up...
	atomic.AddInt32(&w.concurrentProducers, 1)
	defer atomic.AddInt32(&w.concurrentProducers, -1)

	if atomic.LoadInt32(&w.state) != StateConnected {
		err := w.connect()
		if err != nil {
			return err
		}
	}

	t := &ProducerTransaction{
		cmd:      cmd,
		doneChan: doneChan,
		Args:     args,
	}

	select {
	case w.transactionChan <- t:
	case <-w.exitChan:
		return ErrStopped
	}

	return nil
}

func (w *Producer) connect() error {
	w.guard.Lock()
	defer w.guard.Unlock()

	if atomic.LoadInt32(&w.stopFlag) == 1 {
		return ErrStopped
	}

	if !atomic.CompareAndSwapInt32(&w.state, StateInit, StateConnected) {
		return ErrNotConnected
	}

	w.log(LogLevelInfo, "(%s) connecting to nsqd", w.addr)

	conn := NewConn(w.addr, &w.config, &producerConnDelegate{w})
	conn.SetLogger(w.logger, w.logLvl, fmt.Sprintf("%3d (%%s)", w.id))

	_, err := conn.Connect()
	if err != nil {
		conn.Close()
		w.log(LogLevelError, "(%s) error connecting to nsqd - %s", w.addr, err)
		atomic.StoreInt32(&w.state, StateInit)
		return err
	}
	w.conn = conn

	w.wg.Add(1)
	go w.router()

	return nil
}

func (w *Producer) close() {
	if !atomic.CompareAndSwapInt32(&w.state, StateConnected, StateDisconnected) {
		return
	}
	w.conn.Close()
	go func() {
		// we need to handle this in a goroutine so we don't
		// block the caller from making progress
		w.wg.Wait()
		atomic.StoreInt32(&w.state, StateInit)
	}()
}

func (w *Producer) router() {
	for {
		select {
		case t := <-w.transactionChan:
			w.transactions = append(w.transactions, t)
			err := w.conn.WriteCommand(t.cmd)
			if err != nil {
				w.log(LogLevelError, "(%s) sending command - %s", w.conn.String(), err)
				w.close()
			}
		case data := <-w.responseChan:
			w.popTransaction(FrameTypeResponse, data)
		case data := <-w.errorChan:
			w.popTransaction(FrameTypeError, data)
		case <-w.closeChan:
			goto exit
		case <-w.exitChan:
			goto exit
		}
	}

exit:
	w.transactionCleanup()
	w.wg.Done()
	w.log(LogLevelInfo, "exiting router")
}

func (w *Producer) popTransaction(frameType int32, data []byte) {
	t := w.transactions[0]
	w.transactions = w.transactions[1:]
	if frameType == FrameTypeError {
		t.Error = ErrProtocol{string(data)}
	}
	t.finish()
}

func (w *Producer) transactionCleanup() {
	// clean up transactions we can easily account for
	for _, t := range w.transactions {
		t.Error = ErrNotConnected
		t.finish()
	}
	w.transactions = w.transactions[:0]

	// spin and free up any writes that might have raced
	// with the cleanup process (blocked on writing
	// to transactionChan)
	for {
		select {
		case t := <-w.transactionChan:
			t.Error = ErrNotConnected
			t.finish()
		default:
			// keep spinning until there are 0 concurrent producers
			if atomic.LoadInt32(&w.concurrentProducers) == 0 {
				return
			}
			// give the runtime a chance to schedule other racing goroutines
			time.Sleep(5 * time.Millisecond)
			continue
		}
	}
}

func (w *Producer) log(lvl LogLevel, line string, args ...interface{}) {
	var prefix string

	if w.logger == nil {
		return
	}

	if w.logLvl > lvl {
		return
	}

	switch lvl {
	case LogLevelDebug:
		prefix = "DBG"
	case LogLevelInfo:
		prefix = "INF"
	case LogLevelWarning:
		prefix = "WRN"
	case LogLevelError:
		prefix = "ERR"
	}

	w.logger.Printf("%-4s %3d %s", prefix, w.id, fmt.Sprintf(line, args...))
}

func (w *Producer) onConnResponse(c *Conn, data []byte) { w.responseChan <- data }
func (w *Producer) onConnError(c *Conn, data []byte)    { w.errorChan <- data }
func (w *Producer) onConnHeartbeat(c *Conn)             {}
func (w *Producer) onConnIOError(c *Conn, err error)    { w.close() }
func (w *Producer) onConnClose(c *Conn)                 { w.closeChan <- 1 }
<|endoftext|>


<commit_before>package mpb

import (
	"container/heap"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

const (
	// default RefreshRate
	prr = 120 * time.Millisecond
	// default width
	pwidth = 80
	// default format
	pformat = "[=>-]"
)

// Progress represents the container that renders Progress bars
type Progress struct {
	wg           *sync.WaitGroup
	uwg          *sync.WaitGroup
	operateState chan func(*pState)
	done         chan struct{}
}

type (
	pState struct {
		bHeap           *priorityQueue
		shutdownPending []*Bar
		heapUpdated     bool
		zeroWait        bool
		idCounter       int
		width           int
		format          string
		rr              time.Duration
		cw              *cwriter.Writer
		ticker          *time.Ticker
		pMatrix         map[int][]chan int
		aMatrix         map[int][]chan int

		// following are provided by user
		uwg              *sync.WaitGroup
		cancel           <-chan struct{}
		shutdownNotifier chan struct{}
		waitBars         map[*Bar]*Bar
		debugOut         io.Writer
	}
	widthSyncer struct {
		// Public for easy testing
		Accumulator []chan int
		Distributor []chan int
	}
)

// New creates new Progress instance, which orchestrates bars rendering process.
// Accepts mpb.ProgressOption funcs for customization.
func New(options ...ProgressOption) *Progress {
	pq := make(priorityQueue, 0)
	heap.Init(&pq)
	s := &pState{
		bHeap:    &pq,
		width:    pwidth,
		format:   pformat,
		cw:       cwriter.New(os.Stdout),
		rr:       prr,
		ticker:   time.NewTicker(prr),
		waitBars: make(map[*Bar]*Bar),
		debugOut: ioutil.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		uwg:          s.uwg,
		wg:           new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		done:         make(chan struct{}),
	}
	go p.serve(s)
	return p
}

// AddBar creates a new progress bar and adds to the container.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	p.wg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(s *pState) {
		options = append(options, barWidth(s.width), barFormat(s.format))
		b := newBar(p.wg, s.idCounter, total, s.cancel, options...)
		if b.runningBar != nil {
			s.waitBars[b.runningBar] = b
		} else {
			heap.Push(s.bHeap, b)
			s.heapUpdated = true
		}
		s.idCounter++
		result <- b
	}:
		return <-result
	case <-p.done:
		p.wg.Done()
		return nil
	}
}

// Abort is only effective while bar progress is running,
// it means remove bar now without waiting for its completion.
// If bar is already completed, there is nothing to abort.
// If you need to remove bar after completion, use BarRemoveOnComplete BarOption.
func (p *Progress) Abort(b *Bar) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		s.heapUpdated = heap.Remove(s.bHeap, b.index) != nil
		s.shutdownPending = append(s.shutdownPending, b)
	}:
	case <-p.done:
	}
}

// UpdateBarPriority provides a way to change bar's order position.
// Zero is highest priority, i.e. bar will be on top.
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) { s.bHeap.update(b, priority) }:
	case <-p.done:
	}
}

// BarCount returns bars count
func (p *Progress) BarCount() int {
	result := make(chan int, 1)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Wait first waits for user provided *sync.WaitGroup, if any,
// then waits far all bars to complete and finally shutdowns master goroutine.
// After this method has been called, there is no way to reuse *Progress instance.
func (p *Progress) Wait() {
	if p.uwg != nil {
		p.uwg.Wait()
	}

	p.wg.Wait()

	select {
	case p.operateState <- func(s *pState) { s.zeroWait = true }:
		<-p.done
	case <-p.done:
	}
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		column := column
		go func() {
			var maxWidth int
			for _, ch := range column {
				w := <-ch
				if w > maxWidth {
					maxWidth = w
				}
			}
			for _, ch := range column {
				ch <- maxWidth
			}
		}()
	}
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) render(tw int) {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		go bar.render(s.debugOut, tw)
	}

	if err := s.flush(); err != nil {
		fmt.Fprintf(s.debugOut, "%s %s %v\n", "[mpb]", time.Now(), err)
	}
}

func (s *pState) flush() (err error) {
	for s.bHeap.Len() > 0 {
		bar := heap.Pop(s.bHeap).(*Bar)
		reader := <-bar.frameReaderCh
		if _, e := s.cw.ReadFrom(reader); e != nil {
			err = e
		}
		defer func() {
			if frame, ok := reader.(*frameReader); ok && frame.toShutdown {
				// shutdown at next flush, in other words decrement underlying WaitGroup
				// only after the bar with completed state has been flushed.
				// this ensures no bar ends up with less than 100% rendered.
				s.shutdownPending = append(s.shutdownPending, bar)
				if replacementBar, ok := s.waitBars[bar]; ok {
					heap.Push(s.bHeap, replacementBar)
					s.heapUpdated = true
					delete(s.waitBars, bar)
				}
				if frame.removeOnComplete {
					s.heapUpdated = true
					return
				}
			}
			heap.Push(s.bHeap, bar)
		}()
	}

	if e := s.cw.Flush(); err == nil {
		err = e
	}

	for i := len(s.shutdownPending) - 1; i >= 0; i-- {
		close(s.shutdownPending[i].shutdown)
		s.shutdownPending = s.shutdownPending[:i]
	}
	return
}
<commit_msg>func to bottom<commit_after>package mpb

import (
	"container/heap"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

const (
	// default RefreshRate
	prr = 120 * time.Millisecond
	// default width
	pwidth = 80
	// default format
	pformat = "[=>-]"
)

// Progress represents the container that renders Progress bars
type Progress struct {
	wg           *sync.WaitGroup
	uwg          *sync.WaitGroup
	operateState chan func(*pState)
	done         chan struct{}
}

type (
	pState struct {
		bHeap           *priorityQueue
		shutdownPending []*Bar
		heapUpdated     bool
		zeroWait        bool
		idCounter       int
		width           int
		format          string
		rr              time.Duration
		cw              *cwriter.Writer
		ticker          *time.Ticker
		pMatrix         map[int][]chan int
		aMatrix         map[int][]chan int

		// following are provided by user
		uwg              *sync.WaitGroup
		cancel           <-chan struct{}
		shutdownNotifier chan struct{}
		waitBars         map[*Bar]*Bar
		debugOut         io.Writer
	}
	widthSyncer struct {
		// Public for easy testing
		Accumulator []chan int
		Distributor []chan int
	}
)

// New creates new Progress instance, which orchestrates bars rendering process.
// Accepts mpb.ProgressOption funcs for customization.
func New(options ...ProgressOption) *Progress {
	pq := make(priorityQueue, 0)
	heap.Init(&pq)
	s := &pState{
		bHeap:    &pq,
		width:    pwidth,
		format:   pformat,
		cw:       cwriter.New(os.Stdout),
		rr:       prr,
		ticker:   time.NewTicker(prr),
		waitBars: make(map[*Bar]*Bar),
		debugOut: ioutil.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		uwg:          s.uwg,
		wg:           new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		done:         make(chan struct{}),
	}
	go p.serve(s)
	return p
}

// AddBar creates a new progress bar and adds to the container.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	p.wg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(s *pState) {
		options = append(options, barWidth(s.width), barFormat(s.format))
		b := newBar(p.wg, s.idCounter, total, s.cancel, options...)
		if b.runningBar != nil {
			s.waitBars[b.runningBar] = b
		} else {
			heap.Push(s.bHeap, b)
			s.heapUpdated = true
		}
		s.idCounter++
		result <- b
	}:
		return <-result
	case <-p.done:
		p.wg.Done()
		return nil
	}
}

// Abort is only effective while bar progress is running,
// it means remove bar now without waiting for its completion.
// If bar is already completed, there is nothing to abort.
// If you need to remove bar after completion, use BarRemoveOnComplete BarOption.
func (p *Progress) Abort(b *Bar) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		s.heapUpdated = heap.Remove(s.bHeap, b.index) != nil
		s.shutdownPending = append(s.shutdownPending, b)
	}:
	case <-p.done:
	}
}

// UpdateBarPriority provides a way to change bar's order position.
// Zero is highest priority, i.e. bar will be on top.
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) { s.bHeap.update(b, priority) }:
	case <-p.done:
	}
}

// BarCount returns bars count
func (p *Progress) BarCount() int {
	result := make(chan int, 1)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Wait first waits for user provided *sync.WaitGroup, if any,
// then waits far all bars to complete and finally shutdowns master goroutine.
// After this method has been called, there is no way to reuse *Progress instance.
func (p *Progress) Wait() {
	if p.uwg != nil {
		p.uwg.Wait()
	}

	p.wg.Wait()

	select {
	case p.operateState <- func(s *pState) { s.zeroWait = true }:
		<-p.done
	case <-p.done:
	}
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) render(tw int) {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		go bar.render(s.debugOut, tw)
	}

	if err := s.flush(); err != nil {
		fmt.Fprintf(s.debugOut, "%s %s %v\n", "[mpb]", time.Now(), err)
	}
}

func (s *pState) flush() (err error) {
	for s.bHeap.Len() > 0 {
		bar := heap.Pop(s.bHeap).(*Bar)
		reader := <-bar.frameReaderCh
		if _, e := s.cw.ReadFrom(reader); e != nil {
			err = e
		}
		defer func() {
			if frame, ok := reader.(*frameReader); ok && frame.toShutdown {
				// shutdown at next flush, in other words decrement underlying WaitGroup
				// only after the bar with completed state has been flushed.
				// this ensures no bar ends up with less than 100% rendered.
				s.shutdownPending = append(s.shutdownPending, bar)
				if replacementBar, ok := s.waitBars[bar]; ok {
					heap.Push(s.bHeap, replacementBar)
					s.heapUpdated = true
					delete(s.waitBars, bar)
				}
				if frame.removeOnComplete {
					s.heapUpdated = true
					return
				}
			}
			heap.Push(s.bHeap, bar)
		}()
	}

	if e := s.cw.Flush(); err == nil {
		err = e
	}

	for i := len(s.shutdownPending) - 1; i >= 0; i-- {
		close(s.shutdownPending[i].shutdown)
		s.shutdownPending = s.shutdownPending[:i]
	}
	return
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		column := column
		go func() {
			var maxWidth int
			for _, ch := range column {
				w := <-ch
				if w > maxWidth {
					maxWidth = w
				}
			}
			for _, ch := range column {
				ch <- maxWidth
			}
		}()
	}
}
<|endoftext|>


<commit_before>package mpb

import (
	"errors"
	"io"
	"log"
	"os"
	"runtime"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb/cwriter"
)

var logger = log.New(os.Stderr, "mpb: ", log.LstdFlags|log.Lshortfile)

// ErrCallAfterStop thrown by panic, if Progress methods like (*Progress).AddBar()
// are called after (*Progress).Stop() has been called
var ErrCallAfterStop = errors.New("method call on stopped Progress instance")

type (
	// BeforeRender is a func, which gets called before render process
	BeforeRender func([]*Bar)
	barOpType    uint

	operation struct {
		kind   barOpType
		bar    *Bar
		result chan bool
	}

	indexedBarBuffer struct {
		index int
		buf   []byte
	}

	indexedBar struct {
		index     int
		termWidth int
		bar       *Bar
	}

	widthSync struct {
		listen []chan int
		result []chan int
	}
)

const (
	barAdd barOpType = iota
	barRemove
)

const (
	// default RefreshRate
	rr = 100
	// default width
	pwidth = 70
	// number of format runes for bar
	numFmtRunes = 5
)

// Progress represents the container that renders Progress bars
type Progress struct {
	// Context for canceling bars rendering
	// ctx context.Context
	// WaitGroup for internal rendering sync
	wg *sync.WaitGroup

	out    io.Writer
	width  int
	format string

	operationCh    chan *operation
	rrChangeReqCh  chan time.Duration
	outChangeReqCh chan io.Writer
	barCountReqCh  chan chan int
	brCh           chan BeforeRender
	done           chan struct{}
	cancel         <-chan struct{}
}

// New creates new Progress instance, which will orchestrate bars rendering
// process. It acceepts context.Context, for cancellation.
// If you don't plan to cancel, it is safe to feed with nil
func New() *Progress {
	p := &Progress{
		width:          pwidth,
		operationCh:    make(chan *operation),
		rrChangeReqCh:  make(chan time.Duration),
		outChangeReqCh: make(chan io.Writer),
		barCountReqCh:  make(chan chan int),
		brCh:           make(chan BeforeRender),
		done:           make(chan struct{}),
		wg:             new(sync.WaitGroup),
	}
	go p.server(cwriter.New(os.Stdout))
	return p
}

// WithCancel cancellation via channel
func (p *Progress) WithCancel(ch <-chan struct{}) *Progress {
	if ch == nil {
		panic("nil cancel channel")
	}
	p2 := new(Progress)
	*p2 = *p
	p2.cancel = ch
	return p2
}

// SetWidth overrides default (70) width of bar(s)
func (p *Progress) SetWidth(n int) *Progress {
	if n < 0 {
		panic("negative width")
	}
	p2 := new(Progress)
	*p2 = *p
	p2.width = n
	return p2
}

// SetOut sets underlying writer of progress. Default is os.Stdout
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) SetOut(w io.Writer) *Progress {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	if w == nil {
		return p
	}
	p.outChangeReqCh <- w
	return p
}

// RefreshRate overrides default (100ms) refresh rate value
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) RefreshRate(d time.Duration) *Progress {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	p.rrChangeReqCh <- d
	return p
}

// BeforeRenderFunc accepts a func, which gets called before render process.
func (p *Progress) BeforeRenderFunc(f BeforeRender) *Progress {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	p.brCh <- f
	return p
}

// AddBar creates a new progress bar and adds to the container
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) AddBar(total int64) *Bar {
	return p.AddBarWithID(0, total)
}

// AddBarWithID creates a new progress bar and adds to the container
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) AddBarWithID(id int, total int64) *Bar {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	result := make(chan bool)
	bar := newBar(id, total, p.width, p.format, p.wg, p.cancel)
	p.operationCh <- &operation{barAdd, bar, result}
	if <-result {
		p.wg.Add(1)
	}
	return bar
}

// RemoveBar removes bar at any time.
// Pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) RemoveBar(b *Bar) bool {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	result := make(chan bool)
	p.operationCh <- &operation{barRemove, b, result}
	return <-result
}

// BarCount returns bars count in the container.
// Pancis if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) BarCount() int {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	respCh := make(chan int, 1)
	p.barCountReqCh <- respCh
	return <-respCh
}

// Format sets custom format for underlying bar(s).
// The default one is "[=>-]"
func (p *Progress) Format(format string) *Progress {
	if utf8.RuneCountInString(format) != numFmtRunes {
		return p
	}
	p.format = format
	return p
}

// Stop shutdowns Progress' goroutine
// Should be called only after each bor's work done, i.e. bar has reached its
// 100 %. It is NOT for concelation. Use WithContext or WithCancel for
// cancelation purposes.
func (p *Progress) Stop() {
	p.wg.Wait()
	if isClosed(p.done) {
		return
	}
	close(p.operationCh)
}

// server monitors underlying channels and renders any progress bars
func (p *Progress) server(cw *cwriter.Writer) {
	userRR := rr * time.Millisecond
	t := time.NewTicker(userRR)

	defer func() {
		t.Stop()
		close(p.done)
	}()

	var wg sync.WaitGroup
	recoverIfPanic := func() {
		if p := recover(); p != nil {
			logger.Printf("unexpected panic: %+v\n", p)
			var buf [4096]byte
			n := runtime.Stack(buf[:], false)
			os.Stderr.Write(buf[:n])
		}
		wg.Done()
	}
	var beforeRender BeforeRender
	bars := make([]*Bar, 0, 3)

	for {
		select {
		case w := <-p.outChangeReqCh:
			cw.Flush()
			cw = cwriter.New(w)
		case op, ok := <-p.operationCh:
			if !ok {
				return
			}
			switch op.kind {
			case barAdd:
				bars = append(bars, op.bar)
				op.result <- true
			case barRemove:
				var ok bool
				for i, b := range bars {
					if b == op.bar {
						bars = append(bars[:i], bars[i+1:]...)
						ok = true
						b.remove()
						break
					}
				}
				op.result <- ok
			}
		case respCh := <-p.barCountReqCh:
			respCh <- len(bars)
		case beforeRender = <-p.brCh:
		case <-t.C:
			numBars := len(bars)

			if numBars == 0 {
				break
			}

			if beforeRender != nil {
				beforeRender(bars)
			}

			prependWs := newWidthSync(userRR, numBars, bars[0].NumOfPrependers())
			appendWs := newWidthSync(userRR, numBars, bars[0].NumOfAppenders())

			width, _, _ := cwriter.GetTermSize()
			ibars := iBarsGen(bars, width)
			ibbCh := make(chan indexedBarBuffer)
			wg.Add(numBars)
			for i := 0; i < numBars; i++ {
				go func() {
					defer recoverIfPanic()
					drawer(ibars, ibbCh, prependWs, appendWs)
				}()
			}
			go func() {
				wg.Wait()
				close(ibbCh)
				for _, ch := range prependWs.listen {
					close(ch)
				}
				for _, ch := range appendWs.listen {
					close(ch)
				}
			}()

			m := make(map[int][]byte, len(bars))
			for ibb := range ibbCh {
				m[ibb.index] = ibb.buf
			}
			for i := 0; i < len(bars); i++ {
				cw.Write(m[i])
			}

			cw.Flush()

			for _, b := range bars {
				b.flushed()
			}
		case userRR = <-p.rrChangeReqCh:
			t.Stop()
			t = time.NewTicker(userRR)
		case <-p.cancel:
			return
		}
	}
}

func newWidthSync(userRR time.Duration, numBars, numColumn int) *widthSync {
	ws := &widthSync{
		listen: make([]chan int, numColumn),
		result: make([]chan int, numColumn),
	}
	for i := 0; i < numColumn; i++ {
		ws.listen[i] = make(chan int, numBars)
		ws.result[i] = make(chan int, numBars)
	}
	for i := 0; i < numColumn; i++ {
		go func(listenCh <-chan int, resultCh chan<- int) {
			defer close(resultCh)
			widths := make([]int, 0, numBars)
		loop:
			for {
				select {
				case w := <-listenCh:
					widths = append(widths, w)
					if len(widths) == numBars {
						break loop
					}
				case <-time.After(userRR):
					return
				}
			}
			result := max(widths)
			for i := 0; i < numBars; i++ {
				resultCh <- result
			}
		}(ws.listen[i], ws.result[i])
	}
	return ws
}

func drawer(ibars <-chan indexedBar, ibbCh chan<- indexedBarBuffer, prependWs, appendWs *widthSync) {
	for b := range ibars {
		buf := b.bar.bytes(b.termWidth, prependWs, appendWs)
		buf = append(buf, '\n')
		ibbCh <- indexedBarBuffer{b.index, buf}
	}
}

func iBarsGen(bars []*Bar, width int) <-chan indexedBar {
	ibars := make(chan indexedBar)
	go func() {
		defer close(ibars)
		for i, b := range bars {
			ibars <- indexedBar{i, width, b}
		}
	}()
	return ibars
}

// isClosed check if ch closed
// caution see: http://www.tapirgames.com/blog/golang-channel-closing
func isClosed(ch <-chan struct{}) bool {
	select {
	case <-ch:
		return true
	default:
		return false
	}
}

func max(slice []int) int {
	max := slice[0]

	for i := 1; i < len(slice); i++ {
		if slice[i] > max {
			max = slice[i]
		}
	}

	return max
}
<commit_msg>init cwriter inside server()<commit_after>package mpb

import (
	"errors"
	"io"
	"log"
	"os"
	"runtime"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb/cwriter"
)

var logger = log.New(os.Stderr, "mpb: ", log.LstdFlags|log.Lshortfile)

// ErrCallAfterStop thrown by panic, if Progress methods like (*Progress).AddBar()
// are called after (*Progress).Stop() has been called
var ErrCallAfterStop = errors.New("method call on stopped Progress instance")

type (
	// BeforeRender is a func, which gets called before render process
	BeforeRender func([]*Bar)
	barOpType    uint

	operation struct {
		kind   barOpType
		bar    *Bar
		result chan bool
	}

	indexedBarBuffer struct {
		index int
		buf   []byte
	}

	indexedBar struct {
		index     int
		termWidth int
		bar       *Bar
	}

	widthSync struct {
		listen []chan int
		result []chan int
	}
)

const (
	barAdd barOpType = iota
	barRemove
)

const (
	// default RefreshRate
	rr = 100
	// default width
	pwidth = 70
	// number of format runes for bar
	numFmtRunes = 5
)

// Progress represents the container that renders Progress bars
type Progress struct {
	// Context for canceling bars rendering
	// ctx context.Context
	// WaitGroup for internal rendering sync
	wg *sync.WaitGroup

	out    io.Writer
	width  int
	format string

	operationCh    chan *operation
	rrChangeReqCh  chan time.Duration
	outChangeReqCh chan io.Writer
	barCountReqCh  chan chan int
	brCh           chan BeforeRender
	done           chan struct{}
	cancel         <-chan struct{}
}

// New creates new Progress instance, which will orchestrate bars rendering
// process. It acceepts context.Context, for cancellation.
// If you don't plan to cancel, it is safe to feed with nil
func New() *Progress {
	p := &Progress{
		width:          pwidth,
		operationCh:    make(chan *operation),
		rrChangeReqCh:  make(chan time.Duration),
		outChangeReqCh: make(chan io.Writer),
		barCountReqCh:  make(chan chan int),
		brCh:           make(chan BeforeRender),
		done:           make(chan struct{}),
		wg:             new(sync.WaitGroup),
	}
	go p.server()
	return p
}

// WithCancel cancellation via channel
func (p *Progress) WithCancel(ch <-chan struct{}) *Progress {
	if ch == nil {
		panic("nil cancel channel")
	}
	p2 := new(Progress)
	*p2 = *p
	p2.cancel = ch
	return p2
}

// SetWidth overrides default (70) width of bar(s)
func (p *Progress) SetWidth(n int) *Progress {
	if n < 0 {
		panic("negative width")
	}
	p2 := new(Progress)
	*p2 = *p
	p2.width = n
	return p2
}

// SetOut sets underlying writer of progress. Default is os.Stdout
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) SetOut(w io.Writer) *Progress {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	if w == nil {
		return p
	}
	p.outChangeReqCh <- w
	return p
}

// RefreshRate overrides default (100ms) refresh rate value
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) RefreshRate(d time.Duration) *Progress {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	p.rrChangeReqCh <- d
	return p
}

// BeforeRenderFunc accepts a func, which gets called before render process.
func (p *Progress) BeforeRenderFunc(f BeforeRender) *Progress {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	p.brCh <- f
	return p
}

// AddBar creates a new progress bar and adds to the container
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) AddBar(total int64) *Bar {
	return p.AddBarWithID(0, total)
}

// AddBarWithID creates a new progress bar and adds to the container
// pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) AddBarWithID(id int, total int64) *Bar {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	result := make(chan bool)
	bar := newBar(id, total, p.width, p.format, p.wg, p.cancel)
	p.operationCh <- &operation{barAdd, bar, result}
	if <-result {
		p.wg.Add(1)
	}
	return bar
}

// RemoveBar removes bar at any time.
// Pancis, if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) RemoveBar(b *Bar) bool {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	result := make(chan bool)
	p.operationCh <- &operation{barRemove, b, result}
	return <-result
}

// BarCount returns bars count in the container.
// Pancis if called on stopped Progress instance, i.e after (*Progress).Stop()
func (p *Progress) BarCount() int {
	if isClosed(p.done) {
		panic(ErrCallAfterStop)
	}
	respCh := make(chan int, 1)
	p.barCountReqCh <- respCh
	return <-respCh
}

// Format sets custom format for underlying bar(s).
// The default one is "[=>-]"
func (p *Progress) Format(format string) *Progress {
	if utf8.RuneCountInString(format) != numFmtRunes {
		return p
	}
	p.format = format
	return p
}

// Stop shutdowns Progress' goroutine
// Should be called only after each bor's work done, i.e. bar has reached its
// 100 %. It is NOT for concelation. Use WithContext or WithCancel for
// cancelation purposes.
func (p *Progress) Stop() {
	p.wg.Wait()
	if isClosed(p.done) {
		return
	}
	close(p.operationCh)
}

// server monitors underlying channels and renders any progress bars
func (p *Progress) server() {
	userRR := rr * time.Millisecond
	t := time.NewTicker(userRR)

	defer func() {
		t.Stop()
		close(p.done)
	}()

	var wg sync.WaitGroup
	recoverIfPanic := func() {
		if p := recover(); p != nil {
			logger.Printf("unexpected panic: %+v\n", p)
			var buf [4096]byte
			n := runtime.Stack(buf[:], false)
			os.Stderr.Write(buf[:n])
		}
		wg.Done()
	}
	var beforeRender BeforeRender
	cw := cwriter.New(os.Stdout)
	bars := make([]*Bar, 0, 3)

	for {
		select {
		case w := <-p.outChangeReqCh:
			cw.Flush()
			cw = cwriter.New(w)
		case op, ok := <-p.operationCh:
			if !ok {
				return
			}
			switch op.kind {
			case barAdd:
				bars = append(bars, op.bar)
				op.result <- true
			case barRemove:
				var ok bool
				for i, b := range bars {
					if b == op.bar {
						bars = append(bars[:i], bars[i+1:]...)
						ok = true
						b.remove()
						break
					}
				}
				op.result <- ok
			}
		case respCh := <-p.barCountReqCh:
			respCh <- len(bars)
		case beforeRender = <-p.brCh:
		case <-t.C:
			numBars := len(bars)

			if numBars == 0 {
				break
			}

			if beforeRender != nil {
				beforeRender(bars)
			}

			prependWs := newWidthSync(userRR, numBars, bars[0].NumOfPrependers())
			appendWs := newWidthSync(userRR, numBars, bars[0].NumOfAppenders())

			width, _, _ := cwriter.GetTermSize()
			ibars := iBarsGen(bars, width)
			ibbCh := make(chan indexedBarBuffer)
			wg.Add(numBars)
			for i := 0; i < numBars; i++ {
				go func() {
					defer recoverIfPanic()
					drawer(ibars, ibbCh, prependWs, appendWs)
				}()
			}
			go func() {
				wg.Wait()
				close(ibbCh)
				for _, ch := range prependWs.listen {
					close(ch)
				}
				for _, ch := range appendWs.listen {
					close(ch)
				}
			}()

			m := make(map[int][]byte, len(bars))
			for ibb := range ibbCh {
				m[ibb.index] = ibb.buf
			}
			for i := 0; i < len(bars); i++ {
				cw.Write(m[i])
			}

			cw.Flush()

			for _, b := range bars {
				b.flushed()
			}
		case userRR = <-p.rrChangeReqCh:
			t.Stop()
			t = time.NewTicker(userRR)
		case <-p.cancel:
			return
		}
	}
}

func newWidthSync(userRR time.Duration, numBars, numColumn int) *widthSync {
	ws := &widthSync{
		listen: make([]chan int, numColumn),
		result: make([]chan int, numColumn),
	}
	for i := 0; i < numColumn; i++ {
		ws.listen[i] = make(chan int, numBars)
		ws.result[i] = make(chan int, numBars)
	}
	for i := 0; i < numColumn; i++ {
		go func(listenCh <-chan int, resultCh chan<- int) {
			defer close(resultCh)
			widths := make([]int, 0, numBars)
		loop:
			for {
				select {
				case w := <-listenCh:
					widths = append(widths, w)
					if len(widths) == numBars {
						break loop
					}
				case <-time.After(userRR):
					return
				}
			}
			result := max(widths)
			for i := 0; i < numBars; i++ {
				resultCh <- result
			}
		}(ws.listen[i], ws.result[i])
	}
	return ws
}

func drawer(ibars <-chan indexedBar, ibbCh chan<- indexedBarBuffer, prependWs, appendWs *widthSync) {
	for b := range ibars {
		buf := b.bar.bytes(b.termWidth, prependWs, appendWs)
		buf = append(buf, '\n')
		ibbCh <- indexedBarBuffer{b.index, buf}
	}
}

func iBarsGen(bars []*Bar, width int) <-chan indexedBar {
	ibars := make(chan indexedBar)
	go func() {
		defer close(ibars)
		for i, b := range bars {
			ibars <- indexedBar{i, width, b}
		}
	}()
	return ibars
}

// isClosed check if ch closed
// caution see: http://www.tapirgames.com/blog/golang-channel-closing
func isClosed(ch <-chan struct{}) bool {
	select {
	case <-ch:
		return true
	default:
		return false
	}
}

func max(slice []int) int {
	max := slice[0]

	for i := 1; i < len(slice); i++ {
		if slice[i] > max {
			max = slice[i]
		}
	}

	return max
}
<|endoftext|>


<commit_before>package maodou

import (
	"log"
	"net/http"
	"time"

	"github.com/mnhkahn/maodou/dao"
	. "github.com/mnhkahn/maodou/logs"
	"github.com/mnhkahn/maodou/models"
	"github.com/mnhkahn/maodou/supervisor"
)

type Handler interface {
	Init()
	Start()
	Index(resp *Response)
	Detail(resp *Response)
	Result(result *models.Result)
	Config() *HandlerConfig
	Route(http_method, route string, function func(w http.ResponseWriter, req *http.Request))
	Serve(ip string, port int, graceful_timeout int)
}

type HandlerConfig struct {
	dao_name         string
	dsn              string
	cawl_every       time.Duration
	interval         time.Duration
	ip               string
	port             int
	graceful_timeout int
}

type MaoDou struct {
	Dao        dao.DaoContainer
	req        *Request
	resp       *Response
	settings   *HandlerConfig
	supervisor *supervisor.SupervisorController
	Debug      bool
}

func (this *MaoDou) SetRate(times ...time.Duration) {
	if len(times) == 1 {
		this.settings.cawl_every = times[0]
	} else if len(times) == 2 {
		this.settings.cawl_every = times[0]
		this.req.Interval = times[1]
	}
}

func (this *MaoDou) SetServe(ip string, port int, graceful_timeout int) {
	this.settings.ip, this.settings.port, this.settings.graceful_timeout = ip, port, graceful_timeout
}

func (this *MaoDou) SetDao(dao_name, dsn string) {
	var err error
	this.settings.dao_name, this.settings.dsn = dao_name, dsn
	this.Dao, err = dao.NewDao(this.settings.dao_name, this.settings.dsn)
	if err != nil {
		panic(err)
	}
}

func (this *MaoDou) Init() {
	var err error
	this.settings = new(HandlerConfig)
	this.settings.cawl_every = 0
	this.settings.interval = 0
	this.settings.graceful_timeout = 1
	this.supervisor = supervisor.NewSupervisorController()
	this.settings.dao_name = "sqlite"
	this.settings.dsn = "./maodou"
	this.req = NewRequest(0)
	this.Dao, err = dao.NewDao(this.settings.dao_name, this.settings.dsn)
	if err != nil {
		panic(err)
	}
}

func (this *MaoDou) Start() {
	log.Println("Start Method is not override.")
}

func (this *MaoDou) Cawl(url string) (*Response, error) {
	return this.req.Cawl(url)
}

func (this *MaoDou) Index(resp *Response) {
	log.Println("Start Method is not override.")
	this.Detail(nil)
}

func (this *MaoDou) Detail(resp *Response) {
	log.Println("Start Method is not override.")
	this.Result(nil)
}

func (this *MaoDou) Result(result *models.Result) {
	log.Println("Start Method is not override.")
}

func (this *MaoDou) Config() *HandlerConfig {
	return this.settings
}

func (this *MaoDou) Route(http_method, route string, function func(w http.ResponseWriter, req *http.Request)) {
	this.supervisor.Route(http_method, route, function)
}

func (this *MaoDou) Serve(ip string, port int, graceful_timeout int) {
	this.supervisor.Run(ip, port, graceful_timeout)
}

var APP *App

type App struct {
	handler Handler
}

func NewController(handler Handler) *App {
	app := new(App)
	app.handler = handler
	return app
}

func (this *App) Run() error {
	duration := time.Duration(0)
	if this.handler.Config() != nil && this.handler.Config().cawl_every > 0 {
		duration = this.handler.Config().cawl_every
	}
	APP.run()
	if duration > 0 {
		timer := time.NewTicker(duration)
		for {
			select {
			case <-timer.C:
				go func() {
					APP.run()
				}()
			}
		}
	}

	if this.handler.Config().port > 0 {
		this.handler.Serve(this.handler.Config().ip, this.handler.Config().port, this.handler.Config().graceful_timeout)
	}
	return nil
}

func (this *App) run() {
	ColorLog("[INFO] Start parse at %s...\n", time.Now().Format(time.RFC3339))
	this.handler.Start()
	ColorLog("[SUCC] Parse end.\n")
}

func Register(handler Handler) {
	APP = NewController(handler)
	APP.Run()
}
<commit_msg>rm supervisor<commit_after>package maodou

import (
	"log"
	"net/http"
	"time"

	"github.com/mnhkahn/maodou/dao"
	. "github.com/mnhkahn/maodou/logs"
	"github.com/mnhkahn/maodou/models"
)

type Handler interface {
	Init()
	Start()
	Index(resp *Response)
	Detail(resp *Response)
	Result(result *models.Result)
	Config() *HandlerConfig
	Route(http_method, route string, function func(w http.ResponseWriter, req *http.Request))
	Serve(ip string, port int, graceful_timeout int)
}

type HandlerConfig struct {
	dao_name         string
	dsn              string
	cawl_every       time.Duration
	interval         time.Duration
	ip               string
	port             int
	graceful_timeout int
}

type MaoDou struct {
	Dao      dao.DaoContainer
	req      *Request
	resp     *Response
	settings *HandlerConfig
	// supervisor *supervisor.SupervisorController
	Debug bool
}

func (this *MaoDou) SetRate(times ...time.Duration) {
	if len(times) == 1 {
		this.settings.cawl_every = times[0]
	} else if len(times) == 2 {
		this.settings.cawl_every = times[0]
		this.req.Interval = times[1]
	}
}

func (this *MaoDou) SetServe(ip string, port int, graceful_timeout int) {
	this.settings.ip, this.settings.port, this.settings.graceful_timeout = ip, port, graceful_timeout
}

func (this *MaoDou) SetDao(dao_name, dsn string) {
	var err error
	this.settings.dao_name, this.settings.dsn = dao_name, dsn
	this.Dao, err = dao.NewDao(this.settings.dao_name, this.settings.dsn)
	if err != nil {
		panic(err)
	}
}

func (this *MaoDou) Init() {
	var err error
	this.settings = new(HandlerConfig)
	this.settings.cawl_every = 0
	this.settings.interval = 0
	this.settings.graceful_timeout = 1
	// this.supervisor = supervisor.NewSupervisorController()
	this.settings.dao_name = "sqlite"
	this.settings.dsn = "./maodou"
	this.req = NewRequest(0)
	this.Dao, err = dao.NewDao(this.settings.dao_name, this.settings.dsn)
	if err != nil {
		panic(err)
	}
}

func (this *MaoDou) Start() {
	log.Println("Start Method is not override.")
}

func (this *MaoDou) Cawl(url string) (*Response, error) {
	return this.req.Cawl(url)
}

func (this *MaoDou) Index(resp *Response) {
	log.Println("Start Method is not override.")
	this.Detail(nil)
}

func (this *MaoDou) Detail(resp *Response) {
	log.Println("Start Method is not override.")
	this.Result(nil)
}

func (this *MaoDou) Result(result *models.Result) {
	log.Println("Start Method is not override.")
}

func (this *MaoDou) Config() *HandlerConfig {
	return this.settings
}

func (this *MaoDou) Route(http_method, route string, function func(w http.ResponseWriter, req *http.Request)) {
	// this.supervisor.Route(http_method, route, function)
}

func (this *MaoDou) Serve(ip string, port int, graceful_timeout int) {
	// this.supervisor.Run(ip, port, graceful_timeout)
}

var APP *App

type App struct {
	handler Handler
}

func NewController(handler Handler) *App {
	app := new(App)
	app.handler = handler
	return app
}

func (this *App) Run() error {
	duration := time.Duration(0)
	if this.handler.Config() != nil && this.handler.Config().cawl_every > 0 {
		duration = this.handler.Config().cawl_every
	}
	APP.run()
	if duration > 0 {
		timer := time.NewTicker(duration)
		for {
			select {
			case <-timer.C:
				go func() {
					APP.run()
				}()
			}
		}
	}

	if this.handler.Config().port > 0 {
		this.handler.Serve(this.handler.Config().ip, this.handler.Config().port, this.handler.Config().graceful_timeout)
	}
	return nil
}

func (this *App) run() {
	ColorLog("[INFO] Start parse at %s...\n", time.Now().Format(time.RFC3339))
	this.handler.Start()
	ColorLog("[SUCC] Parse end.\n")
}

func Register(handler Handler) {
	APP = NewController(handler)
	APP.Run()
}
<|endoftext|>


<commit_before>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	yaml "gopkg.in/yaml.v2"
)

var (
	identifierRE   = `[a-zA-Z_][a-zA-Z0-9_]+`
	statsdMetricRE = `[a-zA-Z_](-?[a-zA-Z0-9_])+`

	metricLineRE = regexp.MustCompile(`^(\*\.|` + statsdMetricRE + `\.)+(\*|` + statsdMetricRE + `)$`)
	labelLineRE  = regexp.MustCompile(`^(` + identifierRE + `)\s*=\s*"(.*)"$`)
	metricNameRE = regexp.MustCompile(`^` + identifierRE + `$`)
)

type mapperConfigDefaults struct {
	TimerType timerType `yaml:"timer_type"`
	Buckets   []float64 `yaml:"buckets"`
}

type metricMapper struct {
	Defaults mapperConfigDefaults `yaml:"defaults"`
	Mappings []metricMapping      `yaml:"mappings"`
	mutex    sync.Mutex
}

type metricMapping struct {
	Match     string `yaml:"match"`
	regex     *regexp.Regexp
	Labels    prometheus.Labels `yaml:"labels"`
	TimerType timerType         `yaml:"timer_type"`
	Buckets   []float64         `yaml:"buckets"`
}

type configLoadStates int

const (
	SEARCHING configLoadStates = iota
	METRIC_DEFINITION
)

func (m *metricMapper) initFromString(fileContents string) error {
	lines := strings.Split(fileContents, "\n")
	numLines := len(lines)
	state := SEARCHING

	parsedMappings := []metricMapping{}
	currentMapping := metricMapping{Labels: prometheus.Labels{}}
	for i, line := range lines {
		line = strings.TrimSpace(line)

		// skip comments
		if strings.HasPrefix(line, "#") {
			continue
		}

		switch state {
		case SEARCHING:
			if line == "" {
				continue
			}
			if !metricLineRE.MatchString(line) {
				return fmt.Errorf("Line %d: expected metric match line, got: %s", i, line)
			}

			// Translate the glob-style metric match line into a proper regex that we
			// can use to match metrics later on.
			metricRe := strings.Replace(line, ".", "\\.", -1)
			metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
			currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

			state = METRIC_DEFINITION

		case METRIC_DEFINITION:
			if (i == numLines-1) && (line != "") {
				return fmt.Errorf("Line %d: missing terminating newline", i)
			}
			if line == "" {
				if len(currentMapping.Labels) == 0 {
					return fmt.Errorf("Line %d: metric mapping didn't set any labels", i)
				}
				if _, ok := currentMapping.Labels["name"]; !ok {
					return fmt.Errorf("Line %d: metric mapping didn't set a metric name", i)
				}
				parsedMappings = append(parsedMappings, currentMapping)
				state = SEARCHING
				currentMapping = metricMapping{Labels: prometheus.Labels{}}
				continue
			}
			if err := m.updateMapping(line, i, &currentMapping); err != nil {
				return err
			}
		default:
			panic("illegal state")
		}
	}

	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.Mappings = parsedMappings

	mappingsCount.Set(float64(len(parsedMappings)))

	return nil
}

func (m *metricMapper) initFromYAMLString(fileContents string) error {
	var n metricMapper

	if err := yaml.Unmarshal([]byte(fileContents), &n); err != nil {
		return err
	}

	if n.Defaults.Buckets == nil || len(n.Defaults.Buckets) == 0 {
		n.Defaults.Buckets = prometheus.DefBuckets
	}

	for i := range n.Mappings {
		currentMapping := &n.Mappings[i]

		if !metricLineRE.MatchString(currentMapping.Match) {
			return fmt.Errorf("invalid match: %s", currentMapping.Match)
		}

		// check that label is correct
		for k, v := range currentMapping.Labels {
			label := fmt.Sprintf("%s=%q", k, v)
			if len(labelLineRE.FindStringSubmatch(label)) != 3 {
				return fmt.Errorf("invalid label: %s: %s", k, v)
			}
			if k == "name" && !metricNameRE.MatchString(v) {
				return fmt.Errorf("metric name '%s' doesn't match regex '%s'", v, metricNameRE)
			}
		}

		if _, ok := currentMapping.Labels["name"]; !ok {
			return fmt.Errorf("Line %d: metric mapping didn't set a metric name", i)
		}

		// Translate the glob-style metric match line into a proper regex that we
		// can use to match metrics later on.
		metricRe := strings.Replace(currentMapping.Match, ".", "\\.", -1)
		metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
		currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

		if currentMapping.TimerType == "" {
			currentMapping.TimerType = n.Defaults.TimerType
		}

		if currentMapping.Buckets == nil || len(currentMapping.Buckets) == 0 {
			currentMapping.Buckets = n.Defaults.Buckets
		}

	}

	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.Defaults = n.Defaults
	m.Mappings = n.Mappings

	mappingsCount.Set(float64(len(n.Mappings)))

	return nil
}

func (m *metricMapper) initFromFile(fileName string) error {
	mappingStr, err := ioutil.ReadFile(fileName)
	if err != nil {
		return err
	}
	switch strings.ToLower(filepath.Ext(fileName)) {
	case ".yaml", ".yml":
		return m.initFromYAMLString(string(mappingStr))
	default:
		return m.initFromString(string(mappingStr))
	}
}

func (m *metricMapper) getMapping(statsdMetric string) (*metricMapping, prometheus.Labels, bool) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	for _, mapping := range m.Mappings {
		matches := mapping.regex.FindStringSubmatchIndex(statsdMetric)
		if len(matches) == 0 {
			continue
		}

		labels := prometheus.Labels{}
		for label, valueExpr := range mapping.Labels {
			value := mapping.regex.ExpandString([]byte{}, valueExpr, statsdMetric, matches)
			labels[label] = string(value)
		}
		return &mapping, labels, true
	}

	return nil, nil, false
}

func (m *metricMapper) updateMapping(line string, i int, mapping *metricMapping) error {
	matches := labelLineRE.FindStringSubmatch(line)
	if len(matches) != 3 {
		return fmt.Errorf("Line %d: expected label mapping line, got: %s", i, line)
	}
	label, value := matches[1], matches[2]
	if label == "name" && !metricNameRE.MatchString(value) {
		return fmt.Errorf("Line %d: metric name '%s' doesn't match regex '%s'", i, value, metricNameRE)
	}
	(*mapping).Labels[label] = value
	return nil
}
<commit_msg>Style fixes<commit_after>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	yaml "gopkg.in/yaml.v2"
)

var (
	identifierRE   = `[a-zA-Z_][a-zA-Z0-9_]+`
	statsdMetricRE = `[a-zA-Z_](-?[a-zA-Z0-9_])+`

	metricLineRE = regexp.MustCompile(`^(\*\.|` + statsdMetricRE + `\.)+(\*|` + statsdMetricRE + `)$`)
	labelLineRE  = regexp.MustCompile(`^(` + identifierRE + `)\s*=\s*"(.*)"$`)
	metricNameRE = regexp.MustCompile(`^` + identifierRE + `$`)
)

type mapperConfigDefaults struct {
	TimerType timerType `yaml:"timer_type"`
	Buckets   []float64 `yaml:"buckets"`
}

type metricMapper struct {
	Defaults mapperConfigDefaults `yaml:"defaults"`
	Mappings []metricMapping      `yaml:"mappings"`
	mutex    sync.Mutex
}

type metricMapping struct {
	Match     string `yaml:"match"`
	regex     *regexp.Regexp
	Labels    prometheus.Labels `yaml:"labels"`
	TimerType timerType         `yaml:"timer_type"`
	Buckets   []float64         `yaml:"buckets"`
}

type configLoadStates int

const (
	SEARCHING configLoadStates = iota
	METRIC_DEFINITION
)

func (m *metricMapper) initFromString(fileContents string) error {
	lines := strings.Split(fileContents, "\n")
	numLines := len(lines)
	state := SEARCHING

	parsedMappings := []metricMapping{}
	currentMapping := metricMapping{Labels: prometheus.Labels{}}
	for i, line := range lines {
		line = strings.TrimSpace(line)

		// skip comments
		if strings.HasPrefix(line, "#") {
			continue
		}

		switch state {
		case SEARCHING:
			if line == "" {
				continue
			}
			if !metricLineRE.MatchString(line) {
				return fmt.Errorf("line %d: expected metric match line, got: %s", i, line)
			}

			// Translate the glob-style metric match line into a proper regex that we
			// can use to match metrics later on.
			metricRe := strings.Replace(line, ".", "\\.", -1)
			metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
			currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

			state = METRIC_DEFINITION

		case METRIC_DEFINITION:
			if (i == numLines-1) && (line != "") {
				return fmt.Errorf("Line %d: missing terminating newline", i)
			}
			if line == "" {
				if len(currentMapping.Labels) == 0 {
					return fmt.Errorf("Line %d: metric mapping didn't set any labels", i)
				}
				if _, ok := currentMapping.Labels["name"]; !ok {
					return fmt.Errorf("Line %d: metric mapping didn't set a metric name", i)
				}
				parsedMappings = append(parsedMappings, currentMapping)
				state = SEARCHING
				currentMapping = metricMapping{Labels: prometheus.Labels{}}
				continue
			}
			if err := m.updateMapping(line, i, &currentMapping); err != nil {
				return err
			}
		default:
			panic("illegal state")
		}
	}

	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.Mappings = parsedMappings

	mappingsCount.Set(float64(len(parsedMappings)))

	return nil
}

func (m *metricMapper) initFromYAMLString(fileContents string) error {
	var n metricMapper

	if err := yaml.Unmarshal([]byte(fileContents), &n); err != nil {
		return err
	}

	if n.Defaults.Buckets == nil || len(n.Defaults.Buckets) == 0 {
		n.Defaults.Buckets = prometheus.DefBuckets
	}

	for i := range n.Mappings {
		currentMapping := &n.Mappings[i]

		if !metricLineRE.MatchString(currentMapping.Match) {
			return fmt.Errorf("invalid match: %s", currentMapping.Match)
		}

		// Check that label is correct.
		for k, v := range currentMapping.Labels {
			label := fmt.Sprintf("%s=%q", k, v)
			if len(labelLineRE.FindStringSubmatch(label)) != 3 {
				return fmt.Errorf("invalid label: %s: %s", k, v)
			}
			if k == "name" && !metricNameRE.MatchString(v) {
				return fmt.Errorf("metric name '%s' doesn't match regex '%s'", v, metricNameRE)
			}
		}

		if _, ok := currentMapping.Labels["name"]; !ok {
			return fmt.Errorf("line %d: metric mapping didn't set a metric name", i)
		}

		// Translate the glob-style metric match line into a proper regex that we
		// can use to match metrics later on.
		metricRe := strings.Replace(currentMapping.Match, ".", "\\.", -1)
		metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
		currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

		if currentMapping.TimerType == "" {
			currentMapping.TimerType = n.Defaults.TimerType
		}

		if currentMapping.Buckets == nil || len(currentMapping.Buckets) == 0 {
			currentMapping.Buckets = n.Defaults.Buckets
		}

	}

	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.Defaults = n.Defaults
	m.Mappings = n.Mappings

	mappingsCount.Set(float64(len(n.Mappings)))

	return nil
}

func (m *metricMapper) initFromFile(fileName string) error {
	mappingStr, err := ioutil.ReadFile(fileName)
	if err != nil {
		return err
	}
	switch strings.ToLower(filepath.Ext(fileName)) {
	case ".yaml", ".yml":
		return m.initFromYAMLString(string(mappingStr))
	default:
		return m.initFromString(string(mappingStr))
	}
}

func (m *metricMapper) getMapping(statsdMetric string) (*metricMapping, prometheus.Labels, bool) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	for _, mapping := range m.Mappings {
		matches := mapping.regex.FindStringSubmatchIndex(statsdMetric)
		if len(matches) == 0 {
			continue
		}

		labels := prometheus.Labels{}
		for label, valueExpr := range mapping.Labels {
			value := mapping.regex.ExpandString([]byte{}, valueExpr, statsdMetric, matches)
			labels[label] = string(value)
		}
		return &mapping, labels, true
	}

	return nil, nil, false
}

func (m *metricMapper) updateMapping(line string, i int, mapping *metricMapping) error {
	matches := labelLineRE.FindStringSubmatch(line)
	if len(matches) != 3 {
		return fmt.Errorf("Line %d: expected label mapping line, got: %s", i, line)
	}
	label, value := matches[1], matches[2]
	if label == "name" && !metricNameRE.MatchString(value) {
		return fmt.Errorf("Line %d: metric name '%s' doesn't match regex '%s'", i, value, metricNameRE)
	}
	(*mapping).Labels[label] = value
	return nil
}
<|endoftext|>


<commit_before>package views

import (
	"encoding/json"
	"html/template"
	"net/http"
	"strconv"

	"github.com/antihax/evedata/models"
	"github.com/antihax/evedata/server"
	"github.com/antihax/evedata/templates"
)

func init() {
	evedata.AddRoute("arbitrageCalculator", "GET", "/arbitrageCalculator", arbitrageCalculatorPage)
	evedata.AddRoute("arbitrageCalculatorStations", "GET", "/J/arbitrageCalculatorStations", arbitrageCalculatorStations)
	evedata.AddRoute("arbitrageCalculator", "GET", "/J/arbitrageCalculator", arbitrageCalculator)
}

func arbitrageCalculatorPage(w http.ResponseWriter, r *http.Request) {
	setCache(w, 60*60*24)
	p := newPage(r, "Arbitrage Calculator")

	templates.Templates = template.Must(template.ParseFiles("templates/arbitrageCalculator.html", templates.LayoutPath))
	err := templates.Templates.ExecuteTemplate(w, "base", p)

	if err != nil {
		httpErr(w, err)
		return
	}
}

func arbitrageCalculatorStations(w http.ResponseWriter, r *http.Request) {
	setCache(w, 60*30)
	v, err := models.GetArbitrageCalculatorStations()
	if err != nil {
		httpErr(w, err)
		return
	}

	encoder := json.NewEncoder(w)
	encoder.Encode(v)
}

func arbitrageCalculator(w http.ResponseWriter, r *http.Request) {
	setCache(w, 60*30)

	stationID, err := strconv.ParseInt(r.FormValue("stationID"), 10, 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	minVolume, err := strconv.ParseInt(r.FormValue("minVolume"), 10, 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	maxPrice, err := strconv.ParseInt(r.FormValue("maxPrice"), 10, 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	brokersFee, err := strconv.ParseFloat(r.FormValue("brokersFee"), 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	tax, err := strconv.ParseFloat(r.FormValue("tax"), 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	method := r.FormValue("method")

	brokersFee = brokersFee / 100
	tax = tax / 100

	v, err := models.GetArbitrageCalculator(stationID, minVolume, maxPrice, brokersFee, tax, method)
	if err != nil {
		httpErr(w, err)
		return
	}

	encoder := json.NewEncoder(w)
	encoder.Encode(v)
}
<commit_msg>simplify<commit_after>package views

import (
	"encoding/json"
	"html/template"
	"net/http"
	"strconv"

	"github.com/antihax/evedata/models"
	"github.com/antihax/evedata/server"
	"github.com/antihax/evedata/templates"
)

func init() {
	evedata.AddRoute("arbitrageCalculator", "GET", "/arbitrageCalculator", arbitrageCalculatorPage)
	evedata.AddRoute("arbitrageCalculatorStations", "GET", "/J/arbitrageCalculatorStations", arbitrageCalculatorStations)
	evedata.AddRoute("arbitrageCalculator", "GET", "/J/arbitrageCalculator", arbitrageCalculator)
}

func arbitrageCalculatorPage(w http.ResponseWriter, r *http.Request) {
	setCache(w, 60*60*24)
	p := newPage(r, "Arbitrage Calculator")

	templates.Templates = template.Must(template.ParseFiles("templates/arbitrageCalculator.html", templates.LayoutPath))
	err := templates.Templates.ExecuteTemplate(w, "base", p)

	if err != nil {
		httpErr(w, err)
		return
	}
}

func arbitrageCalculatorStations(w http.ResponseWriter, r *http.Request) {
	setCache(w, 60*30)
	v, err := models.GetArbitrageCalculatorStations()
	if err != nil {
		httpErr(w, err)
		return
	}

	err = json.NewEncoder(w).Encode(v)
	if err != nil {
		httpErr(w, err)
		return
	}
}

func arbitrageCalculator(w http.ResponseWriter, r *http.Request) {
	setCache(w, 60*30)

	stationID, err := strconv.ParseInt(r.FormValue("stationID"), 10, 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	minVolume, err := strconv.ParseInt(r.FormValue("minVolume"), 10, 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	maxPrice, err := strconv.ParseInt(r.FormValue("maxPrice"), 10, 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	brokersFee, err := strconv.ParseFloat(r.FormValue("brokersFee"), 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	tax, err := strconv.ParseFloat(r.FormValue("tax"), 64)
	if err != nil {
		httpErr(w, err)
		return
	}

	method := r.FormValue("method")

	brokersFee = brokersFee / 100
	tax = tax / 100

	v, err := models.GetArbitrageCalculator(stationID, minVolume, maxPrice, brokersFee, tax, method)
	if err != nil {
		httpErr(w, err)
		return
	}

	encoder := json.NewEncoder(w)
	encoder.Encode(v)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"fmt"
	Q "github.com/PuerkitoBio/goquery"
	"github.com/keybase/go-libkb"
	"github.com/kr/text"
	"io"
	"os"
	"regexp"
	"strings"
)

var INDENT int = 3

type Paragraph struct {
	data   []byte
	indent int
	cols   int
	prefix string
}

// Buffer adds data to the internal paragraph buffer.
func (p *Paragraph) Buffer(b []byte) {
	p.data = append(p.data, b...)
}

var (
	nl = []byte{'\n'}
	sp = []byte{' '}
)

// makePad makes a whitespace pad that is l bytes long.
func makePad(l int) []byte {
	ret := make([]byte, l)
	for i := 0; i < l; i++ {
		ret[i] = ' '
	}
	return ret
}

// spacify replaces arbitrary strings of whitespace with
// a single ' ' character. Also strips off leading and trailing
// whitespace.
func spacify(s string) string {
	v := regexp.MustCompile(`[[:space:]]+`).Split(s, -1)
	if len(v) > 0 && v[0] == "" {
		v = v[1:]
	}
	l := len(v)
	if l > 0 && v[l-1] == "" {
		v = v[0:(l - 1)]
	}
	return strings.Join(v, " ")
}

// Output a paragraph to the io.Writer, applying the proper
// formatting.
func (p Paragraph) Output(out io.Writer) {
	s := []byte(spacify(string(p.data)))
	if len(s) == 0 {
		out.Write(nl)
		return
	}
	indent := p.indent * INDENT
	width := p.cols - indent - len(p.prefix)
	wrapped := text.WrapBytes(s, width)
	lines := bytes.Split(wrapped, nl)
	gutter := makePad(indent)
	pad := makePad(len(p.prefix))

	for i, line := range lines {
		out.Write(gutter)
		if i == 0 {
			out.Write([]byte(p.prefix))
		} else {
			out.Write(pad)
		}
		out.Write(line)
		out.Write(nl)
	}
}

type Renderer struct {
	indent    int
	cols      int
	out       io.Writer
	paragraph *Paragraph
}

func NewRenderer(out io.Writer) *Renderer {
	width, _ := G_UI.GetTerminalSize()
	if width == 0 {
		width = 80
	}
	return &Renderer{indent: 0, out: out, cols: width}
}

func (r *Renderer) RenderNodes(nodes []*html.Node) {
	for _, n := range nodes {
		r.RenderNode(n)
	}
}

func (r *Renderer) Buffer(d []byte) {
	if r.paragraph == nil {
		d = bytes.TrimSpace(d)
		if len(d) > 0 {
			G.Log.Warning("floating data in Markup is ignored: %v", d)
		}
	} else {
		r.paragraph.Buffer(d)
	}
}

func (r *Renderer) NewParagraph(prefix string) *Paragraph {
	r.FlushParagraph()
	r.paragraph = &Paragraph{indent: r.indent, cols: r.cols, prefix: prefix}
	return r.paragraph
}

func (r *Renderer) FlushParagraph() {
	if r.paragraph != nil {
		r.paragraph.Output(r.out)
		r.paragraph = nil
	}
}

func GetNodeAttrVal(node *html.Node, which string) *string {
	for _, attr := range node.Attr {
		if attr.Key == which {
			return &attr.Val
		}
	}
	return nil
}

func (r *Renderer) RenderNode(node *html.Node) {
	if node.Type == html.TextNode {
		r.Buffer([]byte(node.Data))
		return
	}

	var cp *CodePair

	if node.DataAtom == atom.Ul {
		r.indent++
	} else if node.DataAtom == atom.Li {
		var bullet string
		if bp := GetNodeAttrVal(node, "bullet"); bp != nil {
			bullet = *bp
		} else {
			bullet = "• "
		}
		r.NewParagraph(bullet)
	} else if node.DataAtom == atom.P {
		r.NewParagraph("")
	} else if node.DataAtom == atom.Strong {
		cp = &CpBold
	} else if node.DataAtom == atom.Em {
		cp = &CpItalic
	} else if node.Data == "url" {
		cp = &CpUnderline
	} else if node.Data == "color" {
		if c := GetNodeAttrVal(node, "name"); c != nil {
			cp = GetColorCode(*c)
		}
	}

	if cp != nil {
		r.Buffer(cp.OpenBytes())
	}

	if node.FirstChild != nil {
		for c := node.FirstChild; c != nil; c = c.NextSibling {
			r.RenderNode(c)
		}
	}

	if node.DataAtom == atom.Ul {
		r.indent--
	}
	if node.DataAtom == atom.Li || node.DataAtom == atom.P {
		r.FlushParagraph()
	}
	if cp != nil {
		r.Buffer(cp.CloseBytes())
	}
}

func Render(w io.Writer, m *libkb.Markup) {
	if m == nil {
		return
	}
	doc, err := Q.NewDocumentFromReader(m.ToReader())
	if err != nil {
		fmt.Printf("Cannot render markup: %s\n", err.Error())
		return
	}
	renderer := NewRenderer(os.Stdout)
	renderer.RenderNodes(doc.Nodes)
}
<commit_msg>recompile<commit_after>package main

import (
	"bytes"
	"fmt"
	Q "github.com/PuerkitoBio/goquery"
	"github.com/keybase/go-libkb"
	"github.com/kr/text"
	"golang.org/x/net/html"
	"golang.org/x/net/html/atom"
	"io"
	"os"
	"regexp"
	"strings"
)

var INDENT int = 3

type Paragraph struct {
	data   []byte
	indent int
	cols   int
	prefix string
}

// Buffer adds data to the internal paragraph buffer.
func (p *Paragraph) Buffer(b []byte) {
	p.data = append(p.data, b...)
}

var (
	nl = []byte{'\n'}
	sp = []byte{' '}
)

// makePad makes a whitespace pad that is l bytes long.
func makePad(l int) []byte {
	ret := make([]byte, l)
	for i := 0; i < l; i++ {
		ret[i] = ' '
	}
	return ret
}

// spacify replaces arbitrary strings of whitespace with
// a single ' ' character. Also strips off leading and trailing
// whitespace.
func spacify(s string) string {
	v := regexp.MustCompile(`[[:space:]]+`).Split(s, -1)
	if len(v) > 0 && v[0] == "" {
		v = v[1:]
	}
	l := len(v)
	if l > 0 && v[l-1] == "" {
		v = v[0:(l - 1)]
	}
	return strings.Join(v, " ")
}

// Output a paragraph to the io.Writer, applying the proper
// formatting.
func (p Paragraph) Output(out io.Writer) {
	s := []byte(spacify(string(p.data)))
	if len(s) == 0 {
		out.Write(nl)
		return
	}
	indent := p.indent * INDENT
	width := p.cols - indent - len(p.prefix)
	wrapped := text.WrapBytes(s, width)
	lines := bytes.Split(wrapped, nl)
	gutter := makePad(indent)
	pad := makePad(len(p.prefix))

	for i, line := range lines {
		out.Write(gutter)
		if i == 0 {
			out.Write([]byte(p.prefix))
		} else {
			out.Write(pad)
		}
		out.Write(line)
		out.Write(nl)
	}
}

type Renderer struct {
	indent    int
	cols      int
	out       io.Writer
	paragraph *Paragraph
}

func NewRenderer(out io.Writer) *Renderer {
	width, _ := G_UI.GetTerminalSize()
	if width == 0 {
		width = 80
	}
	return &Renderer{indent: 0, out: out, cols: width}
}

func (r *Renderer) RenderNodes(nodes []*html.Node) {
	for _, n := range nodes {
		r.RenderNode(n)
	}
}

func (r *Renderer) Buffer(d []byte) {
	if r.paragraph == nil {
		d = bytes.TrimSpace(d)
		if len(d) > 0 {
			G.Log.Warning("floating data in Markup is ignored: %v", d)
		}
	} else {
		r.paragraph.Buffer(d)
	}
}

func (r *Renderer) NewParagraph(prefix string) *Paragraph {
	r.FlushParagraph()
	r.paragraph = &Paragraph{indent: r.indent, cols: r.cols, prefix: prefix}
	return r.paragraph
}

func (r *Renderer) FlushParagraph() {
	if r.paragraph != nil {
		r.paragraph.Output(r.out)
		r.paragraph = nil
	}
}

func GetNodeAttrVal(node *html.Node, which string) *string {
	for _, attr := range node.Attr {
		if attr.Key == which {
			return &attr.Val
		}
	}
	return nil
}

func (r *Renderer) RenderNode(node *html.Node) {
	if node.Type == html.TextNode {
		r.Buffer([]byte(node.Data))
		return
	}

	var cp *CodePair

	if node.DataAtom == atom.Ul {
		r.indent++
	} else if node.DataAtom == atom.Li {
		var bullet string
		if bp := GetNodeAttrVal(node, "bullet"); bp != nil {
			bullet = *bp
		} else {
			bullet = "• "
		}
		r.NewParagraph(bullet)
	} else if node.DataAtom == atom.P {
		r.NewParagraph("")
	} else if node.DataAtom == atom.Strong {
		cp = &CpBold
	} else if node.DataAtom == atom.Em {
		cp = &CpItalic
	} else if node.Data == "url" {
		cp = &CpUnderline
	} else if node.Data == "color" {
		if c := GetNodeAttrVal(node, "name"); c != nil {
			cp = GetColorCode(*c)
		}
	}

	if cp != nil {
		r.Buffer(cp.OpenBytes())
	}

	if node.FirstChild != nil {
		for c := node.FirstChild; c != nil; c = c.NextSibling {
			r.RenderNode(c)
		}
	}

	if node.DataAtom == atom.Ul {
		r.indent--
	}
	if node.DataAtom == atom.Li || node.DataAtom == atom.P {
		r.FlushParagraph()
	}
	if cp != nil {
		r.Buffer(cp.CloseBytes())
	}
}

func Render(w io.Writer, m *libkb.Markup) {
	if m == nil {
		return
	}
	doc, err := Q.NewDocumentFromReader(m.ToReader())
	if err != nil {
		fmt.Printf("Cannot render markup: %s\n", err.Error())
		return
	}
	renderer := NewRenderer(os.Stdout)
	renderer.RenderNodes(doc.Nodes)
}
<|endoftext|>