<commit_before>/*
 *  Copyright (c) 2011 NeuStar, Inc.
 *  All rights reserved.
 *
 *  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.
 *
 *  NeuStar, the Neustar logo and related names and logos are registered
 *  trademarks, service marks or tradenames of NeuStar, Inc. All other
 *  product names, company names, marks, logos and symbols may be trademarks
 *  of their respective owners.
 */

package kafka

import (
  "encoding/binary"
  "errors"
  "io"
  "log"
  "net"
  "time"
)

type BrokerConsumer struct {
  broker  *Broker
  offset  uint64
  maxSize uint32
  codecs  map[byte]PayloadCodec
}

// Create a new broker consumer
// hostname - host and optionally port, delimited by ':'
// topic to consume
// partition to consume from
// offset to start consuming from
// maxSize (in bytes) of the message to consume (this should be at least as big as the biggest message to be published)
func NewBrokerConsumer(hostname string, topic string, partition int, offset uint64, maxSize uint32) *BrokerConsumer {
  return &BrokerConsumer{broker: newBroker(hostname, topic, partition),
    offset:  offset,
    maxSize: maxSize,
    codecs:  DefaultCodecsMap}
}

// Simplified consumer that defaults the offset and maxSize to 0.
// hostname - host and optionally port, delimited by ':'
// topic to consume
// partition to consume from
func NewBrokerOffsetConsumer(hostname string, topic string, partition int) *BrokerConsumer {
  return &BrokerConsumer{broker: newBroker(hostname, topic, partition),
    offset:  0,
    maxSize: 0,
    codecs:  DefaultCodecsMap}
}

// Add Custom Payload Codecs for Consumer Decoding
// payloadCodecs - an array of PayloadCodec implementations
func (consumer *BrokerConsumer) AddCodecs(payloadCodecs []PayloadCodec) {
  // merge to the default map, so one 'could' override the default codecs..
  for k, v := range codecsMap(payloadCodecs) {
    consumer.codecs[k] = v
  }
}

func (consumer *BrokerConsumer) ConsumeOnChannel(msgChan chan *Message, pollTimeoutMs int64, quit chan bool) (int, error) {
  conn, err := consumer.broker.connect()
  if err != nil {
    return -1, err
  }

  num := 0
  done := make(chan bool, 1)
  go func() {
    for {
      _, err := consumer.consumeWithConn(conn, func(msg *Message) {
        msgChan <- msg
        num += 1
      })

      if err != nil {
        if err != io.EOF {
          log.Println("Fatal Error: ", err)
          panic(err)
        }
        quit <- true // force quit
        break
      }
      time.Sleep(time.Millisecond * time.Duration(pollTimeoutMs))
    }
    done <- true
  }()
  // wait to be told to stop..
  <-quit
  conn.Close()
  close(msgChan)
  <-done
  return num, err
}

type MessageHandlerFunc func(msg *Message)

func (consumer *BrokerConsumer) Consume(handlerFunc MessageHandlerFunc) (int, error) {
  conn, err := consumer.broker.connect()
  if err != nil {
    return -1, err
  }
  defer conn.Close()

  num, err := consumer.consumeWithConn(conn, handlerFunc)

  if err != nil {
    log.Println("Fatal Error: ", err)
  }

  return num, err
}

func (consumer *BrokerConsumer) consumeWithConn(conn *net.TCPConn, handlerFunc MessageHandlerFunc) (int, error) {
  _, err := conn.Write(consumer.broker.EncodeConsumeRequest(consumer.offset, consumer.maxSize))
  if err != nil {
    return -1, err
  }

  length, payload, err := consumer.broker.readResponse(conn)

  if err != nil {
    return -1, err
  }

  num := 0
  if length > 2 {
    // parse out the messages
    var currentOffset uint64 = 0
    for currentOffset <= uint64(length-4) {
      totalLength, msgs := Decode(payload[currentOffset:], consumer.codecs)
      if msgs == nil {
        // update the broker's offset for next consumption incase they want to skip this message and keep going
        consumer.offset += currentOffset
        return num, errors.New("Error Decoding Message")
      }
      msgOffset := consumer.offset + currentOffset
      for _, msg := range msgs {
        // update all of the messages offset
        // multiple messages can be at the same offset (compressed for example)
        msg.offset = msgOffset
        handlerFunc(&msg)
        num += 1
      }
      currentOffset += uint64(4 + totalLength)
    }
    // update the broker's offset for next consumption
    consumer.offset += currentOffset
  }

  return num, err
}

// Get a list of valid offsets (up to maxNumOffsets) before the given time, where
// time is in milliseconds (-1, from the latest offset available, -2 from the smallest offset available)
// The result is a list of offsets, in descending order.
func (consumer *BrokerConsumer) GetOffsets(time int64, maxNumOffsets uint32) ([]uint64, error) {
  offsets := make([]uint64, 0)

  conn, err := consumer.broker.connect()
  if err != nil {
    return offsets, err
  }

  defer conn.Close()

  _, err = conn.Write(consumer.broker.EncodeOffsetRequest(time, maxNumOffsets))
  if err != nil {
    return offsets, err
  }

  length, payload, err := consumer.broker.readResponse(conn)
  if err != nil {
    return offsets, err
  }

  if length > 4 {
    // get the number of offsets
    numOffsets := binary.BigEndian.Uint32(payload[0:])
    var currentOffset uint64 = 4
    for currentOffset < uint64(length-4) && uint32(len(offsets)) < numOffsets {
      offset := binary.BigEndian.Uint64(payload[currentOffset:])
      offsets = append(offsets, offset)
      currentOffset += 8 // offset size
    }
  }

  return offsets, err
}
<commit_msg>Adding a func to return the current offset for a broker.<commit_after>/*
 *  Copyright (c) 2011 NeuStar, Inc.
 *  All rights reserved.
 *
 *  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.
 *
 *  NeuStar, the Neustar logo and related names and logos are registered
 *  trademarks, service marks or tradenames of NeuStar, Inc. All other
 *  product names, company names, marks, logos and symbols may be trademarks
 *  of their respective owners.
 */

package kafka

import (
  "encoding/binary"
  "errors"
  "io"
  "log"
  "net"
  "time"
)

type BrokerConsumer struct {
  broker  *Broker
  offset  uint64
  maxSize uint32
  codecs  map[byte]PayloadCodec
}

// Create a new broker consumer
// hostname - host and optionally port, delimited by ':'
// topic to consume
// partition to consume from
// offset to start consuming from
// maxSize (in bytes) of the message to consume (this should be at least as big as the biggest message to be published)
func NewBrokerConsumer(hostname string, topic string, partition int, offset uint64, maxSize uint32) *BrokerConsumer {
  return &BrokerConsumer{broker: newBroker(hostname, topic, partition),
    offset:  offset,
    maxSize: maxSize,
    codecs:  DefaultCodecsMap}
}

// Simplified consumer that defaults the offset and maxSize to 0.
// hostname - host and optionally port, delimited by ':'
// topic to consume
// partition to consume from
func NewBrokerOffsetConsumer(hostname string, topic string, partition int) *BrokerConsumer {
  return &BrokerConsumer{broker: newBroker(hostname, topic, partition),
    offset:  0,
    maxSize: 0,
    codecs:  DefaultCodecsMap}
}

// Add Custom Payload Codecs for Consumer Decoding
// payloadCodecs - an array of PayloadCodec implementations
func (consumer *BrokerConsumer) AddCodecs(payloadCodecs []PayloadCodec) {
  // merge to the default map, so one 'could' override the default codecs..
  for k, v := range codecsMap(payloadCodecs) {
    consumer.codecs[k] = v
  }
}

func (consumer *BrokerConsumer) ConsumeOnChannel(msgChan chan *Message, pollTimeoutMs int64, quit chan bool) (int, error) {
  conn, err := consumer.broker.connect()
  if err != nil {
    return -1, err
  }

  num := 0
  done := make(chan bool, 1)
  go func() {
    for {
      _, err := consumer.consumeWithConn(conn, func(msg *Message) {
        msgChan <- msg
        num += 1
      })

      if err != nil {
        if err != io.EOF {
          log.Println("Fatal Error: ", err)
          panic(err)
        }
        quit <- true // force quit
        break
      }
      time.Sleep(time.Millisecond * time.Duration(pollTimeoutMs))
    }
    done <- true
  }()
  // wait to be told to stop..
  <-quit
  conn.Close()
  close(msgChan)
  <-done
  return num, err
}

type MessageHandlerFunc func(msg *Message)

func (consumer *BrokerConsumer) Consume(handlerFunc MessageHandlerFunc) (int, error) {
  conn, err := consumer.broker.connect()
  if err != nil {
    return -1, err
  }
  defer conn.Close()

  num, err := consumer.consumeWithConn(conn, handlerFunc)

  if err != nil {
    log.Println("Fatal Error: ", err)
  }

  return num, err
}

func (consumer *BrokerConsumer) consumeWithConn(conn *net.TCPConn, handlerFunc MessageHandlerFunc) (int, error) {
  _, err := conn.Write(consumer.broker.EncodeConsumeRequest(consumer.offset, consumer.maxSize))
  if err != nil {
    return -1, err
  }

  length, payload, err := consumer.broker.readResponse(conn)

  if err != nil {
    return -1, err
  }

  num := 0
  if length > 2 {
    // parse out the messages
    var currentOffset uint64 = 0
    for currentOffset <= uint64(length-4) {
      totalLength, msgs := Decode(payload[currentOffset:], consumer.codecs)
      if msgs == nil {
        // update the broker's offset for next consumption incase they want to skip this message and keep going
        consumer.offset += currentOffset
        return num, errors.New("Error Decoding Message")
      }
      msgOffset := consumer.offset + currentOffset
      for _, msg := range msgs {
        // update all of the messages offset
        // multiple messages can be at the same offset (compressed for example)
        msg.offset = msgOffset
        handlerFunc(&msg)
        num += 1
      }
      currentOffset += uint64(4 + totalLength)
    }
    // update the broker's offset for next consumption
    consumer.offset += currentOffset
  }

  return num, err
}

// Get a list of valid offsets (up to maxNumOffsets) before the given time, where
// time is in milliseconds (-1, from the latest offset available, -2 from the smallest offset available)
// The result is a list of offsets, in descending order.
func (consumer *BrokerConsumer) GetOffsets(time int64, maxNumOffsets uint32) ([]uint64, error) {
  offsets := make([]uint64, 0)

  conn, err := consumer.broker.connect()
  if err != nil {
    return offsets, err
  }

  defer conn.Close()

  _, err = conn.Write(consumer.broker.EncodeOffsetRequest(time, maxNumOffsets))
  if err != nil {
    return offsets, err
  }

  length, payload, err := consumer.broker.readResponse(conn)
  if err != nil {
    return offsets, err
  }

  if length > 4 {
    // get the number of offsets
    numOffsets := binary.BigEndian.Uint32(payload[0:])
    var currentOffset uint64 = 4
    for currentOffset < uint64(length-4) && uint32(len(offsets)) < numOffsets {
      offset := binary.BigEndian.Uint64(payload[currentOffset:])
      offsets = append(offsets, offset)
      currentOffset += 8 // offset size
    }
  }

  return offsets, err
}

// Get the current offset for a broker.
func (consumer *BrokerConsumer) GetOffset() (uint64) {
  return consumer.offset
}
<|endoftext|>


<commit_before>package sarama

// OffsetMethod is passed in ConsumerConfig to tell the consumer how to determine the starting offset.
type OffsetMethod int

const (
	// OffsetMethodManual causes the consumer to interpret the OffsetValue in the ConsumerConfig as the
	// offset at which to start, allowing the user to manually specify their desired starting offset.
	OffsetMethodManual OffsetMethod = iota
	// OffsetMethodNewest causes the consumer to start at the most recent available offset, as
	// determined by querying the broker.
	OffsetMethodNewest
	// OffsetMethodOldest causes the consumer to start at the oldest available offset, as
	// determined by querying the broker.
	OffsetMethodOldest
)

// ConsumerConfig is used to pass multiple configuration options to NewConsumer.
type ConsumerConfig struct {
	// The default (maximum) amount of data to fetch from the broker in each request. The default of 0 is treated as 1024 bytes.
	DefaultFetchSize int32
	// The minimum amount of data to fetch in a request - the broker will wait until at least this many bytes are available.
	// The default of 0 is treated as 'at least one' to prevent the consumer from spinning when no messages are available.
	MinFetchSize int32
	// The maximum permittable message size - messages larger than this will return MessageTooLarge. The default of 0 is
	// treated as no limit.
	MaxMessageSize int32
	// The maximum amount of time (in ms) the broker will wait for MinFetchSize bytes to become available before it
	// returns fewer than that anyways. The default of 0 causes Kafka to return immediately, which is rarely desirable
	// as it causes the Consumer to spin when no events are available. 100-500ms is a reasonable range for most cases.
	MaxWaitTime int32

	// The method used to determine at which offset to begin consuming messages.
	OffsetMethod OffsetMethod
	// Interpreted differently according to the value of OffsetMethod.
	OffsetValue int64

	// The number of events to buffer in the Events channel. Setting this can let the
	// consumer continue fetching messages in the background while local code consumes events,
	// greatly improving throughput.
	EventBufferSize int
}

// ConsumerEvent is what is provided to the user when an event occurs. It is either an error (in which case Err is non-nil) or
// a message (in which case Err is nil and the other fields are all set).
type ConsumerEvent struct {
	Key, Value []byte
	Offset     int64
	Err        error
}

// Consumer processes Kafka messages from a given topic and partition.
// You MUST call Close() on a consumer to avoid leaks, it will not be garbage-collected automatically when
// it passes out of scope (this is in addition to calling Close on the underlying client, which is still necessary).
type Consumer struct {
	client *Client

	topic     string
	partition int32
	group     string
	config    ConsumerConfig

	offset        int64
	broker        *Broker
	stopper, done chan bool
	events        chan *ConsumerEvent
}

// NewConsumer creates a new consumer attached to the given client. It will read messages from the given topic and partition, as
// part of the named consumer group.
func NewConsumer(client *Client, topic string, partition int32, group string, config *ConsumerConfig) (*Consumer, error) {
	if config == nil {
		config = new(ConsumerConfig)
	}

	if config.DefaultFetchSize < 0 {
		return nil, ConfigurationError("Invalid DefaultFetchSize")
	} else if config.DefaultFetchSize == 0 {
		config.DefaultFetchSize = 1024
	}

	if config.MinFetchSize < 0 {
		return nil, ConfigurationError("Invalid MinFetchSize")
	} else if config.MinFetchSize == 0 {
		config.MinFetchSize = 1
	}

	if config.MaxMessageSize < 0 {
		return nil, ConfigurationError("Invalid MaxMessageSize")
	}

	if config.MaxWaitTime < 0 {
		return nil, ConfigurationError("Invalid MaxWaitTime")
	} else if config.MaxWaitTime < 100 {
		Logger.Println("ConsumerConfig.MaxWaitTime is very low, which can cause high CPU and network usage. See documentation for details.")
	}

	if config.EventBufferSize < 0 {
		return nil, ConfigurationError("Invalid EventBufferSize")
	}

	if topic == "" {
		return nil, ConfigurationError("Empty topic")
	}

	broker, err := client.Leader(topic, partition)
	if err != nil {
		return nil, err
	}

	c := &Consumer{
		client:    client,
		topic:     topic,
		partition: partition,
		group:     group,
		config:    *config,
		broker:    broker,
		stopper:   make(chan bool),
		done:      make(chan bool),
		events:    make(chan *ConsumerEvent, config.EventBufferSize),
	}

	switch config.OffsetMethod {
	case OffsetMethodManual:
		if config.OffsetValue < 0 {
			return nil, ConfigurationError("OffsetValue cannot be < 0 when OffsetMethod is MANUAL")
		}
		c.offset = config.OffsetValue
	case OffsetMethodNewest:
		c.offset, err = c.getOffset(LatestOffsets, true)
		if err != nil {
			return nil, err
		}
	case OffsetMethodOldest:
		c.offset, err = c.getOffset(EarliestOffset, true)
		if err != nil {
			return nil, err
		}
	default:
		return nil, ConfigurationError("Invalid OffsetMethod")
	}

	go withRecover(c.fetchMessages)

	return c, nil
}

// Events returns the read channel for any events (messages or errors) that might be returned by the broker.
func (c *Consumer) Events() <-chan *ConsumerEvent {
	return c.events
}

// Close stops the consumer from fetching messages. It is required to call this function before
// a consumer object passes out of scope, as it will otherwise leak memory. You must call this before
// calling Close on the underlying client.
func (c *Consumer) Close() error {
	close(c.stopper)
	<-c.done
	return nil
}

// helper function for safely sending an error on the errors channel
// if it returns true, the error was sent (or was nil)
// if it returns false, the stopper channel signaled that your goroutine should return!
func (c *Consumer) sendError(err error) bool {
	if err == nil {
		return true
	}

	select {
	case <-c.stopper:
		close(c.events)
		close(c.done)
		return false
	case c.events <- &ConsumerEvent{Err: err}:
		return true
	}

	// For backward compatibility with go1.0
	return true
}

func (c *Consumer) fetchMessages() {

	fetchSize := c.config.DefaultFetchSize

	for {
		request := new(FetchRequest)
		request.MinBytes = c.config.MinFetchSize
		request.MaxWaitTime = c.config.MaxWaitTime
		request.AddBlock(c.topic, c.partition, c.offset, fetchSize)

		response, err := c.broker.Fetch(c.client.id, request)
		switch {
		case err == nil:
			break
		case err == EncodingError:
			if c.sendError(err) {
				continue
			} else {
				return
			}
		default:
			c.client.disconnectBroker(c.broker)
			for c.broker, err = c.client.Leader(c.topic, c.partition); err != nil; c.broker, err = c.client.Leader(c.topic, c.partition) {
				if !c.sendError(err) {
					return
				}
			}
			continue
		}

		block := response.GetBlock(c.topic, c.partition)
		if block == nil {
			if c.sendError(IncompleteResponse) {
				continue
			} else {
				return
			}
		}

		switch block.Err {
		case NoError:
			break
		case UnknownTopicOrPartition, NotLeaderForPartition, LeaderNotAvailable:
			err = c.client.RefreshTopicMetadata(c.topic)
			if c.sendError(err) {
				for c.broker, err = c.client.Leader(c.topic, c.partition); err != nil; c.broker, err = c.client.Leader(c.topic, c.partition) {
					if !c.sendError(err) {
						return
					}
				}
				continue
			} else {
				return
			}
		default:
			if c.sendError(block.Err) {
				continue
			} else {
				return
			}
		}

		if len(block.MsgSet.Messages) == 0 {
			// We got no messages. If we got a trailing one then we need to ask for more data.
			// Otherwise we just poll again and wait for one to be produced...
			if block.MsgSet.PartialTrailingMessage {
				if c.config.MaxMessageSize == 0 {
					fetchSize *= 2
				} else {
					if fetchSize == c.config.MaxMessageSize {
						if c.sendError(MessageTooLarge) {
							continue
						} else {
							return
						}
					} else {
						fetchSize *= 2
						if fetchSize > c.config.MaxMessageSize {
							fetchSize = c.config.MaxMessageSize
						}
					}
				}
			}
			select {
			case <-c.stopper:
				close(c.events)
				close(c.done)
				return
			default:
				continue
			}
		} else {
			fetchSize = c.config.DefaultFetchSize
		}

		for _, msgBlock := range block.MsgSet.Messages {
			select {
			case <-c.stopper:
				close(c.events)
				close(c.done)
				return
			case c.events <- &ConsumerEvent{Key: msgBlock.Msg.Key, Value: msgBlock.Msg.Value, Offset: msgBlock.Offset}:
				c.offset++
			}
		}
	}
}

func (c *Consumer) getOffset(where OffsetTime, retry bool) (int64, error) {
	request := &OffsetRequest{}
	request.AddBlock(c.topic, c.partition, where, 1)

	response, err := c.broker.GetAvailableOffsets(c.client.id, request)
	switch err {
	case nil:
		break
	case EncodingError:
		return -1, err
	default:
		if !retry {
			return -1, err
		}
		c.client.disconnectBroker(c.broker)
		c.broker, err = c.client.Leader(c.topic, c.partition)
		if err != nil {
			return -1, err
		}
		return c.getOffset(where, false)
	}

	block := response.GetBlock(c.topic, c.partition)
	if block == nil {
		return -1, IncompleteResponse
	}

	switch block.Err {
	case NoError:
		if len(block.Offsets) < 1 {
			return -1, IncompleteResponse
		}
		return block.Offsets[0], nil
	case UnknownTopicOrPartition, NotLeaderForPartition, LeaderNotAvailable:
		if !retry {
			return -1, block.Err
		}
		err = c.client.RefreshTopicMetadata(c.topic)
		if err != nil {
			return -1, err
		}
		c.broker, err = c.client.Leader(c.topic, c.partition)
		if err != nil {
			return -1, err
		}
		return c.getOffset(where, false)
	}

	return -1, block.Err
}
<commit_msg>Don't let the caller pass in MaxWaitTime=0.<commit_after>package sarama

// OffsetMethod is passed in ConsumerConfig to tell the consumer how to determine the starting offset.
type OffsetMethod int

const (
	// OffsetMethodManual causes the consumer to interpret the OffsetValue in the ConsumerConfig as the
	// offset at which to start, allowing the user to manually specify their desired starting offset.
	OffsetMethodManual OffsetMethod = iota
	// OffsetMethodNewest causes the consumer to start at the most recent available offset, as
	// determined by querying the broker.
	OffsetMethodNewest
	// OffsetMethodOldest causes the consumer to start at the oldest available offset, as
	// determined by querying the broker.
	OffsetMethodOldest
)

// ConsumerConfig is used to pass multiple configuration options to NewConsumer.
type ConsumerConfig struct {
	// The default (maximum) amount of data to fetch from the broker in each request. The default of 0 is treated as 1024 bytes.
	DefaultFetchSize int32
	// The minimum amount of data to fetch in a request - the broker will wait until at least this many bytes are available.
	// The default of 0 is treated as 'at least one' to prevent the consumer from spinning when no messages are available.
	MinFetchSize int32
	// The maximum permittable message size - messages larger than this will return MessageTooLarge. The default of 0 is
	// treated as no limit.
	MaxMessageSize int32
	// The maximum amount of time (in ms) the broker will wait for MinFetchSize bytes to become available before it
	// returns fewer than that anyways. The default of 0 causes Kafka to return immediately, which is rarely desirable
	// as it causes the Consumer to spin when no events are available. 100-500ms is a reasonable range for most cases.
	MaxWaitTime int32

	// The method used to determine at which offset to begin consuming messages.
	OffsetMethod OffsetMethod
	// Interpreted differently according to the value of OffsetMethod.
	OffsetValue int64

	// The number of events to buffer in the Events channel. Setting this can let the
	// consumer continue fetching messages in the background while local code consumes events,
	// greatly improving throughput.
	EventBufferSize int
}

// ConsumerEvent is what is provided to the user when an event occurs. It is either an error (in which case Err is non-nil) or
// a message (in which case Err is nil and the other fields are all set).
type ConsumerEvent struct {
	Key, Value []byte
	Offset     int64
	Err        error
}

// Consumer processes Kafka messages from a given topic and partition.
// You MUST call Close() on a consumer to avoid leaks, it will not be garbage-collected automatically when
// it passes out of scope (this is in addition to calling Close on the underlying client, which is still necessary).
type Consumer struct {
	client *Client

	topic     string
	partition int32
	group     string
	config    ConsumerConfig

	offset        int64
	broker        *Broker
	stopper, done chan bool
	events        chan *ConsumerEvent
}

// NewConsumer creates a new consumer attached to the given client. It will read messages from the given topic and partition, as
// part of the named consumer group.
func NewConsumer(client *Client, topic string, partition int32, group string, config *ConsumerConfig) (*Consumer, error) {
	if config == nil {
		config = new(ConsumerConfig)
	}

	if config.DefaultFetchSize < 0 {
		return nil, ConfigurationError("Invalid DefaultFetchSize")
	} else if config.DefaultFetchSize == 0 {
		config.DefaultFetchSize = 1024
	}

	if config.MinFetchSize < 0 {
		return nil, ConfigurationError("Invalid MinFetchSize")
	} else if config.MinFetchSize == 0 {
		config.MinFetchSize = 1
	}

	if config.MaxMessageSize < 0 {
		return nil, ConfigurationError("Invalid MaxMessageSize")
	}

	if config.MaxWaitTime <= 0 {
		return nil, ConfigurationError("Invalid MaxWaitTime")
	} else if config.MaxWaitTime < 100 {
		Logger.Println("ConsumerConfig.MaxWaitTime is very low, which can cause high CPU and network usage. See documentation for details.")
	}

	if config.EventBufferSize < 0 {
		return nil, ConfigurationError("Invalid EventBufferSize")
	}

	if topic == "" {
		return nil, ConfigurationError("Empty topic")
	}

	broker, err := client.Leader(topic, partition)
	if err != nil {
		return nil, err
	}

	c := &Consumer{
		client:    client,
		topic:     topic,
		partition: partition,
		group:     group,
		config:    *config,
		broker:    broker,
		stopper:   make(chan bool),
		done:      make(chan bool),
		events:    make(chan *ConsumerEvent, config.EventBufferSize),
	}

	switch config.OffsetMethod {
	case OffsetMethodManual:
		if config.OffsetValue < 0 {
			return nil, ConfigurationError("OffsetValue cannot be < 0 when OffsetMethod is MANUAL")
		}
		c.offset = config.OffsetValue
	case OffsetMethodNewest:
		c.offset, err = c.getOffset(LatestOffsets, true)
		if err != nil {
			return nil, err
		}
	case OffsetMethodOldest:
		c.offset, err = c.getOffset(EarliestOffset, true)
		if err != nil {
			return nil, err
		}
	default:
		return nil, ConfigurationError("Invalid OffsetMethod")
	}

	go withRecover(c.fetchMessages)

	return c, nil
}

// Events returns the read channel for any events (messages or errors) that might be returned by the broker.
func (c *Consumer) Events() <-chan *ConsumerEvent {
	return c.events
}

// Close stops the consumer from fetching messages. It is required to call this function before
// a consumer object passes out of scope, as it will otherwise leak memory. You must call this before
// calling Close on the underlying client.
func (c *Consumer) Close() error {
	close(c.stopper)
	<-c.done
	return nil
}

// helper function for safely sending an error on the errors channel
// if it returns true, the error was sent (or was nil)
// if it returns false, the stopper channel signaled that your goroutine should return!
func (c *Consumer) sendError(err error) bool {
	if err == nil {
		return true
	}

	select {
	case <-c.stopper:
		close(c.events)
		close(c.done)
		return false
	case c.events <- &ConsumerEvent{Err: err}:
		return true
	}

	// For backward compatibility with go1.0
	return true
}

func (c *Consumer) fetchMessages() {

	fetchSize := c.config.DefaultFetchSize

	for {
		request := new(FetchRequest)
		request.MinBytes = c.config.MinFetchSize
		request.MaxWaitTime = c.config.MaxWaitTime
		request.AddBlock(c.topic, c.partition, c.offset, fetchSize)

		response, err := c.broker.Fetch(c.client.id, request)
		switch {
		case err == nil:
			break
		case err == EncodingError:
			if c.sendError(err) {
				continue
			} else {
				return
			}
		default:
			c.client.disconnectBroker(c.broker)
			for c.broker, err = c.client.Leader(c.topic, c.partition); err != nil; c.broker, err = c.client.Leader(c.topic, c.partition) {
				if !c.sendError(err) {
					return
				}
			}
			continue
		}

		block := response.GetBlock(c.topic, c.partition)
		if block == nil {
			if c.sendError(IncompleteResponse) {
				continue
			} else {
				return
			}
		}

		switch block.Err {
		case NoError:
			break
		case UnknownTopicOrPartition, NotLeaderForPartition, LeaderNotAvailable:
			err = c.client.RefreshTopicMetadata(c.topic)
			if c.sendError(err) {
				for c.broker, err = c.client.Leader(c.topic, c.partition); err != nil; c.broker, err = c.client.Leader(c.topic, c.partition) {
					if !c.sendError(err) {
						return
					}
				}
				continue
			} else {
				return
			}
		default:
			if c.sendError(block.Err) {
				continue
			} else {
				return
			}
		}

		if len(block.MsgSet.Messages) == 0 {
			// We got no messages. If we got a trailing one then we need to ask for more data.
			// Otherwise we just poll again and wait for one to be produced...
			if block.MsgSet.PartialTrailingMessage {
				if c.config.MaxMessageSize == 0 {
					fetchSize *= 2
				} else {
					if fetchSize == c.config.MaxMessageSize {
						if c.sendError(MessageTooLarge) {
							continue
						} else {
							return
						}
					} else {
						fetchSize *= 2
						if fetchSize > c.config.MaxMessageSize {
							fetchSize = c.config.MaxMessageSize
						}
					}
				}
			}
			select {
			case <-c.stopper:
				close(c.events)
				close(c.done)
				return
			default:
				continue
			}
		} else {
			fetchSize = c.config.DefaultFetchSize
		}

		for _, msgBlock := range block.MsgSet.Messages {
			select {
			case <-c.stopper:
				close(c.events)
				close(c.done)
				return
			case c.events <- &ConsumerEvent{Key: msgBlock.Msg.Key, Value: msgBlock.Msg.Value, Offset: msgBlock.Offset}:
				c.offset++
			}
		}
	}
}

func (c *Consumer) getOffset(where OffsetTime, retry bool) (int64, error) {
	request := &OffsetRequest{}
	request.AddBlock(c.topic, c.partition, where, 1)

	response, err := c.broker.GetAvailableOffsets(c.client.id, request)
	switch err {
	case nil:
		break
	case EncodingError:
		return -1, err
	default:
		if !retry {
			return -1, err
		}
		c.client.disconnectBroker(c.broker)
		c.broker, err = c.client.Leader(c.topic, c.partition)
		if err != nil {
			return -1, err
		}
		return c.getOffset(where, false)
	}

	block := response.GetBlock(c.topic, c.partition)
	if block == nil {
		return -1, IncompleteResponse
	}

	switch block.Err {
	case NoError:
		if len(block.Offsets) < 1 {
			return -1, IncompleteResponse
		}
		return block.Offsets[0], nil
	case UnknownTopicOrPartition, NotLeaderForPartition, LeaderNotAvailable:
		if !retry {
			return -1, block.Err
		}
		err = c.client.RefreshTopicMetadata(c.topic)
		if err != nil {
			return -1, err
		}
		c.broker, err = c.client.Leader(c.topic, c.partition)
		if err != nil {
			return -1, err
		}
		return c.getOffset(where, false)
	}

	return -1, block.Err
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"
	"log"
	"strconv"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/config"
)

// EvalApply is an EvalNode implementation that writes the diff to
// the full diff.
type EvalApply struct {
	Info      *InstanceInfo
	State     **InstanceState
	Diff      **InstanceDiff
	Provider  *ResourceProvider
	Output    **InstanceState
	CreateNew *bool
	Error     *error
}

// TODO: test
func (n *EvalApply) Eval(ctx EvalContext) (interface{}, error) {
	diff := *n.Diff
	provider := *n.Provider
	state := *n.State

	// If we have no diff, we have nothing to do!
	if diff.Empty() {
		log.Printf(
			"[DEBUG] apply: %s: diff is empty, doing nothing.", n.Info.Id)
		return nil, nil
	}

	// Remove any output values from the diff
	for k, ad := range diff.CopyAttributes() {
		if ad.Type == DiffAttrOutput {
			diff.DelAttribute(k)
		}
	}

	// If the state is nil, make it non-nil
	if state == nil {
		state = new(InstanceState)
	}
	state.init()

	// Flag if we're creating a new instance
	if n.CreateNew != nil {
		*n.CreateNew = state.ID == "" && !diff.GetDestroy() || diff.RequiresNew()
	}

	// With the completed diff, apply!
	log.Printf("[DEBUG] apply: %s: executing Apply", n.Info.Id)
	state, err := provider.Apply(n.Info, state, diff)
	if state == nil {
		state = new(InstanceState)
	}
	state.init()

	// Force the "id" attribute to be our ID
	if state.ID != "" {
		state.Attributes["id"] = state.ID
	}

	// If the value is the unknown variable value, then it is an error.
	// In this case we record the error and remove it from the state
	for ak, av := range state.Attributes {
		if av == config.UnknownVariableValue {
			err = multierror.Append(err, fmt.Errorf(
				"Attribute with unknown value: %s", ak))
			delete(state.Attributes, ak)
		}
	}

	// Write the final state
	if n.Output != nil {
		*n.Output = state
	}

	// If there are no errors, then we append it to our output error
	// if we have one, otherwise we just output it.
	if err != nil {
		if n.Error != nil {
			helpfulErr := fmt.Errorf("%s: %s", n.Info.Id, err.Error())
			*n.Error = multierror.Append(*n.Error, helpfulErr)
		} else {
			return nil, err
		}
	}

	return nil, nil
}

// EvalApplyPre is an EvalNode implementation that does the pre-Apply work
type EvalApplyPre struct {
	Info  *InstanceInfo
	State **InstanceState
	Diff  **InstanceDiff
}

// TODO: test
func (n *EvalApplyPre) Eval(ctx EvalContext) (interface{}, error) {
	state := *n.State
	diff := *n.Diff

	// If the state is nil, make it non-nil
	if state == nil {
		state = new(InstanceState)
	}
	state.init()

	if resourceHasUserVisibleApply(n.Info) {
		// Call post-apply hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PreApply(n.Info, state, diff)
		})
		if err != nil {
			return nil, err
		}
	}

	return nil, nil
}

// EvalApplyPost is an EvalNode implementation that does the post-Apply work
type EvalApplyPost struct {
	Info  *InstanceInfo
	State **InstanceState
	Error *error
}

// TODO: test
func (n *EvalApplyPost) Eval(ctx EvalContext) (interface{}, error) {
	state := *n.State

	if resourceHasUserVisibleApply(n.Info) {
		// Call post-apply hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PostApply(n.Info, state, *n.Error)
		})
		if err != nil {
			return nil, err
		}
	}

	return nil, *n.Error
}

// resourceHasUserVisibleApply returns true if the given resource is one where
// apply actions should be exposed to the user.
//
// Certain resources do apply actions only as an implementation detail, so
// these should not be advertised to code outside of this package.
func resourceHasUserVisibleApply(info *InstanceInfo) bool {
	addr := info.ResourceAddress()

	// Only managed resources have user-visible apply actions.
	// In particular, this excludes data resources since we "apply" these
	// only as an implementation detail of removing them from state when
	// they are destroyed. (When reading, they don't get here at all because
	// we present them as "Refresh" actions.)
	return addr.Mode == config.ManagedResourceMode
}

// EvalApplyProvisioners is an EvalNode implementation that executes
// the provisioners for a resource.
//
// TODO(mitchellh): This should probably be split up into a more fine-grained
// ApplyProvisioner (single) that is looped over.
type EvalApplyProvisioners struct {
	Info           *InstanceInfo
	State          **InstanceState
	Resource       *config.Resource
	InterpResource *Resource
	CreateNew      *bool
	Error          *error

	// When is the type of provisioner to run at this point
	When config.ProvisionerWhen
}

// TODO: test
func (n *EvalApplyProvisioners) Eval(ctx EvalContext) (interface{}, error) {
	state := *n.State

	if n.CreateNew != nil && !*n.CreateNew {
		// If we're not creating a new resource, then don't run provisioners
		return nil, nil
	}

	provs := n.filterProvisioners()
	if len(provs) == 0 {
		// We have no provisioners, so don't do anything
		return nil, nil
	}

	// taint tells us whether to enable tainting.
	taint := n.When == config.ProvisionerWhenCreate

	if n.Error != nil && *n.Error != nil {
		if taint {
			state.Tainted = true
		}

		// We're already tainted, so just return out
		return nil, nil
	}

	{
		// Call pre hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PreProvisionResource(n.Info, state)
		})
		if err != nil {
			return nil, err
		}
	}

	// If there are no errors, then we append it to our output error
	// if we have one, otherwise we just output it.
	err := n.apply(ctx, provs)
	if err != nil {
		if taint {
			state.Tainted = true
		}

		if n.Error != nil {
			*n.Error = multierror.Append(*n.Error, err)
		} else {
			return nil, err
		}
	}

	{
		// Call post hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PostProvisionResource(n.Info, state)
		})
		if err != nil {
			return nil, err
		}
	}

	return nil, nil
}

// filterProvisioners filters the provisioners on the resource to only
// the provisioners specified by the "when" option.
func (n *EvalApplyProvisioners) filterProvisioners() []*config.Provisioner {
	// Fast path the zero case
	if n.Resource == nil {
		return nil
	}

	if len(n.Resource.Provisioners) == 0 {
		return nil
	}

	result := make([]*config.Provisioner, 0, len(n.Resource.Provisioners))
	for _, p := range n.Resource.Provisioners {
		if p.When == n.When {
			result = append(result, p)
		}
	}

	return result
}

func (n *EvalApplyProvisioners) apply(ctx EvalContext, provs []*config.Provisioner) error {
	state := *n.State

	// Store the original connection info, restore later
	origConnInfo := state.Ephemeral.ConnInfo
	defer func() {
		state.Ephemeral.ConnInfo = origConnInfo
	}()

	for _, prov := range provs {
		// Get the provisioner
		provisioner := ctx.Provisioner(prov.Type)

		// Interpolate the provisioner config
		provConfig, err := ctx.Interpolate(prov.RawConfig.Copy(), n.InterpResource)
		if err != nil {
			return err
		}

		// Interpolate the conn info, since it may contain variables
		connInfo, err := ctx.Interpolate(prov.ConnInfo.Copy(), n.InterpResource)
		if err != nil {
			return err
		}

		// Merge the connection information
		overlay := make(map[string]string)
		if origConnInfo != nil {
			for k, v := range origConnInfo {
				overlay[k] = v
			}
		}
		for k, v := range connInfo.Config {
			switch vt := v.(type) {
			case string:
				overlay[k] = vt
			case int64:
				overlay[k] = strconv.FormatInt(vt, 10)
			case int32:
				overlay[k] = strconv.FormatInt(int64(vt), 10)
			case int:
				overlay[k] = strconv.FormatInt(int64(vt), 10)
			case float32:
				overlay[k] = strconv.FormatFloat(float64(vt), 'f', 3, 32)
			case float64:
				overlay[k] = strconv.FormatFloat(vt, 'f', 3, 64)
			case bool:
				overlay[k] = strconv.FormatBool(vt)
			default:
				overlay[k] = fmt.Sprintf("%v", vt)
			}
		}
		state.Ephemeral.ConnInfo = overlay

		{
			// Call pre hook
			err := ctx.Hook(func(h Hook) (HookAction, error) {
				return h.PreProvision(n.Info, prov.Type)
			})
			if err != nil {
				return err
			}
		}

		// The output function
		outputFn := func(msg string) {
			ctx.Hook(func(h Hook) (HookAction, error) {
				h.ProvisionOutput(n.Info, prov.Type, msg)
				return HookActionContinue, nil
			})
		}

		// Invoke the Provisioner
		output := CallbackUIOutput{OutputFn: outputFn}
		applyErr := provisioner.Apply(&output, state, provConfig)

		// Call post hook
		hookErr := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PostProvision(n.Info, prov.Type, applyErr)
		})

		// Handle the error before we deal with the hook
		if applyErr != nil {
			// Determine failure behavior
			switch prov.OnFailure {
			case config.ProvisionerOnFailureContinue:
				log.Printf(
					"[INFO] apply: %s [%s]: error during provision, continue requested",
					n.Info.Id, prov.Type)

			case config.ProvisionerOnFailureFail:
				return applyErr
			}
		}

		// Deal with the hook
		if hookErr != nil {
			return hookErr
		}
	}

	return nil

}
<commit_msg>return provisioner Apply errors<commit_after>package terraform

import (
	"fmt"
	"log"
	"strconv"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/config"
)

// EvalApply is an EvalNode implementation that writes the diff to
// the full diff.
type EvalApply struct {
	Info      *InstanceInfo
	State     **InstanceState
	Diff      **InstanceDiff
	Provider  *ResourceProvider
	Output    **InstanceState
	CreateNew *bool
	Error     *error
}

// TODO: test
func (n *EvalApply) Eval(ctx EvalContext) (interface{}, error) {
	diff := *n.Diff
	provider := *n.Provider
	state := *n.State

	// If we have no diff, we have nothing to do!
	if diff.Empty() {
		log.Printf(
			"[DEBUG] apply: %s: diff is empty, doing nothing.", n.Info.Id)
		return nil, nil
	}

	// Remove any output values from the diff
	for k, ad := range diff.CopyAttributes() {
		if ad.Type == DiffAttrOutput {
			diff.DelAttribute(k)
		}
	}

	// If the state is nil, make it non-nil
	if state == nil {
		state = new(InstanceState)
	}
	state.init()

	// Flag if we're creating a new instance
	if n.CreateNew != nil {
		*n.CreateNew = state.ID == "" && !diff.GetDestroy() || diff.RequiresNew()
	}

	// With the completed diff, apply!
	log.Printf("[DEBUG] apply: %s: executing Apply", n.Info.Id)
	state, err := provider.Apply(n.Info, state, diff)
	if state == nil {
		state = new(InstanceState)
	}
	state.init()

	// Force the "id" attribute to be our ID
	if state.ID != "" {
		state.Attributes["id"] = state.ID
	}

	// If the value is the unknown variable value, then it is an error.
	// In this case we record the error and remove it from the state
	for ak, av := range state.Attributes {
		if av == config.UnknownVariableValue {
			err = multierror.Append(err, fmt.Errorf(
				"Attribute with unknown value: %s", ak))
			delete(state.Attributes, ak)
		}
	}

	// Write the final state
	if n.Output != nil {
		*n.Output = state
	}

	// If there are no errors, then we append it to our output error
	// if we have one, otherwise we just output it.
	if err != nil {
		if n.Error != nil {
			helpfulErr := fmt.Errorf("%s: %s", n.Info.Id, err.Error())
			*n.Error = multierror.Append(*n.Error, helpfulErr)
		} else {
			return nil, err
		}
	}

	return nil, nil
}

// EvalApplyPre is an EvalNode implementation that does the pre-Apply work
type EvalApplyPre struct {
	Info  *InstanceInfo
	State **InstanceState
	Diff  **InstanceDiff
}

// TODO: test
func (n *EvalApplyPre) Eval(ctx EvalContext) (interface{}, error) {
	state := *n.State
	diff := *n.Diff

	// If the state is nil, make it non-nil
	if state == nil {
		state = new(InstanceState)
	}
	state.init()

	if resourceHasUserVisibleApply(n.Info) {
		// Call post-apply hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PreApply(n.Info, state, diff)
		})
		if err != nil {
			return nil, err
		}
	}

	return nil, nil
}

// EvalApplyPost is an EvalNode implementation that does the post-Apply work
type EvalApplyPost struct {
	Info  *InstanceInfo
	State **InstanceState
	Error *error
}

// TODO: test
func (n *EvalApplyPost) Eval(ctx EvalContext) (interface{}, error) {
	state := *n.State

	if resourceHasUserVisibleApply(n.Info) {
		// Call post-apply hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PostApply(n.Info, state, *n.Error)
		})
		if err != nil {
			return nil, err
		}
	}

	return nil, *n.Error
}

// resourceHasUserVisibleApply returns true if the given resource is one where
// apply actions should be exposed to the user.
//
// Certain resources do apply actions only as an implementation detail, so
// these should not be advertised to code outside of this package.
func resourceHasUserVisibleApply(info *InstanceInfo) bool {
	addr := info.ResourceAddress()

	// Only managed resources have user-visible apply actions.
	// In particular, this excludes data resources since we "apply" these
	// only as an implementation detail of removing them from state when
	// they are destroyed. (When reading, they don't get here at all because
	// we present them as "Refresh" actions.)
	return addr.Mode == config.ManagedResourceMode
}

// EvalApplyProvisioners is an EvalNode implementation that executes
// the provisioners for a resource.
//
// TODO(mitchellh): This should probably be split up into a more fine-grained
// ApplyProvisioner (single) that is looped over.
type EvalApplyProvisioners struct {
	Info           *InstanceInfo
	State          **InstanceState
	Resource       *config.Resource
	InterpResource *Resource
	CreateNew      *bool
	Error          *error

	// When is the type of provisioner to run at this point
	When config.ProvisionerWhen
}

// TODO: test
func (n *EvalApplyProvisioners) Eval(ctx EvalContext) (interface{}, error) {
	state := *n.State

	if n.CreateNew != nil && !*n.CreateNew {
		// If we're not creating a new resource, then don't run provisioners
		return nil, nil
	}

	provs := n.filterProvisioners()
	if len(provs) == 0 {
		// We have no provisioners, so don't do anything
		return nil, nil
	}

	// taint tells us whether to enable tainting.
	taint := n.When == config.ProvisionerWhenCreate

	if n.Error != nil && *n.Error != nil {
		if taint {
			state.Tainted = true
		}

		// We're already tainted, so just return out
		return nil, nil
	}

	{
		// Call pre hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PreProvisionResource(n.Info, state)
		})
		if err != nil {
			return nil, err
		}
	}

	// If there are no errors, then we append it to our output error
	// if we have one, otherwise we just output it.
	err := n.apply(ctx, provs)
	if err != nil {
		if taint {
			state.Tainted = true
		}

		*n.Error = multierror.Append(*n.Error, err)
		return nil, err
	}

	{
		// Call post hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PostProvisionResource(n.Info, state)
		})
		if err != nil {
			return nil, err
		}
	}

	return nil, nil
}

// filterProvisioners filters the provisioners on the resource to only
// the provisioners specified by the "when" option.
func (n *EvalApplyProvisioners) filterProvisioners() []*config.Provisioner {
	// Fast path the zero case
	if n.Resource == nil {
		return nil
	}

	if len(n.Resource.Provisioners) == 0 {
		return nil
	}

	result := make([]*config.Provisioner, 0, len(n.Resource.Provisioners))
	for _, p := range n.Resource.Provisioners {
		if p.When == n.When {
			result = append(result, p)
		}
	}

	return result
}

func (n *EvalApplyProvisioners) apply(ctx EvalContext, provs []*config.Provisioner) error {
	state := *n.State

	// Store the original connection info, restore later
	origConnInfo := state.Ephemeral.ConnInfo
	defer func() {
		state.Ephemeral.ConnInfo = origConnInfo
	}()

	for _, prov := range provs {
		// Get the provisioner
		provisioner := ctx.Provisioner(prov.Type)

		// Interpolate the provisioner config
		provConfig, err := ctx.Interpolate(prov.RawConfig.Copy(), n.InterpResource)
		if err != nil {
			return err
		}

		// Interpolate the conn info, since it may contain variables
		connInfo, err := ctx.Interpolate(prov.ConnInfo.Copy(), n.InterpResource)
		if err != nil {
			return err
		}

		// Merge the connection information
		overlay := make(map[string]string)
		if origConnInfo != nil {
			for k, v := range origConnInfo {
				overlay[k] = v
			}
		}
		for k, v := range connInfo.Config {
			switch vt := v.(type) {
			case string:
				overlay[k] = vt
			case int64:
				overlay[k] = strconv.FormatInt(vt, 10)
			case int32:
				overlay[k] = strconv.FormatInt(int64(vt), 10)
			case int:
				overlay[k] = strconv.FormatInt(int64(vt), 10)
			case float32:
				overlay[k] = strconv.FormatFloat(float64(vt), 'f', 3, 32)
			case float64:
				overlay[k] = strconv.FormatFloat(vt, 'f', 3, 64)
			case bool:
				overlay[k] = strconv.FormatBool(vt)
			default:
				overlay[k] = fmt.Sprintf("%v", vt)
			}
		}
		state.Ephemeral.ConnInfo = overlay

		{
			// Call pre hook
			err := ctx.Hook(func(h Hook) (HookAction, error) {
				return h.PreProvision(n.Info, prov.Type)
			})
			if err != nil {
				return err
			}
		}

		// The output function
		outputFn := func(msg string) {
			ctx.Hook(func(h Hook) (HookAction, error) {
				h.ProvisionOutput(n.Info, prov.Type, msg)
				return HookActionContinue, nil
			})
		}

		// Invoke the Provisioner
		output := CallbackUIOutput{OutputFn: outputFn}
		applyErr := provisioner.Apply(&output, state, provConfig)

		// Call post hook
		hookErr := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PostProvision(n.Info, prov.Type, applyErr)
		})

		// Handle the error before we deal with the hook
		if applyErr != nil {
			// Determine failure behavior
			switch prov.OnFailure {
			case config.ProvisionerOnFailureContinue:
				log.Printf(
					"[INFO] apply: %s [%s]: error during provision, continue requested",
					n.Info.Id, prov.Type)

			case config.ProvisionerOnFailureFail:
				return applyErr
			}
		}

		// Deal with the hook
		if hookErr != nil {
			return hookErr
		}
	}

	return nil

}
<|endoftext|>


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

// Package morass implements file system-backed sorting.
//
// Use morass when you don't want your data to be a quagmire.
//
// Sort data larger than can fit in memory.
//
//  morass məˈras/
//  1. An area of muddy or boggy ground.
//  2. A complicated or confused situation.
package morass

import (
	"container/heap"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"runtime"
	"sort"
	"sync"
)

var (
	registerLock = &sync.Mutex{}
	registered   = make(map[reflect.Type]struct{})
	nextID       = 0
)

func register(e interface{}, t reflect.Type) {
	registerLock.Lock()
	defer registerLock.Unlock()
	defer func() {
		recover()                  // The only panic that we can get is from trying to register a base type.
		registered[t] = struct{}{} // Remember for next time.
	}()

	if _, exists := registered[t]; !exists {
		registered[t] = struct{}{}
		gob.RegisterName(fmt.Sprintf("ℳ%d", nextID), e)
		nextID++
	}
}

// LessInterface wraps the Less method.
type LessInterface interface {
	// Is the receiver less than the parameterised interface
	Less(i interface{}) bool
}

type sorter []LessInterface

func (s sorter) Len() int { return len(s) }

func (s sorter) Less(i, j int) bool { return s[i].Less(s[j]) }

func (s sorter) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type file struct {
	head    LessInterface
	file    *os.File
	encoder *gob.Encoder
	decoder *gob.Decoder
}

type files []*file

func (f files) Len() int { return len(f) }

func (f files) Less(i, j int) bool { return f[i].head.Less(f[j].head) }

func (f files) Swap(i, j int) { f[i], f[j] = f[j], f[i] }

func (f *files) Pop() (i interface{}) {
	i = (*f)[len(*f)-1]
	*f = (*f)[:len(*f)-1]
	return
}

func (f *files) Push(x interface{}) { *f = append(*f, x.(*file)) }

// Morass implements sorting of very large data sets.
type Morass struct {
	typ reflect.Type

	pos, len int64

	// dir and prefix specify the location
	// of temporary files.
	dir    string
	prefix string

	// AutoClear specifies that the Morass
	// should call Clear when emptied by
	// a call to Pull.
	AutoClear bool

	// AutoClean specifies that the Morass
	// should call delete temporary sort
	// files when it has been emptied by
	// a call to Pull.
	AutoClean bool

	// fast indicates sorting was performed
	// entirely in memory.
	fast bool

	chunk     sorter
	chunkSize int
	pool      chan sorter
	writable  chan sorter

	filesLock sync.Mutex
	files     files

	errLock sync.Mutex
	_err    error
}

// New creates a new Morass. prefix and dir are passed to ioutil.TempDir. chunkSize specifies
// the amount of sorting to be done in memory, concurrent specifies that temporary file
// writing occurs concurrently with sorting.
// An error is returned if no temporary directory can be created.
// Note that the type is registered with the underlying gob encoder using the name ℳn, where
// n is a sequentially assigned integer string, when the type registered. This is done to avoid using
// too much space and will cause problems when using gob itself on this type. If you intend
// use gob itself with this the type, preregister with gob and morass will use the existing
// registration.
func New(e interface{}, prefix, dir string, chunkSize int, concurrent bool) (*Morass, error) {
	d, err := ioutil.TempDir(dir, prefix)
	if err != nil {
		return nil, err
	}

	m := &Morass{
		chunkSize: chunkSize,
		prefix:    prefix,
		dir:       d,
		pool:      make(chan sorter, 2),
		writable:  make(chan sorter, 1),
		files:     files{},
	}

	m.typ = reflect.TypeOf(e)
	register(e, m.typ)

	m.chunk = make(sorter, 0, chunkSize)
	if concurrent {
		m.pool <- nil
	}

	runtime.SetFinalizer(m, func(x *Morass) {
		if x.AutoClean {
			x.CleanUp()
		}
	})

	return m, nil
}

// Push a value on to the Morass. Returns any error that occurs.
func (m *Morass) Push(e LessInterface) error {
	if typ := reflect.TypeOf(e); typ != m.typ {
		return fmt.Errorf("morass: type mismatch: %s != %s", typ, m.typ)
	}

	if err := m.err(); err != nil {
		return err
	}

	if m.chunk == nil {
		return errors.New("morass: push on finalised morass")
	}

	if len(m.chunk) == m.chunkSize {
		m.writable <- m.chunk
		go m.write()
		m.chunk = <-m.pool
		if err := m.err(); err != nil {
			return err
		}
		if cap(m.chunk) == 0 {
			m.chunk = make(sorter, 0, m.chunkSize)
		}
	}

	m.chunk = append(m.chunk, e)
	m.pos++
	m.len++

	return nil
}

func (m *Morass) write() {
	writing := <-m.writable
	defer func() {
		m.pool <- writing[:0]
	}()

	sort.Sort(writing)

	tf, err := ioutil.TempFile(m.dir, m.prefix)
	if err != nil {
		m.setErr(err)
		return
	}

	enc := gob.NewEncoder(tf)
	dec := gob.NewDecoder(tf)
	f := &file{head: nil, file: tf, encoder: enc, decoder: dec}

	m.filesLock.Lock()
	m.files = append(m.files, f)
	m.filesLock.Unlock()

	for _, e := range writing {
		if err := enc.Encode(&e); err != nil {
			m.setErr(err)
			return
		}
	}

	m.setErr(tf.Sync())
}

func (m *Morass) setErr(err error) {
	m.errLock.Lock()
	m._err = err
	m.errLock.Unlock()
}

func (m *Morass) err() error {
	m.errLock.Lock()
	defer m.errLock.Unlock()
	return m._err
}

// Pos returns the current position of the cursor in the Morass.
func (m *Morass) Pos() int64 { return m.pos }

// Len returns the current length of the Morass.
func (m *Morass) Len() int64 { return m.len }

// Finalise is called to indicate that the last element has been pushed on to the Morass
// and write out final data.
func (m *Morass) Finalise() error {
	if err := m.err(); err != nil {
		return err
	}

	if m.chunk != nil {
		if m.pos < int64(cap(m.chunk)) {
			m.fast = true
			sort.Sort(m.chunk)
		} else {
			if len(m.chunk) > 0 {
				m.writable <- m.chunk
				m.chunk = nil
				m.write()
				if err := m.err(); err != nil {
					return err
				}
			}
		}
		m.pos = 0
	} else {
		return nil
	}

	if !m.fast {
		for _, f := range m.files {
			_, err := f.file.Seek(0, 0)
			if err != nil {
				return err
			}
			err = f.decoder.Decode(&f.head)
			if err != nil && err != io.EOF {
				return err
			}
		}

		heap.Init(&m.files)
	}

	return nil
}

// Clear resets the Morass to an empty state.
func (m *Morass) Clear() error {
	var err error
	for _, f := range m.files {
		err = f.file.Close()
		if err != nil {
			return err
		}
		err = os.Remove(f.file.Name())
		if err != nil {
			return err
		}
	}
	m._err = nil
	m.files = m.files[:0]
	m.pos = 0
	m.len = 0
	select {
	case m.chunk = <-m.pool:
		if m.chunk == nil {
			m.chunk = make(sorter, 0, m.chunkSize)
		}
	default:
	}

	return nil
}

// CleanUp deletes the file system components of the Morass. After this call
// the Morass is not usable.
func (m *Morass) CleanUp() error {
	return os.RemoveAll(m.dir)
}

// Pull sets the settable value e to the lowest value in the Morass.
// If io.EOF is returned the Morass is empty. Any other error results
// in no value being set on e.
func (m *Morass) Pull(e LessInterface) error {
	var err error
	v := reflect.ValueOf(e)
	if !reflect.Indirect(v).CanSet() {
		return errors.New("morass: Cannot set e")
	}

	if m.fast {
		switch {
		case m.chunk != nil && m.pos < int64(len(m.chunk)):
			e = m.chunk[m.pos].(LessInterface)
			m.pos++
		case m.chunk != nil:
			m.pool <- m.chunk[:0]
			m.chunk = nil
			fallthrough
		default:
			if m.AutoClear {
				m.Clear()
			}
			err = io.EOF
		}
	} else {
		if m.files.Len() > 0 {
			low := heap.Pop(&m.files).(*file)
			e = low.head
			m.pos++
			switch err = low.decoder.Decode(&low.head); err {
			case nil:
				heap.Push(&m.files, low)
			case io.EOF:
				err = nil
				fallthrough
			default:
				low.file.Close()
				if m.AutoClear {
					os.Remove(low.file.Name())
				}
			}
		} else {
			if m.AutoClear {
				m.Clear()
			}
			if m.AutoClean {
				os.RemoveAll(m.dir)
			}
			err = io.EOF
		}
	}

	if err != io.EOF {
		reflect.Indirect(v).Set(reflect.ValueOf(e))
	}

	return err
}
<commit_msg>Fix error<commit_after>// Copyright ©2011-2012 The bíogo Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package morass implements file system-backed sorting.
//
// Use morass when you don't want your data to be a quagmire.
//
// Sort data larger than can fit in memory.
//
//  morass məˈras/
//  1. An area of muddy or boggy ground.
//  2. A complicated or confused situation.
package morass

import (
	"container/heap"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"runtime"
	"sort"
	"sync"
)

var (
	registerLock = &sync.Mutex{}
	registered   = make(map[reflect.Type]struct{})
	nextID       = 0
)

func register(e interface{}, t reflect.Type) {
	registerLock.Lock()
	defer registerLock.Unlock()
	defer func() {
		recover()                  // The only panic that we can get is from trying to register a base type.
		registered[t] = struct{}{} // Remember for next time.
	}()

	if _, exists := registered[t]; !exists {
		registered[t] = struct{}{}
		gob.RegisterName(fmt.Sprintf("ℳ%d", nextID), e)
		nextID++
	}
}

// LessInterface wraps the Less method.
type LessInterface interface {
	// Is the receiver less than the parameterised interface
	Less(i interface{}) bool
}

type sorter []LessInterface

func (s sorter) Len() int { return len(s) }

func (s sorter) Less(i, j int) bool { return s[i].Less(s[j]) }

func (s sorter) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type file struct {
	head    LessInterface
	file    *os.File
	encoder *gob.Encoder
	decoder *gob.Decoder
}

type files []*file

func (f files) Len() int { return len(f) }

func (f files) Less(i, j int) bool { return f[i].head.Less(f[j].head) }

func (f files) Swap(i, j int) { f[i], f[j] = f[j], f[i] }

func (f *files) Pop() (i interface{}) {
	i = (*f)[len(*f)-1]
	*f = (*f)[:len(*f)-1]
	return
}

func (f *files) Push(x interface{}) { *f = append(*f, x.(*file)) }

// Morass implements sorting of very large data sets.
type Morass struct {
	typ reflect.Type

	pos, len int64

	// dir and prefix specify the location
	// of temporary files.
	dir    string
	prefix string

	// AutoClear specifies that the Morass
	// should call Clear when emptied by
	// a call to Pull.
	AutoClear bool

	// AutoClean specifies that the Morass
	// should call delete temporary sort
	// files when it has been emptied by
	// a call to Pull.
	AutoClean bool

	// fast indicates sorting was performed
	// entirely in memory.
	fast bool

	chunk     sorter
	chunkSize int
	pool      chan sorter
	writable  chan sorter

	filesLock sync.Mutex
	files     files

	errLock sync.Mutex
	_err    error
}

// New creates a new Morass. prefix and dir are passed to ioutil.TempDir. chunkSize specifies
// the amount of sorting to be done in memory, concurrent specifies that temporary file
// writing occurs concurrently with sorting.
// An error is returned if no temporary directory can be created.
// Note that the type is registered with the underlying gob encoder using the name ℳn, where
// n is a sequentially assigned integer string, when the type registered. This is done to avoid using
// too much space and will cause problems when using gob itself on this type. If you intend
// use gob itself with this the type, preregister with gob and morass will use the existing
// registration.
func New(e interface{}, prefix, dir string, chunkSize int, concurrent bool) (*Morass, error) {
	d, err := ioutil.TempDir(dir, prefix)
	if err != nil {
		return nil, err
	}

	m := &Morass{
		chunkSize: chunkSize,
		prefix:    prefix,
		dir:       d,
		pool:      make(chan sorter, 2),
		writable:  make(chan sorter, 1),
		files:     files{},
	}

	m.typ = reflect.TypeOf(e)
	register(e, m.typ)

	m.chunk = make(sorter, 0, chunkSize)
	if concurrent {
		m.pool <- nil
	}

	runtime.SetFinalizer(m, func(x *Morass) {
		if x.AutoClean {
			x.CleanUp()
		}
	})

	return m, nil
}

// Push a value on to the Morass. Returns any error that occurs.
func (m *Morass) Push(e LessInterface) error {
	if typ := reflect.TypeOf(e); typ != m.typ {
		return fmt.Errorf("morass: type mismatch: %s != %s", typ, m.typ)
	}

	if err := m.err(); err != nil {
		return err
	}

	if m.chunk == nil {
		return errors.New("morass: push on finalised morass")
	}

	if len(m.chunk) == m.chunkSize {
		m.writable <- m.chunk
		go m.write()
		m.chunk = <-m.pool
		if err := m.err(); err != nil {
			return err
		}
		if cap(m.chunk) == 0 {
			m.chunk = make(sorter, 0, m.chunkSize)
		}
	}

	m.chunk = append(m.chunk, e)
	m.pos++
	m.len++

	return nil
}

func (m *Morass) write() {
	writing := <-m.writable
	defer func() {
		m.pool <- writing[:0]
	}()

	sort.Sort(writing)

	tf, err := ioutil.TempFile(m.dir, m.prefix)
	if err != nil {
		m.setErr(err)
		return
	}

	enc := gob.NewEncoder(tf)
	dec := gob.NewDecoder(tf)
	f := &file{head: nil, file: tf, encoder: enc, decoder: dec}

	m.filesLock.Lock()
	m.files = append(m.files, f)
	m.filesLock.Unlock()

	for _, e := range writing {
		if err := enc.Encode(&e); err != nil {
			m.setErr(err)
			return
		}
	}

	m.setErr(tf.Sync())
}

func (m *Morass) setErr(err error) {
	m.errLock.Lock()
	m._err = err
	m.errLock.Unlock()
}

func (m *Morass) err() error {
	m.errLock.Lock()
	defer m.errLock.Unlock()
	return m._err
}

// Pos returns the current position of the cursor in the Morass.
func (m *Morass) Pos() int64 { return m.pos }

// Len returns the current length of the Morass.
func (m *Morass) Len() int64 { return m.len }

// Finalise is called to indicate that the last element has been pushed on to the Morass
// and write out final data.
func (m *Morass) Finalise() error {
	if err := m.err(); err != nil {
		return err
	}

	if m.chunk != nil {
		if m.pos < int64(cap(m.chunk)) {
			m.fast = true
			sort.Sort(m.chunk)
		} else {
			if len(m.chunk) > 0 {
				m.writable <- m.chunk
				m.chunk = nil
				m.write()
				if err := m.err(); err != nil {
					return err
				}
			}
		}
		m.pos = 0
	} else {
		return nil
	}

	if !m.fast {
		for _, f := range m.files {
			_, err := f.file.Seek(0, 0)
			if err != nil {
				return err
			}
			err = f.decoder.Decode(&f.head)
			if err != nil && err != io.EOF {
				return err
			}
		}

		heap.Init(&m.files)
	}

	return nil
}

// Clear resets the Morass to an empty state.
func (m *Morass) Clear() error {
	var err error
	for _, f := range m.files {
		err = f.file.Close()
		if err != nil {
			return err
		}
		err = os.Remove(f.file.Name())
		if err != nil {
			return err
		}
	}
	m._err = nil
	m.files = m.files[:0]
	m.pos = 0
	m.len = 0
	select {
	case m.chunk = <-m.pool:
		if m.chunk == nil {
			m.chunk = make(sorter, 0, m.chunkSize)
		}
	default:
	}

	return nil
}

// CleanUp deletes the file system components of the Morass. After this call
// the Morass is not usable.
func (m *Morass) CleanUp() error {
	return os.RemoveAll(m.dir)
}

// Pull sets the settable value e to the lowest value in the Morass.
// If io.EOF is returned the Morass is empty. Any other error results
// in no value being set on e.
func (m *Morass) Pull(e LessInterface) error {
	var err error
	v := reflect.ValueOf(e)
	if !reflect.Indirect(v).CanSet() {
		return errors.New("morass: cannot set e")
	}

	if m.fast {
		switch {
		case m.chunk != nil && m.pos < int64(len(m.chunk)):
			e = m.chunk[m.pos].(LessInterface)
			m.pos++
		case m.chunk != nil:
			m.pool <- m.chunk[:0]
			m.chunk = nil
			fallthrough
		default:
			if m.AutoClear {
				m.Clear()
			}
			err = io.EOF
		}
	} else {
		if m.files.Len() > 0 {
			low := heap.Pop(&m.files).(*file)
			e = low.head
			m.pos++
			switch err = low.decoder.Decode(&low.head); err {
			case nil:
				heap.Push(&m.files, low)
			case io.EOF:
				err = nil
				fallthrough
			default:
				low.file.Close()
				if m.AutoClear {
					os.Remove(low.file.Name())
				}
			}
		} else {
			if m.AutoClear {
				m.Clear()
			}
			if m.AutoClean {
				os.RemoveAll(m.dir)
			}
			err = io.EOF
		}
	}

	if err != io.EOF {
		reflect.Indirect(v).Set(reflect.ValueOf(e))
	}

	return err
}
<|endoftext|>


<commit_before>// Copyright 2013 Google Inc.  All rights reserved.
// Copyright 2016 the gousb Authors.  All rights reserved.
//
// 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.

// rawread attempts to read from the specified USB device.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"github.com/kylelemons/gousb/usb"
	"github.com/kylelemons/gousb/usbid"
)

var (
	vid      = flag.Uint("vid", 0, "VID of the device to which to connect. Exclusive with bus/addr flags.")
	pid      = flag.Uint("pid", 0, "PID of the device to which to connect. Exclusive with bus/addr flags.")
	bus      = flag.Uint("bus", 0, "Bus number for the device to which to connect. Exclusive with vid/pid flags.")
	addr     = flag.Uint("addr", 0, "Address of the device to which to connect. Exclusive with vid/pid flags.")
	config   = flag.Uint("config", 1, "Endpoint to which to connect")
	iface    = flag.Uint("interface", 0, "Endpoint to which to connect")
	setup    = flag.Uint("setup", 0, "Endpoint to which to connect")
	endpoint = flag.Uint("endpoint", 1, "Endpoint to which to connect")
	debug    = flag.Int("debug", 3, "Debug level for libusb")
	size     = flag.Uint("read_size", 1024, "Maximum number of bytes of data to read. Collected will be printed to STDOUT.")
)

func main() {
	flag.Parse()

	// Only one context should be needed for an application.  It should always be closed.
	ctx := usb.NewContext()
	defer ctx.Close()

	ctx.Debug(*debug)

	var devName string
	switch {
	case *vid == 0 && *pid == 0 && *bus == 0 && *addr == 0:
		log.Fatal("You need to specify the device, either through --vid/--pid flags or through --bus/--addr flags.")
	case (*vid > 0 || *pid > 0) && (*bus > 0 || *addr > 0):
		log.Fatal("You can't use --vid/--pid flags at the same time as --bus/--addr.")
	case *vid > 0 || *pid > 0:
		devName = fmt.Sprintf("VID:PID %04x:%04x", *vid, *pid)
	default:
		devName = fmt.Sprintf("bus:addr %d:%d", *bus, *addr)
	}

	log.Printf("Scanning for device %q...", devName)
	// ListDevices is used to find the devices to open.
	devs, err := ctx.ListDevices(func(desc *usb.Descriptor) bool {
		switch {
		case usb.ID(*vid) == desc.Vendor && usb.ID(*pid) == desc.Product:
			return true
		case uint8(*bus) == desc.Bus && uint8(*addr) == desc.Address:
			return true
		}
		return false
	})
	// All Devices returned from ListDevices must be closed.
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()

	// ListDevices can occasionally fail, so be sure to check its return value.
	if err != nil {
		log.Printf("Warning: ListDevices: %s.", err)
	}
	switch {
	case len(devs) == 0:
		log.Fatal("No matching devices found.")
	case len(devs) > 1:
		log.Printf("Warning: multiple devices found. Using bus %d, addr %d.", devs[0].Bus, devs[0].Address)
		for _, d := range devs[1:] {
			d.Close()
		}
		devs = devs[:1]
	}
	dev := devs[0]

	// The usbid package can be used to print out human readable information.
	log.Printf("  Protocol: %s\n", usbid.Classify(dev.Descriptor))

	// The configurations can be examined from the Descriptor, though they can only
	// be set once the device is opened.  All configuration references must be closed,
	// to free up the memory in libusb.
	for _, cfg := range dev.Configs {
		// This loop just uses more of the built-in and usbid pretty printing to list
		// the USB devices.
		log.Printf("  %s:\n", cfg)
		for _, alt := range cfg.Interfaces {
			log.Printf("    --------------\n")
			for _, iface := range alt.Setups {
				log.Printf("    %s\n", iface)
				log.Printf("      %s\n", usbid.Classify(iface))
				for _, end := range iface.Endpoints {
					log.Printf("      %s\n", end)
				}
			}
		}
		log.Printf("    --------------\n")
	}

	log.Printf("Connecting to endpoint...")
	ep, err := dev.OpenEndpoint(uint8(*config), uint8(*iface), uint8(*setup), uint8(*endpoint)|uint8(usb.ENDPOINT_DIR_IN))
	if err != nil {
		log.Fatalf("open: %s", err)
	}

	buf := make([]byte, *size)
	num, err := ep.Read(buf)
	if err != nil {
		log.Fatalf("Reading from device failed: %v", err)
	}
	log.Printf("Read %d bytes of data", num)
	os.Stdout.Write(buf[:num])
}
<commit_msg>Add a benchmark option.<commit_after>// Copyright 2013 Google Inc.  All rights reserved.
// Copyright 2016 the gousb Authors.  All rights reserved.
//
// 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.

// rawread attempts to read from the specified USB device.
package main

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

	"github.com/kylelemons/gousb/usb"
	"github.com/kylelemons/gousb/usbid"
)

var (
	vid      = flag.Uint("vid", 0, "VID of the device to which to connect. Exclusive with bus/addr flags.")
	pid      = flag.Uint("pid", 0, "PID of the device to which to connect. Exclusive with bus/addr flags.")
	bus      = flag.Uint("bus", 0, "Bus number for the device to which to connect. Exclusive with vid/pid flags.")
	addr     = flag.Uint("addr", 0, "Address of the device to which to connect. Exclusive with vid/pid flags.")
	config   = flag.Uint("config", 1, "Endpoint to which to connect")
	iface    = flag.Uint("interface", 0, "Endpoint to which to connect")
	setup    = flag.Uint("setup", 0, "Endpoint to which to connect")
	endpoint = flag.Uint("endpoint", 1, "Endpoint to which to connect")
	debug    = flag.Int("debug", 3, "Debug level for libusb")
	size     = flag.Uint("read_size", 1024, "Number of bytes of data to read in a single transaction.")
	bench    = flag.Bool("benchmark", false, "When true, keep reading from the endpoint and periodically report the measured throughput. If false,  only one read operation is performed and obtained data is printed to STDOUT.")
)

func main() {
	flag.Parse()

	// Only one context should be needed for an application.  It should always be closed.
	ctx := usb.NewContext()
	defer ctx.Close()

	ctx.Debug(*debug)

	var devName string
	switch {
	case *vid == 0 && *pid == 0 && *bus == 0 && *addr == 0:
		log.Fatal("You need to specify the device, either through --vid/--pid flags or through --bus/--addr flags.")
	case (*vid > 0 || *pid > 0) && (*bus > 0 || *addr > 0):
		log.Fatal("You can't use --vid/--pid flags at the same time as --bus/--addr.")
	case *vid > 0 || *pid > 0:
		devName = fmt.Sprintf("VID:PID %04x:%04x", *vid, *pid)
	default:
		devName = fmt.Sprintf("bus:addr %d:%d", *bus, *addr)
	}

	log.Printf("Scanning for device %q...", devName)
	// ListDevices is used to find the devices to open.
	devs, err := ctx.ListDevices(func(desc *usb.Descriptor) bool {
		switch {
		case usb.ID(*vid) == desc.Vendor && usb.ID(*pid) == desc.Product:
			return true
		case uint8(*bus) == desc.Bus && uint8(*addr) == desc.Address:
			return true
		}
		return false
	})
	// All Devices returned from ListDevices must be closed.
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()

	// ListDevices can occasionally fail, so be sure to check its return value.
	if err != nil {
		log.Printf("Warning: ListDevices: %s.", err)
	}
	switch {
	case len(devs) == 0:
		log.Fatal("No matching devices found.")
	case len(devs) > 1:
		log.Printf("Warning: multiple devices found. Using bus %d, addr %d.", devs[0].Bus, devs[0].Address)
		for _, d := range devs[1:] {
			d.Close()
		}
		devs = devs[:1]
	}
	dev := devs[0]

	// The usbid package can be used to print out human readable information.
	log.Printf("  Protocol: %s\n", usbid.Classify(dev.Descriptor))

	// The configurations can be examined from the Descriptor, though they can only
	// be set once the device is opened.  All configuration references must be closed,
	// to free up the memory in libusb.
	for _, cfg := range dev.Configs {
		// This loop just uses more of the built-in and usbid pretty printing to list
		// the USB devices.
		log.Printf("  %s:\n", cfg)
		for _, alt := range cfg.Interfaces {
			log.Printf("    --------------\n")
			for _, iface := range alt.Setups {
				log.Printf("    %s\n", iface)
				log.Printf("      %s\n", usbid.Classify(iface))
				for _, end := range iface.Endpoints {
					log.Printf("      %s\n", end)
				}
			}
		}
		log.Printf("    --------------\n")
	}

	log.Print("Connecting to endpoint...")
	ep, err := dev.OpenEndpoint(uint8(*config), uint8(*iface), uint8(*setup), uint8(*endpoint)|uint8(usb.ENDPOINT_DIR_IN))
	if err != nil {
		log.Fatalf("open: %s", err)
	}
	log.Print("Reading...")

	var total uint64
	if *bench {
		go func() {
			var last uint64
			var before = time.Now()
			for {
				time.Sleep(4 * time.Second)
				cur := atomic.LoadUint64(&total)
				now := time.Now()
				dur := now.Sub(before)
				log.Printf("%.2f B/s\n", float64(cur-last)/(float64(dur)/float64(time.Second)))
				before = now
				last = cur
			}
		}()
	}

	for {
		buf := make([]byte, *size)
		num, err := ep.Read(buf)
		if err != nil {
			log.Fatalf("Reading from device failed: %v", err)
		}
		if !*bench {
			log.Printf("Read %d bytes of data", num)
			os.Stdout.Write(buf[:num])
			return
		}
		atomic.AddUint64(&total, uint64(num))
	}
}
<|endoftext|>


<commit_before>package elastic

import ()

const (
	MAPPING           = "mapping"
	MAPPINGS          = "mappings"
	TYPE              = "type"
	INDEX             = "index"
	PROPERTIES        = "properties"
	MATCH             = "match"
	MatchMappingType  = "match_mapping_type"
	DynamicTemplates  = "dynamic_templates"
	DEFAULT           = "_default_"
	PositionOffsetGap = "position_offset_gap"
	// IndexAnalyzer index-time analyzer
	IndexAnalyzer = "index_analyzer"
	// SearchAnalyzer search-time analyzer
	SearchAnalyzer = "search_analyzer"
)

// Mapping maps between the json fields and how Elasticsearch store them
type Mapping struct {
	client *Elasticsearch
	parser *MappingResultParser
	url    string
	query  Dict
}

// NewMapping creates a new mapping query
func NewMapping() *Mapping {
	return &Mapping{
		query: make(Dict),
	}
}

// newMapping creates a new mapping query
func newMapping(client *Elasticsearch, url string) *Mapping {
	return &Mapping{
		client: client,
		parser: &MappingResultParser{},
		url:    url,
		query:  make(Dict),
	}
}

// Mapping creates request mappings between the json fields and how Elasticsearch store them
// GET /:index/:type/_mapping
func (client *Elasticsearch) Mapping(index, doctype string) *Mapping {
	url := client.request(index, doctype, -1, MAPPING)
	return newMapping(client, url)
}

// String returns a string representation of this mapping API
func (mapping *Mapping) String() string {
	return String(mapping.query)
}

// AddProperty adds a mapping for a type's property (e.g. type, index, analyzer, etc.)
func (mapping *Mapping) AddProperty(fieldname, propertyname string, propertyvalue interface{}) *Mapping {
	if mapping.query[PROPERTIES] == nil {
		mapping.query[PROPERTIES] = make(Dict)
	}
	property := mapping.query[PROPERTIES].(Dict)[fieldname]
	if property == nil {
		property = make(Dict)
	}
	property.(Dict)[propertyname] = propertyvalue
	mapping.query[PROPERTIES].(Dict)[fieldname] = property
	return mapping
}

// AddField adds a mapping for a field
func (mapping *Mapping) AddField(name string, body Dict) *Mapping {
	if mapping.query[PROPERTIES] == nil {
		mapping.query[PROPERTIES] = make(Dict)
	}
	mapping.query[PROPERTIES].(Dict)[name] = body
	return mapping
}

// AddDocumentType adds a mapping for a type of objects
func (mapping *Mapping) AddDocumentType(class *DocType) *Mapping {
	if mapping.query[MAPPINGS] == nil {
		mapping.query[MAPPINGS] = Dict{}
	}
	mapping.query[MAPPINGS].(Dict)[class.name] = class.dict
	return mapping
}

// Get submits a get request mappings between the json fields and how Elasticsearch store them
// GET /:index/_mapping/:type
func (mapping *Mapping) Get() {
	mapping.client.Execute("GET", mapping.url, "", mapping.parser)
}

// Put submits a request for updating the mappings between the json fields and how Elasticsearch store them
// PUT /:index/_mapping/:type
func (mapping *Mapping) Put() {
	url := mapping.url
	query := mapping.String()
	mapping.client.Execute("PUT", url, query, mapping.parser)
}

// DocType a structure for document type
type DocType struct {
	name string
	dict Dict
}

// NewDefaultType returns a '_default_' type that encapsulates shared/default settings
// e.g. specify index wide dynamic templates
func NewDefaultType() *DocType {
	return NewDocType(DEFAULT)
}

// NewDocType  a new mapping template
func NewDocType(name string) *DocType {
	return &DocType{name: name, dict: make(Dict)}
}

// AddProperty adds a property to this document type
func (doctype *DocType) AddProperty(name string, value interface{}) *DocType {
	doctype.dict[name] = value
	return doctype
}

// AddTemplate adds a template to this document type
func (doctype *DocType) AddTemplate(tmpl *Template) *DocType {
	doctype.dict[tmpl.name] = tmpl.dict
	return doctype
}

// AddDynamicTemplate adds a dynamic template to this mapping
func (doctype *DocType) AddDynamicTemplate(tmpl *Template) *DocType {
	if doctype.dict[DynamicTemplates] == nil {
		doctype.dict[DynamicTemplates] = []Dict{}
	}
	dict := make(Dict)
	dict[tmpl.name] = tmpl.dict
	doctype.dict[DynamicTemplates] = append(doctype.dict[DynamicTemplates].([]Dict), dict)
	return doctype
}

// String returns a string representation of this document type
func (doctype *DocType) String() string {
	dict := make(Dict)
	dict[doctype.name] = doctype.dict
	return String(dict)
}

// Template a structure for mapping template
type Template struct {
	name string
	dict Dict
}

// NewAllTemplate returns an new '_all' template
func NewAllTemplate() *Template {
	return NewTemplate(ALL)
}

// NewTemplate creates a new named mapping template
func NewTemplate(name string) *Template {
	return &Template{name: name, dict: make(Dict)}
}

// AddMatch adds a match string (e.g. '*', '_es')
func (template *Template) AddMatch(match string) *Template {
	template.dict[MATCH] = match
	return template
}

// AddProperty adds a property to this template
func (template *Template) AddProperty(name string, value interface{}) *Template {
	template.dict[name] = value
	return template
}

// AddMappingProperty adds a property to the `mapping` object
func (template *Template) AddMappingProperty(name string, value interface{}) *Template {
	if template.dict[MAPPING] == nil {
		template.dict[MAPPING] = make(Dict)
	}
	template.dict[MAPPING].(Dict)[name] = value
	return template
}

// String returns a string representation of this template
func (template *Template) String() string {
	dict := make(Dict)
	dict[template.name] = template.dict
	return String(dict)
}
<commit_msg>fix for golint<commit_after>package elastic

import ()

const (
	// MAPPING part of Mapping API path url
	MAPPING = "mapping"
	// MAPPINGS body of Mapping API query
	MAPPINGS = "mappings"
	// TYPE constant name of data type property of field
	TYPE = "type"
	// INDEX constant name of index name
	INDEX = "index"
	// PROPERTIES constant name of Mapping query body that defines properties
	PROPERTIES = "properties"
	// MATCH
	MATCH = "match"
	// MatchMappingType
	MatchMappingType = "match_mapping_type"
	// DynamicTemplates
	DynamicTemplates = "dynamic_templates"
	// DEFAULT
	DEFAULT = "_default_"
	// PositionOffsetGap constant name for defining acceptable offset gap
	PositionOffsetGap = "position_offset_gap"
	// IndexAnalyzer index-time analyzer
	IndexAnalyzer = "index_analyzer"
	// SearchAnalyzer search-time analyzer
	SearchAnalyzer = "search_analyzer"
	// IndexOptions defines indexing options in Mapping query
	IndexOptions = "index_options"
	// Norms constant name for configuring field length normalization
	Norms = "norms"
)

// Mapping maps between the json fields and how Elasticsearch store them
type Mapping struct {
	client *Elasticsearch
	parser *MappingResultParser
	url    string
	query  Dict
}

// NewMapping creates a new mapping query
func NewMapping() *Mapping {
	return &Mapping{
		query: make(Dict),
	}
}

// newMapping creates a new mapping query
func newMapping(client *Elasticsearch, url string) *Mapping {
	return &Mapping{
		client: client,
		parser: &MappingResultParser{},
		url:    url,
		query:  make(Dict),
	}
}

// Mapping creates request mappings between the json fields and how Elasticsearch store them
// GET /:index/:type/_mapping
func (client *Elasticsearch) Mapping(index, doctype string) *Mapping {
	url := client.request(index, doctype, -1, MAPPING)
	return newMapping(client, url)
}

// String returns a string representation of this mapping API
func (mapping *Mapping) String() string {
	return String(mapping.query)
}

// AddProperty adds a mapping for a type's property (e.g. type, index, analyzer, etc.)
func (mapping *Mapping) AddProperty(fieldname, propertyname string, propertyvalue interface{}) *Mapping {
	if mapping.query[PROPERTIES] == nil {
		mapping.query[PROPERTIES] = make(Dict)
	}
	property := mapping.query[PROPERTIES].(Dict)[fieldname]
	if property == nil {
		property = make(Dict)
	}
	property.(Dict)[propertyname] = propertyvalue
	mapping.query[PROPERTIES].(Dict)[fieldname] = property
	return mapping
}

// AddField adds a mapping for a field
func (mapping *Mapping) AddField(name string, body Dict) *Mapping {
	if mapping.query[PROPERTIES] == nil {
		mapping.query[PROPERTIES] = make(Dict)
	}
	mapping.query[PROPERTIES].(Dict)[name] = body
	return mapping
}

// AddDocumentType adds a mapping for a type of objects
func (mapping *Mapping) AddDocumentType(class *DocType) *Mapping {
	if mapping.query[MAPPINGS] == nil {
		mapping.query[MAPPINGS] = Dict{}
	}
	mapping.query[MAPPINGS].(Dict)[class.name] = class.dict
	return mapping
}

// Get submits a get request mappings between the json fields and how Elasticsearch store them
// GET /:index/_mapping/:type
func (mapping *Mapping) Get() {
	mapping.client.Execute("GET", mapping.url, "", mapping.parser)
}

// Put submits a request for updating the mappings between the json fields and how Elasticsearch store them
// PUT /:index/_mapping/:type
func (mapping *Mapping) Put() {
	url := mapping.url
	query := mapping.String()
	mapping.client.Execute("PUT", url, query, mapping.parser)
}

// DocType a structure for document type
type DocType struct {
	name string
	dict Dict
}

// NewDefaultType returns a '_default_' type that encapsulates shared/default settings
// e.g. specify index wide dynamic templates
func NewDefaultType() *DocType {
	return NewDocType(DEFAULT)
}

// NewDocType  a new mapping template
func NewDocType(name string) *DocType {
	return &DocType{name: name, dict: make(Dict)}
}

// AddProperty adds a property to this document type
func (doctype *DocType) AddProperty(name string, value interface{}) *DocType {
	doctype.dict[name] = value
	return doctype
}

// AddTemplate adds a template to this document type
func (doctype *DocType) AddTemplate(tmpl *Template) *DocType {
	doctype.dict[tmpl.name] = tmpl.dict
	return doctype
}

// AddDynamicTemplate adds a dynamic template to this mapping
func (doctype *DocType) AddDynamicTemplate(tmpl *Template) *DocType {
	if doctype.dict[DynamicTemplates] == nil {
		doctype.dict[DynamicTemplates] = []Dict{}
	}
	dict := make(Dict)
	dict[tmpl.name] = tmpl.dict
	doctype.dict[DynamicTemplates] = append(doctype.dict[DynamicTemplates].([]Dict), dict)
	return doctype
}

// String returns a string representation of this document type
func (doctype *DocType) String() string {
	dict := make(Dict)
	dict[doctype.name] = doctype.dict
	return String(dict)
}

// Template a structure for mapping template
type Template struct {
	name string
	dict Dict
}

// NewAllTemplate returns an new '_all' template
func NewAllTemplate() *Template {
	return NewTemplate(ALL)
}

// NewTemplate creates a new named mapping template
func NewTemplate(name string) *Template {
	return &Template{name: name, dict: make(Dict)}
}

// AddMatch adds a match string (e.g. '*', '_es')
func (template *Template) AddMatch(match string) *Template {
	template.dict[MATCH] = match
	return template
}

// AddProperty adds a property to this template
func (template *Template) AddProperty(name string, value interface{}) *Template {
	template.dict[name] = value
	return template
}

// AddMappingProperty adds a property to the `mapping` object
func (template *Template) AddMappingProperty(name string, value interface{}) *Template {
	if template.dict[MAPPING] == nil {
		template.dict[MAPPING] = make(Dict)
	}
	template.dict[MAPPING].(Dict)[name] = value
	return template
}

// String returns a string representation of this template
func (template *Template) String() string {
	dict := make(Dict)
	dict[template.name] = template.dict
	return String(dict)
}
<|endoftext|>


<commit_before>package pangu

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"regexp"
	"text/template"
)

const VERSION = "2.5.6"

// CJK is short for Chinese, Japanese and Korean.
//
// The constant cjk contains following Unicode blocks:
// 	\u2e80-\u2eff CJK Radicals Supplement
// 	\u2f00-\u2fdf Kangxi Radicals
// 	\u3040-\u309f Hiragana
// 	\u30a0-\u30ff Katakana
// 	\u3100-\u312f Bopomofo
// 	\u3200-\u32ff Enclosed CJK Letters and Months
// 	\u3400-\u4dbf CJK Unified Ideographs Extension A
// 	\u4e00-\u9fff CJK Unified Ideographs
// 	\uf900-\ufaff CJK Compatibility Ideographs
//
// For more information about Unicode blocks, see
// 	http://unicode-table.com/en/
const cjk = "" +
	"\u2e80-\u2eff" +
	"\u2f00-\u2fdf" +
	"\u3040-\u309f" +
	"\u30a0-\u30ff" +
	"\u3100-\u312f" +
	"\u3200-\u32ff" +
	"\u3400-\u4dbf" +
	"\u4e00-\u9fff" +
	"\uf900-\ufaff"

// ANS is short for Alphabets, Numbers
// and Symbols (`~!@#$%^&*()-_=+[]{}\|;:'",<.>/?).
//
// The constant ans doesn't contain all symbols above.
const ans = "A-Za-z0-9`\\$%\\^&\\*\\-=\\+\\\\|/\u00a1-\u00ff\u2022\u2027\u2150-\u218f"

var cjk_quote = regexp.MustCompile(re("([{{ .CJK }}])" + "([\"'])"))
var quote_cjk = regexp.MustCompile(re("([\"'])" + "([{{ .CJK }}])"))
var fix_quote = regexp.MustCompile(re("([\"'\\(\\[\\{<\u201c])" + "(\\s*)" + "(.+?)" + "(\\s*)" + "([\"'\\)\\]\\}>\u201d])"))
var fix_single_quote = regexp.MustCompile(re("([{{ .CJK }}])" + "( )" + "(')" + "([A-Za-z])"))

var cjk_hash = regexp.MustCompile(re("([{{ .CJK }}])" + "(#(\\S+))"))
var hash_cjk = regexp.MustCompile(re("((\\S+)#)" + "([{{ .CJK }}])"))

var cjk_operator_ans = regexp.MustCompile(re("([{{ .CJK }}])" + "([\\+\\-\\*/=&\\|<>])" + "([A-Za-z0-9])"))
var ans_operator_cjk = regexp.MustCompile(re("([A-Za-z0-9])" + "([\\+\\-\\*/=&\\|<>])" + "([{{ .CJK }}])"))

var cjk_bracket_cjk = regexp.MustCompile(re("([{{ .CJK }}])" + "([\\(\\[\\{<\u201c]+(.*?)[\\)\\]\\}>\u201d]+)" + "([{{ .CJK }}])"))
var cjk_bracket = regexp.MustCompile(re("([{{ .CJK }}])" + "([\\(\\[\\{<\u201c>])"))
var bracket_cjk = regexp.MustCompile(re("([\\)\\]\\}>\u201d<])" + "([{{ .CJK }}])"))
var fix_bracket = regexp.MustCompile(re("([\\(\\[\\{<\u201c]+)" + "(\\s*)" + "(.+?)" + "(\\s*)" + "([\\)\\]\\}>\u201d]+)"))

var fix_symbol = regexp.MustCompile(re("([{{ .CJK }}])" + "([~!;:,\\.\\?\u2026])" + "([A-Za-z0-9])"))

var cjk_ans = regexp.MustCompile(re("([{{ .CJK }}])([{{ .ANS }}@])"))
var ans_cjk = regexp.MustCompile(re("([{{ .ANS }}~!;:,\\.\\?\u2026])([{{ .CJK }}])"))

var context = map[string]string{
	"CJK": cjk,
	"ANS": ans,
}

func re(exp string) string {
	var buf bytes.Buffer

	var tmpl = template.New("pangu")
	tmpl, _ = tmpl.Parse(exp)
	tmpl.Execute(&buf, context)
	expr := buf.String()

	return expr
}

// TextSpacing performs paranoid text spacing on text.
// It returns the processed text, with love.
func TextSpacing(text string) string {
	if len(text) < 2 {
		return text
	}

	text = cjk_quote.ReplaceAllString(text, "$1 $2")
	text = quote_cjk.ReplaceAllString(text, "$1 $2")
	text = fix_quote.ReplaceAllString(text, "$1$3$5")
	text = fix_single_quote.ReplaceAllString(text, "$1$3$4")

	text = cjk_hash.ReplaceAllString(text, "$1 $2")
	text = hash_cjk.ReplaceAllString(text, "$1 $3")

	text = cjk_operator_ans.ReplaceAllString(text, "$1 $2 $3")
	text = ans_operator_cjk.ReplaceAllString(text, "$1 $2 $3")

	oldText := text
	newText := cjk_bracket_cjk.ReplaceAllString(oldText, "$1 $2 $4")
	text = newText
	if oldText == newText {
		text = cjk_bracket.ReplaceAllString(text, "$1 $2")
		text = bracket_cjk.ReplaceAllString(text, "$1 $2")
	}
	text = fix_bracket.ReplaceAllString(text, "$1$3$5")

	text = fix_symbol.ReplaceAllString(text, "$1$2 $3")

	text = cjk_ans.ReplaceAllString(text, "$1 $2")
	text = ans_cjk.ReplaceAllString(text, "$1 $2")

	return text
}

// FileSpacing reads the file named by filename, performs paranoid text
// spacing on its contents and writes the processed content to w.
// A successful call returns err == nil.
func FileSpacing(filename string, w io.Writer) (err error) {
	fr, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fr.Close()

	br := bufio.NewReader(fr)
	bw := bufio.NewWriter(w)

	for {
		line, err := br.ReadString('\n')
		if err == nil {
			fmt.Fprint(bw, TextSpacing(line))
		} else {
			if err == io.EOF {
				fmt.Fprint(bw, TextSpacing(line))
				break
			}
			return err
		}
	}
	defer bw.Flush()

	return nil
}
<commit_msg>add comment<commit_after>package pangu

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"regexp"
	"text/template"
)

// VERSION is the version number of pangu
const VERSION = "2.5.6"

// CJK is short for Chinese, Japanese and Korean.
//
// The constant cjk contains following Unicode blocks:
// 	\u2e80-\u2eff CJK Radicals Supplement
// 	\u2f00-\u2fdf Kangxi Radicals
// 	\u3040-\u309f Hiragana
// 	\u30a0-\u30ff Katakana
// 	\u3100-\u312f Bopomofo
// 	\u3200-\u32ff Enclosed CJK Letters and Months
// 	\u3400-\u4dbf CJK Unified Ideographs Extension A
// 	\u4e00-\u9fff CJK Unified Ideographs
// 	\uf900-\ufaff CJK Compatibility Ideographs
//
// For more information about Unicode blocks, see
// 	http://unicode-table.com/en/
const cjk = "" +
	"\u2e80-\u2eff" +
	"\u2f00-\u2fdf" +
	"\u3040-\u309f" +
	"\u30a0-\u30ff" +
	"\u3100-\u312f" +
	"\u3200-\u32ff" +
	"\u3400-\u4dbf" +
	"\u4e00-\u9fff" +
	"\uf900-\ufaff"

// ANS is short for Alphabets, Numbers
// and Symbols (`~!@#$%^&*()-_=+[]{}\|;:'",<.>/?).
//
// The constant ans doesn't contain all symbols above.
const ans = "A-Za-z0-9`\\$%\\^&\\*\\-=\\+\\\\|/\u00a1-\u00ff\u2022\u2027\u2150-\u218f"

var cjk_quote = regexp.MustCompile(re("([{{ .CJK }}])" + "([\"'])"))
var quote_cjk = regexp.MustCompile(re("([\"'])" + "([{{ .CJK }}])"))
var fix_quote = regexp.MustCompile(re("([\"'\\(\\[\\{<\u201c])" + "(\\s*)" + "(.+?)" + "(\\s*)" + "([\"'\\)\\]\\}>\u201d])"))
var fix_single_quote = regexp.MustCompile(re("([{{ .CJK }}])" + "( )" + "(')" + "([A-Za-z])"))

var cjk_hash = regexp.MustCompile(re("([{{ .CJK }}])" + "(#(\\S+))"))
var hash_cjk = regexp.MustCompile(re("((\\S+)#)" + "([{{ .CJK }}])"))

var cjk_operator_ans = regexp.MustCompile(re("([{{ .CJK }}])" + "([\\+\\-\\*/=&\\|<>])" + "([A-Za-z0-9])"))
var ans_operator_cjk = regexp.MustCompile(re("([A-Za-z0-9])" + "([\\+\\-\\*/=&\\|<>])" + "([{{ .CJK }}])"))

var cjk_bracket_cjk = regexp.MustCompile(re("([{{ .CJK }}])" + "([\\(\\[\\{<\u201c]+(.*?)[\\)\\]\\}>\u201d]+)" + "([{{ .CJK }}])"))
var cjk_bracket = regexp.MustCompile(re("([{{ .CJK }}])" + "([\\(\\[\\{<\u201c>])"))
var bracket_cjk = regexp.MustCompile(re("([\\)\\]\\}>\u201d<])" + "([{{ .CJK }}])"))
var fix_bracket = regexp.MustCompile(re("([\\(\\[\\{<\u201c]+)" + "(\\s*)" + "(.+?)" + "(\\s*)" + "([\\)\\]\\}>\u201d]+)"))

var fix_symbol = regexp.MustCompile(re("([{{ .CJK }}])" + "([~!;:,\\.\\?\u2026])" + "([A-Za-z0-9])"))

var cjk_ans = regexp.MustCompile(re("([{{ .CJK }}])([{{ .ANS }}@])"))
var ans_cjk = regexp.MustCompile(re("([{{ .ANS }}~!;:,\\.\\?\u2026])([{{ .CJK }}])"))

var context = map[string]string{
	"CJK": cjk,
	"ANS": ans,
}

func re(exp string) string {
	var buf bytes.Buffer

	var tmpl = template.New("pangu")
	tmpl, _ = tmpl.Parse(exp)
	tmpl.Execute(&buf, context)
	expr := buf.String()

	return expr
}

// TextSpacing performs paranoid text spacing on text.
// It returns the processed text, with love.
func TextSpacing(text string) string {
	if len(text) < 2 {
		return text
	}

	text = cjk_quote.ReplaceAllString(text, "$1 $2")
	text = quote_cjk.ReplaceAllString(text, "$1 $2")
	text = fix_quote.ReplaceAllString(text, "$1$3$5")
	text = fix_single_quote.ReplaceAllString(text, "$1$3$4")

	text = cjk_hash.ReplaceAllString(text, "$1 $2")
	text = hash_cjk.ReplaceAllString(text, "$1 $3")

	text = cjk_operator_ans.ReplaceAllString(text, "$1 $2 $3")
	text = ans_operator_cjk.ReplaceAllString(text, "$1 $2 $3")

	oldText := text
	newText := cjk_bracket_cjk.ReplaceAllString(oldText, "$1 $2 $4")
	text = newText
	if oldText == newText {
		text = cjk_bracket.ReplaceAllString(text, "$1 $2")
		text = bracket_cjk.ReplaceAllString(text, "$1 $2")
	}
	text = fix_bracket.ReplaceAllString(text, "$1$3$5")

	text = fix_symbol.ReplaceAllString(text, "$1$2 $3")

	text = cjk_ans.ReplaceAllString(text, "$1 $2")
	text = ans_cjk.ReplaceAllString(text, "$1 $2")

	return text
}

// FileSpacing reads the file named by filename, performs paranoid text
// spacing on its contents and writes the processed content to w.
// A successful call returns err == nil.
func FileSpacing(filename string, w io.Writer) (err error) {
	fr, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fr.Close()

	br := bufio.NewReader(fr)
	bw := bufio.NewWriter(w)

	for {
		line, err := br.ReadString('\n')
		if err == nil {
			fmt.Fprint(bw, TextSpacing(line))
		} else {
			if err == io.EOF {
				fmt.Fprint(bw, TextSpacing(line))
				break
			}
			return err
		}
	}
	defer bw.Flush()

	return nil
}
<|endoftext|>