<commit_before>package mdns

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/miekg/dns"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

// ServiceEntry is returned after we query for a service
type ServiceEntry struct {
	Name   string
	Host   string
	AddrV4 net.IP
	AddrV6 net.IP
	Port   int
	Info   string

	Addr net.IP // @Deprecated

	hasTXT bool
	sent   bool
}

// complete is used to check if we have all the info we need
func (s *ServiceEntry) complete() bool {
	return (s.AddrV4 != nil || s.AddrV6 != nil || s.Addr != nil) && s.Port != 0 && s.hasTXT
}

// QueryParam is used to customize how a Lookup is performed
type QueryParam struct {
	Service             string               // Service to lookup
	Domain              string               // Lookup domain, default "local"
	Timeout             time.Duration        // Lookup timeout, default 1 second
	Interface           *net.Interface       // Multicast interface to use
	Entries             chan<- *ServiceEntry // Entries Channel
	WantUnicastResponse bool                 // Unicast response desired, as per 5.4 in RFC
}

// DefaultParams is used to return a default set of QueryParam's
func DefaultParams(service string) *QueryParam {
	return &QueryParam{
		Service:             service,
		Domain:              "local",
		Timeout:             time.Second,
		Entries:             make(chan *ServiceEntry),
		WantUnicastResponse: false, // TODO(reddaly): Change this default.
	}
}

// Query looks up a given service, in a domain, waiting at most
// for a timeout before finishing the query. The results are streamed
// to a channel. Sends will not block, so clients should make sure to
// either read or buffer.
func Query(params *QueryParam) error {
	// Create a new client
	client, err := newClient()
	if err != nil {
		return err
	}
	defer client.Close()

	// Set the multicast interface
	if params.Interface != nil {
		if err := client.setInterface(params.Interface); err != nil {
			return err
		}
	}

	// Ensure defaults are set
	if params.Domain == "" {
		params.Domain = "local"
	}
	if params.Timeout == 0 {
		params.Timeout = time.Second
	}

	// Run the query
	return client.query(params)
}

// Lookup is the same as Query, however it uses all the default parameters
func Lookup(service string, entries chan<- *ServiceEntry) error {
	params := DefaultParams(service)
	params.Entries = entries
	return Query(params)
}

// Client provides a query interface that can be used to
// search for service providers using mDNS
type client struct {
	ipv4UnicastConn *net.UDPConn
	ipv6UnicastConn *net.UDPConn

	ipv4MulticastConn *net.UDPConn
	ipv6MulticastConn *net.UDPConn

	closed    bool
	closedCh  chan struct{} // TODO(reddaly): This doesn't appear to be used.
	closeLock sync.Mutex
}

// NewClient creates a new mdns Client that can be used to query
// for records
func newClient() (*client, error) {
	// TODO(reddaly): At least attempt to bind to the port required in the spec.
	// Create a IPv4 listener
	uconn4, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	uconn6, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if uconn4 == nil && uconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any unicast udp port")
	}

	mconn4, err := net.ListenMulticastUDP("udp4", nil, ipv4Addr)
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	mconn6, err := net.ListenMulticastUDP("udp6", nil, ipv6Addr)
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if mconn4 == nil && mconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any multicast udp port")
	}

	c := &client{
		ipv4MulticastConn: mconn4,
		ipv6MulticastConn: mconn6,
		ipv4UnicastConn:   uconn4,
		ipv6UnicastConn:   uconn6,
		closedCh:          make(chan struct{}),
	}
	return c, nil
}

// Close is used to cleanup the client
func (c *client) Close() error {
	c.closeLock.Lock()
	defer c.closeLock.Unlock()

	if c.closed {
		return nil
	}
	c.closed = true

	log.Printf("[INFO] mdns: Closing client %v", *c)
	close(c.closedCh)

	if c.ipv4UnicastConn != nil {
		c.ipv4UnicastConn.Close()
	}
	if c.ipv6UnicastConn != nil {
		c.ipv6UnicastConn.Close()
	}
	if c.ipv4MulticastConn != nil {
		c.ipv4MulticastConn.Close()
	}
	if c.ipv6MulticastConn != nil {
		c.ipv6MulticastConn.Close()
	}

	return nil
}

// setInterface is used to set the query interface, uses sytem
// default if not provided
func (c *client) setInterface(iface *net.Interface) error {
	p := ipv4.NewPacketConn(c.ipv4UnicastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 := ipv6.NewPacketConn(c.ipv6UnicastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	p = ipv4.NewPacketConn(c.ipv4MulticastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 = ipv6.NewPacketConn(c.ipv6MulticastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	return nil
}

// query is used to perform a lookup and stream results
func (c *client) query(params *QueryParam) error {
	// Create the service name
	serviceAddr := fmt.Sprintf("%s.%s.", trimDot(params.Service), trimDot(params.Domain))

	// Start listening for response packets
	msgCh := make(chan *dns.Msg, 32)
	go c.recv(c.ipv4UnicastConn, msgCh)
	go c.recv(c.ipv6UnicastConn, msgCh)
	go c.recv(c.ipv4MulticastConn, msgCh)
	go c.recv(c.ipv6MulticastConn, msgCh)

	// Send the query
	m := new(dns.Msg)
	m.SetQuestion(serviceAddr, dns.TypePTR)
	// RFC 6762, section 18.12.  Repurposing of Top Bit of qclass in Question
	// Section
	//
	// In the Question Section of a Multicast DNS query, the top bit of the qclass
	// field is used to indicate that unicast responses are preferred for this
	// particular question.  (See Section 5.4.)
	if params.WantUnicastResponse {
		m.Question[0].Qclass |= 1 << 15
	}
	m.RecursionDesired = false
	if err := c.sendQuery(m); err != nil {
		return err
	}

	// Map the in-progress responses
	inprogress := make(map[string]*ServiceEntry)

	// Listen until we reach the timeout
	finish := time.After(params.Timeout)
	for {
		select {
		case resp := <-msgCh:
			var inp *ServiceEntry
			for _, answer := range resp.Answer {
				// TODO(reddaly): Check that response corresponds to serviceAddr?
				switch rr := answer.(type) {
				case *dns.PTR:
					// Create new entry for this
					inp = ensureName(inprogress, rr.Ptr)

				case *dns.SRV:
					// Check for a target mismatch
					if rr.Target != rr.Hdr.Name {
						alias(inprogress, rr.Hdr.Name, rr.Target)
					}

					// Get the port
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Host = rr.Target
					inp.Port = int(rr.Port)

				case *dns.TXT:
					// Pull out the txt
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Info = strings.Join(rr.Txt, "|")
					inp.hasTXT = true

				case *dns.A:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.A // @Deprecated
					inp.AddrV4 = rr.A

				case *dns.AAAA:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.AAAA // @Deprecated
					inp.AddrV6 = rr.AAAA
				}
			}

			if inp == nil {
				continue
			}

			// Check if this entry is complete
			if inp.complete() {
				if inp.sent {
					continue
				}
				inp.sent = true
				select {
				case params.Entries <- inp:
				default:
				}
			} else {
				// Fire off a node specific query
				m := new(dns.Msg)
				m.SetQuestion(inp.Name, dns.TypePTR)
				m.RecursionDesired = false
				if err := c.sendQuery(m); err != nil {
					log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
				}
			}
		case <-finish:
			return nil
		}
	}
}

// sendQuery is used to multicast a query out
func (c *client) sendQuery(q *dns.Msg) error {
	buf, err := q.Pack()
	if err != nil {
		return err
	}
	if c.ipv4UnicastConn != nil {
		c.ipv4UnicastConn.WriteToUDP(buf, ipv4Addr)
	}
	if c.ipv6UnicastConn != nil {
		c.ipv6UnicastConn.WriteToUDP(buf, ipv6Addr)
	}
	return nil
}

// recv is used to receive until we get a shutdown
func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) {
	if l == nil {
		return
	}
	buf := make([]byte, 65536)
	for !c.closed {
		n, err := l.Read(buf)
		if err != nil {
			log.Printf("[ERR] mdns: Failed to read packet: %v", err)
			continue
		}
		msg := new(dns.Msg)
		if err := msg.Unpack(buf[:n]); err != nil {
			log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
			continue
		}
		select {
		case msgCh <- msg:
		case <-c.closedCh:
			return
		}
	}
}

// ensureName is used to ensure the named node is in progress
func ensureName(inprogress map[string]*ServiceEntry, name string) *ServiceEntry {
	if inp, ok := inprogress[name]; ok {
		return inp
	}
	inp := &ServiceEntry{
		Name: name,
	}
	inprogress[name] = inp
	return inp
}

// alias is used to setup an alias between two entries
func alias(inprogress map[string]*ServiceEntry, src, dst string) {
	srcEntry := ensureName(inprogress, src)
	inprogress[dst] = srcEntry
}
<commit_msg>Revert "Change import path from code.google.com/p/go.net/context to golang.org/x/net/context."<commit_after>package mdns

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"code.google.com/p/go.net/ipv4"
	"code.google.com/p/go.net/ipv6"
	"github.com/miekg/dns"
)

// ServiceEntry is returned after we query for a service
type ServiceEntry struct {
	Name   string
	Host   string
	AddrV4 net.IP
	AddrV6 net.IP
	Port   int
	Info   string

	Addr net.IP // @Deprecated

	hasTXT bool
	sent   bool
}

// complete is used to check if we have all the info we need
func (s *ServiceEntry) complete() bool {
	return (s.AddrV4 != nil || s.AddrV6 != nil || s.Addr != nil) && s.Port != 0 && s.hasTXT
}

// QueryParam is used to customize how a Lookup is performed
type QueryParam struct {
	Service             string               // Service to lookup
	Domain              string               // Lookup domain, default "local"
	Timeout             time.Duration        // Lookup timeout, default 1 second
	Interface           *net.Interface       // Multicast interface to use
	Entries             chan<- *ServiceEntry // Entries Channel
	WantUnicastResponse bool                 // Unicast response desired, as per 5.4 in RFC
}

// DefaultParams is used to return a default set of QueryParam's
func DefaultParams(service string) *QueryParam {
	return &QueryParam{
		Service:             service,
		Domain:              "local",
		Timeout:             time.Second,
		Entries:             make(chan *ServiceEntry),
		WantUnicastResponse: false, // TODO(reddaly): Change this default.
	}
}

// Query looks up a given service, in a domain, waiting at most
// for a timeout before finishing the query. The results are streamed
// to a channel. Sends will not block, so clients should make sure to
// either read or buffer.
func Query(params *QueryParam) error {
	// Create a new client
	client, err := newClient()
	if err != nil {
		return err
	}
	defer client.Close()

	// Set the multicast interface
	if params.Interface != nil {
		if err := client.setInterface(params.Interface); err != nil {
			return err
		}
	}

	// Ensure defaults are set
	if params.Domain == "" {
		params.Domain = "local"
	}
	if params.Timeout == 0 {
		params.Timeout = time.Second
	}

	// Run the query
	return client.query(params)
}

// Lookup is the same as Query, however it uses all the default parameters
func Lookup(service string, entries chan<- *ServiceEntry) error {
	params := DefaultParams(service)
	params.Entries = entries
	return Query(params)
}

// Client provides a query interface that can be used to
// search for service providers using mDNS
type client struct {
	ipv4UnicastConn *net.UDPConn
	ipv6UnicastConn *net.UDPConn

	ipv4MulticastConn *net.UDPConn
	ipv6MulticastConn *net.UDPConn

	closed    bool
	closedCh  chan struct{} // TODO(reddaly): This doesn't appear to be used.
	closeLock sync.Mutex
}

// NewClient creates a new mdns Client that can be used to query
// for records
func newClient() (*client, error) {
	// TODO(reddaly): At least attempt to bind to the port required in the spec.
	// Create a IPv4 listener
	uconn4, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	uconn6, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if uconn4 == nil && uconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any unicast udp port")
	}

	mconn4, err := net.ListenMulticastUDP("udp4", nil, ipv4Addr)
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	mconn6, err := net.ListenMulticastUDP("udp6", nil, ipv6Addr)
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if mconn4 == nil && mconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any multicast udp port")
	}

	c := &client{
		ipv4MulticastConn: mconn4,
		ipv6MulticastConn: mconn6,
		ipv4UnicastConn:   uconn4,
		ipv6UnicastConn:   uconn6,
		closedCh:          make(chan struct{}),
	}
	return c, nil
}

// Close is used to cleanup the client
func (c *client) Close() error {
	c.closeLock.Lock()
	defer c.closeLock.Unlock()

	if c.closed {
		return nil
	}
	c.closed = true

	log.Printf("[INFO] mdns: Closing client %v", *c)
	close(c.closedCh)

	if c.ipv4UnicastConn != nil {
		c.ipv4UnicastConn.Close()
	}
	if c.ipv6UnicastConn != nil {
		c.ipv6UnicastConn.Close()
	}
	if c.ipv4MulticastConn != nil {
		c.ipv4MulticastConn.Close()
	}
	if c.ipv6MulticastConn != nil {
		c.ipv6MulticastConn.Close()
	}

	return nil
}

// setInterface is used to set the query interface, uses sytem
// default if not provided
func (c *client) setInterface(iface *net.Interface) error {
	p := ipv4.NewPacketConn(c.ipv4UnicastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 := ipv6.NewPacketConn(c.ipv6UnicastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	p = ipv4.NewPacketConn(c.ipv4MulticastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 = ipv6.NewPacketConn(c.ipv6MulticastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	return nil
}

// query is used to perform a lookup and stream results
func (c *client) query(params *QueryParam) error {
	// Create the service name
	serviceAddr := fmt.Sprintf("%s.%s.", trimDot(params.Service), trimDot(params.Domain))

	// Start listening for response packets
	msgCh := make(chan *dns.Msg, 32)
	go c.recv(c.ipv4UnicastConn, msgCh)
	go c.recv(c.ipv6UnicastConn, msgCh)
	go c.recv(c.ipv4MulticastConn, msgCh)
	go c.recv(c.ipv6MulticastConn, msgCh)

	// Send the query
	m := new(dns.Msg)
	m.SetQuestion(serviceAddr, dns.TypePTR)
	// RFC 6762, section 18.12.  Repurposing of Top Bit of qclass in Question
	// Section
	//
	// In the Question Section of a Multicast DNS query, the top bit of the qclass
	// field is used to indicate that unicast responses are preferred for this
	// particular question.  (See Section 5.4.)
	if params.WantUnicastResponse {
		m.Question[0].Qclass |= 1 << 15
	}
	m.RecursionDesired = false
	if err := c.sendQuery(m); err != nil {
		return err
	}

	// Map the in-progress responses
	inprogress := make(map[string]*ServiceEntry)

	// Listen until we reach the timeout
	finish := time.After(params.Timeout)
	for {
		select {
		case resp := <-msgCh:
			var inp *ServiceEntry
			for _, answer := range resp.Answer {
				// TODO(reddaly): Check that response corresponds to serviceAddr?
				switch rr := answer.(type) {
				case *dns.PTR:
					// Create new entry for this
					inp = ensureName(inprogress, rr.Ptr)

				case *dns.SRV:
					// Check for a target mismatch
					if rr.Target != rr.Hdr.Name {
						alias(inprogress, rr.Hdr.Name, rr.Target)
					}

					// Get the port
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Host = rr.Target
					inp.Port = int(rr.Port)

				case *dns.TXT:
					// Pull out the txt
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Info = strings.Join(rr.Txt, "|")
					inp.hasTXT = true

				case *dns.A:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.A // @Deprecated
					inp.AddrV4 = rr.A

				case *dns.AAAA:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.AAAA // @Deprecated
					inp.AddrV6 = rr.AAAA
				}
			}

			if inp == nil {
				continue
			}

			// Check if this entry is complete
			if inp.complete() {
				if inp.sent {
					continue
				}
				inp.sent = true
				select {
				case params.Entries <- inp:
				default:
				}
			} else {
				// Fire off a node specific query
				m := new(dns.Msg)
				m.SetQuestion(inp.Name, dns.TypePTR)
				m.RecursionDesired = false
				if err := c.sendQuery(m); err != nil {
					log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
				}
			}
		case <-finish:
			return nil
		}
	}
}

// sendQuery is used to multicast a query out
func (c *client) sendQuery(q *dns.Msg) error {
	buf, err := q.Pack()
	if err != nil {
		return err
	}
	if c.ipv4UnicastConn != nil {
		c.ipv4UnicastConn.WriteToUDP(buf, ipv4Addr)
	}
	if c.ipv6UnicastConn != nil {
		c.ipv6UnicastConn.WriteToUDP(buf, ipv6Addr)
	}
	return nil
}

// recv is used to receive until we get a shutdown
func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) {
	if l == nil {
		return
	}
	buf := make([]byte, 65536)
	for !c.closed {
		n, err := l.Read(buf)
		if err != nil {
			log.Printf("[ERR] mdns: Failed to read packet: %v", err)
			continue
		}
		msg := new(dns.Msg)
		if err := msg.Unpack(buf[:n]); err != nil {
			log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
			continue
		}
		select {
		case msgCh <- msg:
		case <-c.closedCh:
			return
		}
	}
}

// ensureName is used to ensure the named node is in progress
func ensureName(inprogress map[string]*ServiceEntry, name string) *ServiceEntry {
	if inp, ok := inprogress[name]; ok {
		return inp
	}
	inp := &ServiceEntry{
		Name: name,
	}
	inprogress[name] = inp
	return inp
}

// alias is used to setup an alias between two entries
func alias(inprogress map[string]*ServiceEntry, src, dst string) {
	srcEntry := ensureName(inprogress, src)
	inprogress[dst] = srcEntry
}
<|endoftext|>


<commit_before>// Package instrumented_http provides a drop-in metrics-enabled replacement for
// any http.Client or http.RoundTripper.
package instrumented_http

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

// Transport is a http.RoundTripper that collects Prometheus metrics of every
// request it processes. It allows to be configured with callbacks that process
// request path and query into a suitable label value.
type Transport struct {
	next http.RoundTripper
	cbs  *Callbacks
}

// Callbacks is a collection of callbacks passed to Transport.
type Callbacks struct {
	PathProcessor  func(string) string
	QueryProcessor func(string) string
}

const (
	// Metrics created can be identified by this label value.
	handlerName = "instrumented_http"
)

var (
	// RequestDurationMicroseconds is a Prometheus summary to collect request times.
	RequestDurationMicroseconds = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Name:        "request_duration_microseconds",
			Help:        "The HTTP request latencies in microseconds.",
			Subsystem:   "http",
			ConstLabels: prometheus.Labels{"handler": handlerName},
			Objectives:  map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		},
		[]string{"scheme", "host", "path", "query", "method", "status"},
	)

	// EliminatingProcessor is a callback that returns a blank string on any input.
	EliminatingProcessor = func(_ string) string { return "" }
	// IdentityProcessor is callback that returns whatever is passed to it.
	IdentityProcessor = func(input string) string { return input }
	// LastPathElementProcessor is callback that returns the last element of a URL path.
	LastPathElementProcessor = func(path string) string {
		parts := strings.Split(path, "/")
		return parts[len(parts)-1]
	}
)

// init registers the Prometheus metric globally when the package is loaded.
func init() {
	prometheus.MustRegister(RequestDurationMicroseconds)
}

// RoundTrip implements http.RoundTripper. It forwards the request to the
// next RoundTripper and measures the time it took in Prometheus summary.
func (it *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	// Remember the current time.
	now := time.Now()

	// Make the request using the next RoundTripper.
	resp, err := it.next.RoundTrip(req)

	// Observe the time it took to make the request.
	RequestDurationMicroseconds.WithLabelValues(
		req.URL.Scheme,
		req.URL.Host,
		it.cbs.PathProcessor(req.URL.Path),
		it.cbs.QueryProcessor(req.URL.RawQuery),
		req.Method,
		fmt.Sprintf("%d", resp.StatusCode),
	).Observe(float64(time.Since(now).Nanoseconds() / 1000))

	// return the response and error reported from the next RoundTripper.
	return resp, err
}

// NewClient takes a *http.Client and returns a *http.Client that has its
// RoundTripper wrapped with instrumentation. Optionally, It can receive a
// collection of callbacks that process request path and query into a suitable
// label value.
func NewClient(next *http.Client, cbs *Callbacks) *http.Client {
	// If next client is not defined we'll use http.DefaultClient.
	if next == nil {
		next = http.DefaultClient
	}

	next.Transport = NewTransport(next.Transport, cbs)

	return next
}

// NewTransport takes a http.RoundTripper, wraps it with instrumentation and
// returns it as a new http.RoundTripper. Optionally, It can receive a
// collection of callbacks that process request path and query into a suitable
// label value.
func NewTransport(next http.RoundTripper, cbs *Callbacks) http.RoundTripper {
	// If next RoundTripper is not defined we'll use http.DefaultTransport.
	if next == nil {
		next = http.DefaultTransport
	}

	// If cbs is not defined we'll initilialize it with defaults.
	if cbs == nil {
		cbs = &Callbacks{}
	}
	// By default, path and query will be ignored.
	if cbs.PathProcessor == nil {
		cbs.PathProcessor = EliminatingProcessor
	}
	if cbs.QueryProcessor == nil {
		cbs.QueryProcessor = EliminatingProcessor
	}

	return &Transport{next: next, cbs: cbs}
}
<commit_msg>fix: don't panic when there's no response<commit_after>// Package instrumented_http provides a drop-in metrics-enabled replacement for
// any http.Client or http.RoundTripper.
package instrumented_http

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

// Transport is a http.RoundTripper that collects Prometheus metrics of every
// request it processes. It allows to be configured with callbacks that process
// request path and query into a suitable label value.
type Transport struct {
	next http.RoundTripper
	cbs  *Callbacks
}

// Callbacks is a collection of callbacks passed to Transport.
type Callbacks struct {
	PathProcessor  func(string) string
	QueryProcessor func(string) string
}

const (
	// Metrics created can be identified by this label value.
	handlerName = "instrumented_http"
)

var (
	// RequestDurationMicroseconds is a Prometheus summary to collect request times.
	RequestDurationMicroseconds = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Name:        "request_duration_microseconds",
			Help:        "The HTTP request latencies in microseconds.",
			Subsystem:   "http",
			ConstLabels: prometheus.Labels{"handler": handlerName},
			Objectives:  map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		},
		[]string{"scheme", "host", "path", "query", "method", "status"},
	)

	// EliminatingProcessor is a callback that returns a blank string on any input.
	EliminatingProcessor = func(_ string) string { return "" }
	// IdentityProcessor is callback that returns whatever is passed to it.
	IdentityProcessor = func(input string) string { return input }
	// LastPathElementProcessor is callback that returns the last element of a URL path.
	LastPathElementProcessor = func(path string) string {
		parts := strings.Split(path, "/")
		return parts[len(parts)-1]
	}
)

// init registers the Prometheus metric globally when the package is loaded.
func init() {
	prometheus.MustRegister(RequestDurationMicroseconds)
}

// RoundTrip implements http.RoundTripper. It forwards the request to the
// next RoundTripper and measures the time it took in Prometheus summary.
func (it *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	var statusCode int

	// Remember the current time.
	now := time.Now()

	// Make the request using the next RoundTripper.
	resp, err := it.next.RoundTrip(req)
	if resp != nil {
		statusCode = resp.StatusCode
	}

	// Observe the time it took to make the request.
	RequestDurationMicroseconds.WithLabelValues(
		req.URL.Scheme,
		req.URL.Host,
		it.cbs.PathProcessor(req.URL.Path),
		it.cbs.QueryProcessor(req.URL.RawQuery),
		req.Method,
		fmt.Sprintf("%d", statusCode),
	).Observe(float64(time.Since(now).Nanoseconds() / 1000))

	// return the response and error reported from the next RoundTripper.
	return resp, err
}

// NewClient takes a *http.Client and returns a *http.Client that has its
// RoundTripper wrapped with instrumentation. Optionally, It can receive a
// collection of callbacks that process request path and query into a suitable
// label value.
func NewClient(next *http.Client, cbs *Callbacks) *http.Client {
	// If next client is not defined we'll use http.DefaultClient.
	if next == nil {
		next = http.DefaultClient
	}

	next.Transport = NewTransport(next.Transport, cbs)

	return next
}

// NewTransport takes a http.RoundTripper, wraps it with instrumentation and
// returns it as a new http.RoundTripper. Optionally, It can receive a
// collection of callbacks that process request path and query into a suitable
// label value.
func NewTransport(next http.RoundTripper, cbs *Callbacks) http.RoundTripper {
	// If next RoundTripper is not defined we'll use http.DefaultTransport.
	if next == nil {
		next = http.DefaultTransport
	}

	// If cbs is not defined we'll initilialize it with defaults.
	if cbs == nil {
		cbs = &Callbacks{}
	}
	// By default, path and query will be ignored.
	if cbs.PathProcessor == nil {
		cbs.PathProcessor = EliminatingProcessor
	}
	if cbs.QueryProcessor == nil {
		cbs.QueryProcessor = EliminatingProcessor
	}

	return &Transport{next: next, cbs: cbs}
}
<|endoftext|>


<commit_before>package trousseau

import (
	"github.com/codegangsta/cli"
)

func CreateCommand() cli.Command {
	return cli.Command{
		Name:   "create",
		Usage:  "create the trousseau data store",
		Action: CreateAction,
	}
}

func PushCommand() cli.Command {
	return cli.Command{
		Name:   "push",
		Usage:  "pushes the trousseau to remote storage",
		Action: PushAction,
		Flags: []cli.Flag{
			OverwriteFlag(),
			AskPassword(),
			SshPrivateKeyPathFlag(),
		},
	}
}

func PullCommand() cli.Command {
	return cli.Command{
		Name:   "pull",
		Usage:  "pull the trousseau from remote storage",
		Action: PullAction,
		Flags: []cli.Flag{
			OverwriteFlag(),
			AskPassword(),
			SshPrivateKeyPathFlag(),
		},
	}
}

func ExportCommand() cli.Command {
	return cli.Command{
		Name:   "export",
		Usage:  "export the encrypted trousseau to local fs",
		Action: ExportAction,
		Flags: []cli.Flag{
			OverwriteFlag(),
		},
	}
}

func ImportCommand() cli.Command {
	return cli.Command{
		Name:   "import",
		Usage:  "import an encrypted trousseau from local fs",
		Action: ImportAction,
		Flags: []cli.Flag{
			OverwriteFlag(),
			TheirsFlag(),
			YoursFlag(),
		},
	}
}

func AddRecipientCommand() cli.Command {
	return cli.Command{
		Name:   "add-recipient",
		Usage:  "add a recipient to the encrypted trousseau",
		Action: AddRecipientAction,
	}
}

func RemoveRecipientCommand() cli.Command {
	return cli.Command{
		Name:   "remove-recipient",
		Usage:  "remove a recipient of the encrypted trousseau",
		Action: RemoveRecipientAction,
	}
}

func SetCommand() cli.Command {
	return cli.Command{
		Name:   "set",
		Usage:  "sets a key value pair in the store",
		Action: SetAction,
		Flags: []cli.Flag{
			FileFlag(),
		},
	}
}

func GetCommand() cli.Command {
	return cli.Command{
		Name:   "get",
		Usage:  "get a value from the trousseau",
		Action: GetAction,
	}
}

func DelCommand() cli.Command {
	return cli.Command{
		Name:   "del",
		Usage:  "delete the point key pair from the store",
		Action: DelAction,
	}
}

func KeysCommand() cli.Command {
	return cli.Command{
		Name:   "keys",
		Usage:  "Lists the store keys",
		Action: KeysAction,
	}
}

func ShowCommand() cli.Command {
	return cli.Command{
		Name:   "show",
		Usage:  "shows trousseau content",
		Action: ShowAction,
	}
}

func MetaCommand() cli.Command {
	return cli.Command{
		Name:   "meta",
		Usage:  "shows trousseau metadata",
		Action: MetaAction,
	}
}
<commit_msg>Update commands with verbose flag<commit_after>package trousseau

import (
	"github.com/codegangsta/cli"
)

func CreateCommand() cli.Command {
	return cli.Command{
		Name:   "create",
		Usage:  "create the trousseau data store",
		Action: CreateAction,
	}
}

func PushCommand() cli.Command {
	return cli.Command{
		Name:   "push",
		Usage:  "pushes the trousseau to remote storage",
		Action: PushAction,
		Flags: []cli.Flag{
			OverwriteFlag(),
			AskPassword(),
			VerboseFlag(),
			SshPrivateKeyPathFlag(),
		},
	}
}

func PullCommand() cli.Command {
	return cli.Command{
		Name:   "pull",
		Usage:  "pull the trousseau from remote storage",
		Action: PullAction,
		Flags: []cli.Flag{
			OverwriteFlag(),
			AskPassword(),
			VerboseFlag(),
			SshPrivateKeyPathFlag(),
		},
	}
}

func ExportCommand() cli.Command {
	return cli.Command{
		Name:   "export",
		Usage:  "export the encrypted trousseau to local fs",
		Action: ExportAction,
		Flags: []cli.Flag{
			OverwriteFlag(),
			VerboseFlag(),
		},
	}
}

func ImportCommand() cli.Command {
	return cli.Command{
		Name:   "import",
		Usage:  "import an encrypted trousseau from local fs",
		Action: ImportAction,
		Flags: []cli.Flag{
			VerboseFlag(),
			OverwriteFlag(),
			TheirsFlag(),
			YoursFlag(),
		},
	}
}

func AddRecipientCommand() cli.Command {
	return cli.Command{
		Name:   "add-recipient",
		Usage:  "add a recipient to the encrypted trousseau",
		Action: AddRecipientAction,
		Flags: []cli.Flag{
			VerboseFlag(),
		},
	}
}

func RemoveRecipientCommand() cli.Command {
	return cli.Command{
		Name:   "remove-recipient",
		Usage:  "remove a recipient of the encrypted trousseau",
		Action: RemoveRecipientAction,
		Flags: []cli.Flag{
			VerboseFlag(),
		},
	}
}

func SetCommand() cli.Command {
	return cli.Command{
		Name:   "set",
		Usage:  "sets a key value pair in the store",
		Action: SetAction,
		Flags: []cli.Flag{
			FileFlag(),
			VerboseFlag(),
		},
	}
}

func GetCommand() cli.Command {
	return cli.Command{
		Name:   "get",
		Usage:  "get a value from the trousseau",
		Action: GetAction,
	}
}

func DelCommand() cli.Command {
	return cli.Command{
		Name:   "del",
		Usage:  "delete the point key pair from the store",
		Action: DelAction,
		Flags: []cli.Flag{
			VerboseFlag(),
		},
	}
}

func KeysCommand() cli.Command {
	return cli.Command{
		Name:   "keys",
		Usage:  "Lists the store keys",
		Action: KeysAction,
		Flags: []cli.Flag{
			VerboseFlag(),
		},
	}
}

func ShowCommand() cli.Command {
	return cli.Command{
		Name:   "show",
		Usage:  "shows trousseau content",
		Action: ShowAction,
	}
}

func MetaCommand() cli.Command {
	return cli.Command{
		Name:   "meta",
		Usage:  "shows trousseau metadata",
		Action: MetaAction,
	}
}
<|endoftext|>


<commit_before>package winrmtest

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/antchfx/xquery/xml"
	"github.com/satori/go.uuid"
)

func Test_creating_a_shell(t *testing.T) {
	w := &wsman{}

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/transfer/Create</a:Action>
			</env:Header>
			<env:Body>
				<rsp:Shell>
					<rsp:InputStream>stdin</rsp:InputStream>
					<rsp:OutputStreams>stdout stderr</rsp:OutputStreams>
				</rsp:Shell>
			</env:Body>
		</env:Envelope>`))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Errorf("Expected 200 OK but was %d.\n", res.Code)
	}

	if contentType := res.HeaderMap.Get("Content-Type"); contentType != "application/soap+xml" {
		t.Errorf("Expected ContentType application/soap+xml was %s.\n", contentType)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't parse XML: %s", err)
	}
	result := xmlquery.FindOne(doc, "//rsp:ShellId").InnerText()
	if result == "" {
		t.Error("Expected a Shell identifier.")
	}
}

func Test_executing_a_command(t *testing.T) {
	w := &wsman{}
	id := w.HandleCommand(MatchText("echo tacos"), func(out, err io.Writer) int {
		return 0
	})

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing" xmlns:rsp="http://schemas.microsoft.com/wbem/wsman/1/windows/shell">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/shell/Command</a:Action>
			</env:Header>
			<env:Body>
				<rsp:CommandLine><rsp:Command>"echo tacos"</rsp:Command></rsp:CommandLine>
			</env:Body>
		</env:Envelope>`))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't compile the SOAP response: %s", err)
	}
	result := xmlquery.FindOne(doc, "//rsp:CommandId").InnerText()

	if result != id {
		t.Errorf("Expected CommandId=%s but was %q", id, result)
	}
}

func Test_executing_a_regex_command(t *testing.T) {
	w := &wsman{}
	id := w.HandleCommand(MatchPattern(`echo .* >> C:\file.cmd`), func(out, err io.Writer) int {
		return 0
	})

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(fmt.Sprintf(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing" xmlns:rsp="http://schemas.microsoft.com/wbem/wsman/1/windows/shell">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/shell/Command</a:Action>
			</env:Header>
			<env:Body>
				<rsp:CommandLine><rsp:Command>"echo %d >> C:\file.cmd"</rsp:Command></rsp:CommandLine>
			</env:Body>
		</env:Envelope>`, uuid.NewV4().String())))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't compile the SOAP response.")
	}
	result := xmlquery.FindOne(doc, "//rsp:CommandId").InnerText()

	if result != id {
		t.Errorf("Expected CommandId=%s, but was %q", id, result)
	}
}

func Test_receiving_command_results(t *testing.T) {
	w := &wsman{}
	id := w.HandleCommand(MatchText("echo tacos"), func(out, err io.Writer) int {
		out.Write([]byte("tacos"))
		return 0
	})

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(fmt.Sprintf(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing" xmlns:rsp="http://schemas.microsoft.com/wbem/wsman/1/windows/shell">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive</a:Action>
			</env:Header>
			<env:Body>
				<rsp:Receive><rsp:DesiredStream CommandId="%s">stdout stderr</rsp:DesiredStream></rsp:Receive>
			</env:Body>
		</env:Envelope>`, id)))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't compile the SOAP response: %s", err)
	}

	result := xmlquery.FindOne(doc, "//rsp:ReceiveResponse")
	iter := xmlquery.CreateXPathNavigator(result)
	if !iter.MoveToChild() {
		t.Error("Expected a ReceiveResponse element.")
	}

	xresp := iter.Current()
	result = xmlquery.FindOne(xresp, fmt.Sprintf("rsp:Stream[@CommandId='%s']", id))
	iter = xmlquery.CreateXPathNavigator(xresp)

	testText := "dGFjb3M="
	if !iter.MoveToNext() ||
		iter.Current().SelectAttr("Name") != "stdout" ||
		iter.Current().InnerText() != testText {
		t.Errorf("Expected an stdout Stream with the text %q", testText)
	}

	if !iter.MoveToNext() ||
		iter.Current().SelectAttr("Name") != "stdout" ||
		iter.Current().SelectAttr("End") != "true" {
		t.Errorf("Expected an stdout Stream with an %q attribute", "end")
	}

	if !iter.MoveToNext() ||
		iter.Current().SelectAttr("Name") != "stderr" ||
		iter.Current().SelectAttr("End") != "true" {
		t.Errorf("Expected an stderr Stream with an %q attribute", "end")
	}

	result = xmlquery.FindOne(doc, "//rsp:CommandState")
	if result == nil {
		t.Errorf("Expected CommandState=%q, got: nil", "Done")
	}

	if result.SelectAttr("State") != "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/CommandState/Done" {
		t.Errorf("Expected CommandState=%q", "Done")
	}

	result = xmlquery.FindOne(doc,
		"//rsp:CommandState/rsp:ExitCode")
	if result.InnerText() != "0" {
		t.Errorf("Expected ExitCode=0 but found %q", result.InnerText())
	}
}

func Test_deleting_a_shell(t *testing.T) {
	w := &wsman{}

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/transfer/Delete</a:Action>
			</env:Header>
		</env:Envelope>`))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	if res.Body.Len() != 0 {
		t.Errorf("Expected body to be empty but was \"%v\".", res.Body)
	}
}
<commit_msg>Update fmt/test<commit_after>package winrmtest

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	xmlquery "github.com/antchfx/xquery/xml"
	uuid "github.com/satori/go.uuid"
)

func Test_creating_a_shell(t *testing.T) {
	w := &wsman{}

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/transfer/Create</a:Action>
			</env:Header>
			<env:Body>
				<rsp:Shell>
					<rsp:InputStream>stdin</rsp:InputStream>
					<rsp:OutputStreams>stdout stderr</rsp:OutputStreams>
				</rsp:Shell>
			</env:Body>
		</env:Envelope>`))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Errorf("Expected 200 OK but was %d.\n", res.Code)
	}

	if contentType := res.HeaderMap.Get("Content-Type"); contentType != "application/soap+xml" {
		t.Errorf("Expected ContentType application/soap+xml was %s.\n", contentType)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't parse XML: %s", err)
	}
	result := xmlquery.FindOne(doc, "//rsp:ShellId").InnerText()
	if result == "" {
		t.Error("Expected a Shell identifier.")
	}
}

func Test_executing_a_command(t *testing.T) {
	w := &wsman{}
	id := w.HandleCommand(MatchText("echo tacos"), func(out, err io.Writer) int {
		return 0
	})

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing" xmlns:rsp="http://schemas.microsoft.com/wbem/wsman/1/windows/shell">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/shell/Command</a:Action>
			</env:Header>
			<env:Body>
				<rsp:CommandLine><rsp:Command>"echo tacos"</rsp:Command></rsp:CommandLine>
			</env:Body>
		</env:Envelope>`))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't compile the SOAP response: %s", err)
	}
	result := xmlquery.FindOne(doc, "//rsp:CommandId").InnerText()

	if result != id {
		t.Errorf("Expected CommandId=%s but was %q", id, result)
	}
}

func Test_executing_a_regex_command(t *testing.T) {
	w := &wsman{}
	id := w.HandleCommand(MatchPattern(`echo .* >> C:\file.cmd`), func(out, err io.Writer) int {
		return 0
	})

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(fmt.Sprintf(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing" xmlns:rsp="http://schemas.microsoft.com/wbem/wsman/1/windows/shell">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/shell/Command</a:Action>
			</env:Header>
			<env:Body>
				<rsp:CommandLine><rsp:Command>"echo %s >> C:\file.cmd"</rsp:Command></rsp:CommandLine>
			</env:Body>
		</env:Envelope>`, uuid.NewV4().String())))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't compile the SOAP response.")
	}
	result := xmlquery.FindOne(doc, "//rsp:CommandId").InnerText()

	if result != id {
		t.Errorf("Expected CommandId=%s, but was %q", id, result)
	}
}

func Test_receiving_command_results(t *testing.T) {
	w := &wsman{}
	id := w.HandleCommand(MatchText("echo tacos"), func(out, err io.Writer) int {
		out.Write([]byte("tacos"))
		return 0
	})

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(fmt.Sprintf(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing" xmlns:rsp="http://schemas.microsoft.com/wbem/wsman/1/windows/shell">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive</a:Action>
			</env:Header>
			<env:Body>
				<rsp:Receive><rsp:DesiredStream CommandId="%s">stdout stderr</rsp:DesiredStream></rsp:Receive>
			</env:Body>
		</env:Envelope>`, id)))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	doc, err := xmlquery.Parse(res.Body)
	if err != nil {
		t.Errorf("Couldn't compile the SOAP response: %s", err)
	}

	result := xmlquery.FindOne(doc, "//rsp:ReceiveResponse")
	iter := xmlquery.CreateXPathNavigator(result)
	if !iter.MoveToChild() {
		t.Error("Expected a ReceiveResponse element.")
	}

	xresp := iter.Current()
	result = xmlquery.FindOne(xresp, fmt.Sprintf("rsp:Stream[@CommandId='%s']", id))
	iter = xmlquery.CreateXPathNavigator(xresp)

	testText := "dGFjb3M="
	if !iter.MoveToNext() ||
		iter.Current().SelectAttr("Name") != "stdout" ||
		iter.Current().InnerText() != testText {
		t.Errorf("Expected an stdout Stream with the text %q", testText)
	}

	if !iter.MoveToNext() ||
		iter.Current().SelectAttr("Name") != "stdout" ||
		iter.Current().SelectAttr("End") != "true" {
		t.Errorf("Expected an stdout Stream with an %q attribute", "end")
	}

	if !iter.MoveToNext() ||
		iter.Current().SelectAttr("Name") != "stderr" ||
		iter.Current().SelectAttr("End") != "true" {
		t.Errorf("Expected an stderr Stream with an %q attribute", "end")
	}

	result = xmlquery.FindOne(doc, "//rsp:CommandState")
	if result == nil {
		t.Errorf("Expected CommandState=%q, got: nil", "Done")
	}

	if result.SelectAttr("State") != "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/CommandState/Done" {
		t.Errorf("Expected CommandState=%q", "Done")
	}

	result = xmlquery.FindOne(doc,
		"//rsp:CommandState/rsp:ExitCode")
	if result.InnerText() != "0" {
		t.Errorf("Expected ExitCode=0 but found %q", result.InnerText())
	}
}

func Test_deleting_a_shell(t *testing.T) {
	w := &wsman{}

	res := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "", strings.NewReader(`
		<env:Envelope xmlns:env="http://www.w3.org/2003/05/soap-envelope" xmlns:a="http://schemas.xmlsoap.org/ws/2004/08/addressing">
			<env:Header>
				<a:Action mustUnderstand="true">http://schemas.xmlsoap.org/ws/2004/09/transfer/Delete</a:Action>
			</env:Header>
		</env:Envelope>`))

	w.ServeHTTP(res, req)
	if res.Code != http.StatusOK {
		t.Fatalf("Expected 200 OK but was %d.\n", res.Code)
	}

	if res.Body.Len() != 0 {
		t.Errorf("Expected body to be empty but was \"%v\".", res.Body)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ 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.

package service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
//"github.com/surgemq/message"
	"encoding/binary"
	"math"
	"sync"
)

var (
	bufcnt int64
	DefaultBufferSize int64

	DeviceInBufferSize int64
	DeviceOutBufferSize int64

	MasterInBufferSize int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize = 512
	defaultReadBlockSize = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
 */
type buffer struct {
	readIndex  int64        //读序号
	writeIndex int64        //写序号
	ringBuffer []*ByteArray //环形buffer指针数组
	bufferSize int64        //初始化环形buffer指针数组大小
	mask       int64        //掩码：bufferSize-1
	done       int64        //是否完成
	rcond      *sync.Cond
	wcond      *sync.Cond
}

type ByteArray struct {
	bArray []byte
}

func (this *buffer)ReadCommit(index int64) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()
	this.ringBuffer[index] = nil
	this.rcond.Broadcast()
	this.wcond.Broadcast()
}

func (this *buffer) Len() int {
	cpos := this.GetCurrentReadIndex()
	ppos := this.GetCurrentWriteIndex()
	return int(ppos - cpos)
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
 */
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex: int64(0), //读序号
		writeIndex: int64(0), //写序号
		ringBuffer: make([]*ByteArray, size), //环形buffer指针数组
		bufferSize: size, //初始化环形buffer指针数组大小
		mask:size - 1,
		rcond:         sync.NewCond(new(sync.Mutex)),
		wcond:         sync.NewCond(new(sync.Mutex)),
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
 */
func (this *buffer)GetCurrentReadIndex() (int64) {
	return atomic.LoadInt64(&this.readIndex)
}
/**
2016.03.03 添加
获取当前写序号
 */
func (this *buffer)GetCurrentWriteIndex() (int64) {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
 */
func (this *buffer)ReadBuffer() ([]byte, int64, bool) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch  {
		case readIndex >= writeIndex:
			this.rcond.Wait()
		case writeIndex - readIndex > this.bufferSize:
			this.rcond.Wait()
		default:
			if readIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.readIndex, int64(0))
			}else {
				atomic.AddInt64(&this.readIndex, int64(1))
			}

			index := readIndex & this.mask

			p_ := this.ringBuffer[index]
			//this.ringBuffer[index] = nil
			if p_ == nil {
				return nil, -1, false
			}
			p := p_.bArray

			if p == nil {
				return nil, -1, false
			}

			return p, index, true
		}
	}

}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
 */
func (this *buffer)WriteBuffer(in []byte) (bool) {
	this.wcond.L.Lock()
	defer this.wcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch  {
		case writeIndex - readIndex < 0:
			this.wcond.Wait()
		default:
			index := writeIndex & this.mask
			if writeIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.writeIndex, int64(0))
			}else {
				atomic.AddInt64(&this.writeIndex, int64(1))
			}
			if this.ringBuffer[index] == nil {
				this.ringBuffer[index] = &ByteArray{bArray:in}
				this.rcond.Broadcast()
				return true
			}else {
				this.rcond.Broadcast()
				return false
			}
		}
	}

}

/**
2016.03.03 修改
完成
 */
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	atomic.StoreInt64(&this.done, 1)

	this.wcond.L.Lock()
	this.rcond.Broadcast()
	this.wcond.L.Unlock()

	this.rcond.L.Lock()
	this.wcond.Broadcast()
	this.rcond.L.Unlock()

	return nil
}
/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 5)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			Log.Infoc(func() string {
				return fmt.Sprintf("sendrecv/peekMessageSize: %d=========", cnt)
			})
			// Peek cnt bytes from the input buffer.
			_, err := r.Read(b[cnt:cnt + 1])
			if err != nil {
				return 0, err
			}

			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, m := binary.Uvarint(b[1:])
		Log.Infoc(func() string {
			return fmt.Sprintf("remlen: %d===%d", remlen,n)
		})
		// Total message length is remlen + 1 (msg type) + m (remlen bytes)

		remlen_ := int64(remlen)
		total = remlen_ + int64(1) + int64(m)
		Log.Infoc(func() string {
			return fmt.Sprintf("remlen===n===totle: %d===%d===%d", remlen,n,total)
		})
		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		if this.WriteBuffer(b) {
			return total, err
		}

		return total, nil
	}
}

/**
2016.03.03 修改
 */
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, index, ok := this.ReadBuffer()
		defer this.ReadCommit(index)
		if !ok {
			return total, io.EOF
		}

		Log.Debugc(func() string {
			return fmt.Sprintf("defer this.ReadCommit(%s)", index)
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(p) > 0 {
			n, err := w.Write(p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n & (n - 1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<commit_msg>添加测试代码<commit_after>// Copyright (c) 2014 The SurgeMQ 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.

package service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
//"github.com/surgemq/message"
	"encoding/binary"
	"math"
	"sync"
)

var (
	bufcnt int64
	DefaultBufferSize int64

	DeviceInBufferSize int64
	DeviceOutBufferSize int64

	MasterInBufferSize int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize = 512
	defaultReadBlockSize = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
 */
type buffer struct {
	readIndex  int64        //读序号
	writeIndex int64        //写序号
	ringBuffer []*ByteArray //环形buffer指针数组
	bufferSize int64        //初始化环形buffer指针数组大小
	mask       int64        //掩码：bufferSize-1
	done       int64        //是否完成
	rcond      *sync.Cond
	wcond      *sync.Cond
}

type ByteArray struct {
	bArray []byte
}

func (this *buffer)ReadCommit(index int64) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()
	this.ringBuffer[index] = nil
	this.rcond.Broadcast()
	this.wcond.Broadcast()
}

func (this *buffer) Len() int {
	cpos := this.GetCurrentReadIndex()
	ppos := this.GetCurrentWriteIndex()
	return int(ppos - cpos)
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
 */
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex: int64(0), //读序号
		writeIndex: int64(0), //写序号
		ringBuffer: make([]*ByteArray, size), //环形buffer指针数组
		bufferSize: size, //初始化环形buffer指针数组大小
		mask:size - 1,
		rcond:         sync.NewCond(new(sync.Mutex)),
		wcond:         sync.NewCond(new(sync.Mutex)),
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
 */
func (this *buffer)GetCurrentReadIndex() (int64) {
	return atomic.LoadInt64(&this.readIndex)
}
/**
2016.03.03 添加
获取当前写序号
 */
func (this *buffer)GetCurrentWriteIndex() (int64) {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
 */
func (this *buffer)ReadBuffer() ([]byte, int64, bool) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch  {
		case readIndex >= writeIndex:
			this.rcond.Wait()
		case writeIndex - readIndex > this.bufferSize:
			this.rcond.Wait()
		default:
			if readIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.readIndex, int64(0))
			}else {
				atomic.AddInt64(&this.readIndex, int64(1))
			}

			index := readIndex & this.mask

			p_ := this.ringBuffer[index]
			//this.ringBuffer[index] = nil
			if p_ == nil {
				return nil, -1, false
			}
			p := p_.bArray

			if p == nil {
				return nil, -1, false
			}

			return p, index, true
		}
	}

}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
 */
func (this *buffer)WriteBuffer(in []byte) (bool) {
	this.wcond.L.Lock()
	defer this.wcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch  {
		case writeIndex - readIndex < 0:
			this.wcond.Wait()
		default:
			index := writeIndex & this.mask
			if writeIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.writeIndex, int64(0))
			}else {
				atomic.AddInt64(&this.writeIndex, int64(1))
			}
			if this.ringBuffer[index] == nil {
				this.ringBuffer[index] = &ByteArray{bArray:in}
				this.rcond.Broadcast()
				return true
			}else {
				this.rcond.Broadcast()
				return false
			}
		}
	}

}

/**
2016.03.03 修改
完成
 */
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	atomic.StoreInt64(&this.done, 1)

	this.wcond.L.Lock()
	this.rcond.Broadcast()
	this.wcond.L.Unlock()

	this.rcond.L.Lock()
	this.wcond.Broadcast()
	this.rcond.L.Unlock()

	return nil
}
/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 5)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			Log.Infoc(func() string {
				return fmt.Sprintf("sendrecv/peekMessageSize: %d=========", cnt)
			})
			// Peek cnt bytes from the input buffer.
			_, err := r.Read(b[cnt:cnt + 1])
			if err != nil {
				return 0, err
			}

			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, m := binary.Uvarint(b[1:])
		Log.Infoc(func() string {
			return fmt.Sprintf("remlen: %d===%d", remlen,m)
		})
		// Total message length is remlen + 1 (msg type) + m (remlen bytes)

		remlen_ := int64(remlen)
		total = remlen_ + int64(1) + int64(m)
		Log.Infoc(func() string {
			return fmt.Sprintf("remlen===n===totle: %d===%d===%d", remlen,m,total)
		})
		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		if this.WriteBuffer(b) {
			return total, err
		}

		return total, nil
	}
}

/**
2016.03.03 修改
 */
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, index, ok := this.ReadBuffer()
		defer this.ReadCommit(index)
		if !ok {
			return total, io.EOF
		}

		Log.Debugc(func() string {
			return fmt.Sprintf("defer this.ReadCommit(%s)", index)
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(p) > 0 {
			n, err := w.Write(p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n & (n - 1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ 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.

package service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
	//"github.com/surgemq/message"
	"encoding/binary"
	"math"
	"sync"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize    = 512
	defaultReadBlockSize  = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
*/
type buffer struct {
	readIndex  int64        //读序号
	writeIndex int64        //写序号
	ringBuffer []*ByteArray //环形buffer指针数组
	bufferSize int64        //初始化环形buffer指针数组大小
	mask       int64        //掩码：bufferSize-1
	done       int64        //是否完成
	rcond      *sync.Cond
	wcond      *sync.Cond
}

type ByteArray struct {
	bArray []byte
}

func (this *buffer) ReadCommit(index int64) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()
	this.ringBuffer[index] = nil
	this.rcond.Broadcast()
	this.wcond.Broadcast()
}

func (this *buffer) Len() int {
	cpos := this.GetCurrentReadIndex()
	ppos := this.GetCurrentWriteIndex()
	return int(ppos - cpos)
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
*/
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex:  int64(0),                 //读序号
		writeIndex: int64(0),                 //写序号
		ringBuffer: make([]*ByteArray, size), //环形buffer指针数组
		bufferSize: size,                     //初始化环形buffer指针数组大小
		mask:       size - 1,
		rcond:      sync.NewCond(new(sync.Mutex)),
		wcond:      sync.NewCond(new(sync.Mutex)),
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
2016.03.03 添加
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() ([]byte, int64, bool) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch {
		case readIndex >= writeIndex:
			this.rcond.Wait()
		case writeIndex-readIndex > this.bufferSize:
			this.rcond.Wait()
		default:
			if readIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.readIndex, int64(0))
			} else {
				atomic.AddInt64(&this.readIndex, int64(1))
			}

			index := readIndex & this.mask

			p_ := this.ringBuffer[index]
			//this.ringBuffer[index] = nil
			if p_ == nil {
				return nil, -1, false
			}
			p := p_.bArray

			if p == nil {
				return nil, -1, false
			}

			return p, index, true
		}
	}

}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in []byte) bool {
	this.wcond.L.Lock()
	defer this.wcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch {
		case writeIndex-readIndex < 0:
			this.wcond.Wait()
		default:
			index := writeIndex & this.mask
			if writeIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.writeIndex, int64(0))
			} else {
				atomic.AddInt64(&this.writeIndex, int64(1))
			}
			if this.ringBuffer[index] == nil {
				this.ringBuffer[index] = &ByteArray{bArray: in}
				this.rcond.Broadcast()
				return true
			} else {
				this.rcond.Broadcast()
				return false
			}
		}
	}

}

/**
2016.03.03 修改
完成
*/
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	atomic.StoreInt64(&this.done, 1)

	this.wcond.L.Lock()
	this.rcond.Broadcast()
	this.wcond.L.Unlock()

	this.rcond.L.Lock()
	this.wcond.Broadcast()
	this.rcond.L.Unlock()

	return nil
}

/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 0, 6)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			Log.Infoc(func() string {
				return fmt.Sprintf("sendrecv/peekMessageSize: %d=========", cnt)
			})
			// Peek cnt bytes from the input buffer.

			_, err := r.Read(b[cnt:(cnt + 1)])
			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("_, err := r.Read(b[cnt:(cnt + 1)]) is error")
				})
				return 0, err
			}
			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, m := binary.Uvarint(b[1 : cnt+1])
		Log.Infoc(func() string {
			return fmt.Sprintf("b[cnt:(cnt + 1)]==end")
		})
		// Total message length is remlen + 1 (msg type) + m (remlen bytes)

		remlen_ := int64(remlen)
		total = remlen_ + int64(1) + int64(m)
		Log.Infoc(func() string {
			return fmt.Sprintf("remlen===n===totle: %d===%d===%d", remlen, m, total)
		})
		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, 0, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		Log.Infoc(func() string {
			return fmt.Sprintf("len(b):%d", len(b))
		})
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		if this.WriteBuffer(b) {
			return total, err
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, index, ok := this.ReadBuffer()
		defer this.ReadCommit(index)
		if !ok {
			return total, io.EOF
		}

		Log.Debugc(func() string {
			return fmt.Sprintf("defer this.ReadCommit(%s)", index)
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(p) > 0 {
			n, err := w.Write(p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<commit_msg>添加测试代码<commit_after>// Copyright (c) 2014 The SurgeMQ 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.

package service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
	//"github.com/surgemq/message"
	"encoding/binary"
	"math"
	"sync"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize    = 512
	defaultReadBlockSize  = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
*/
type buffer struct {
	readIndex  int64        //读序号
	writeIndex int64        //写序号
	ringBuffer []*ByteArray //环形buffer指针数组
	bufferSize int64        //初始化环形buffer指针数组大小
	mask       int64        //掩码：bufferSize-1
	done       int64        //是否完成
	rcond      *sync.Cond
	wcond      *sync.Cond
}

type ByteArray struct {
	bArray []byte
}

func (this *buffer) ReadCommit(index int64) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()
	this.ringBuffer[index] = nil
	this.rcond.Broadcast()
	this.wcond.Broadcast()
}

func (this *buffer) Len() int {
	cpos := this.GetCurrentReadIndex()
	ppos := this.GetCurrentWriteIndex()
	return int(ppos - cpos)
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
*/
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex:  int64(0),                 //读序号
		writeIndex: int64(0),                 //写序号
		ringBuffer: make([]*ByteArray, size), //环形buffer指针数组
		bufferSize: size,                     //初始化环形buffer指针数组大小
		mask:       size - 1,
		rcond:      sync.NewCond(new(sync.Mutex)),
		wcond:      sync.NewCond(new(sync.Mutex)),
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
2016.03.03 添加
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() ([]byte, int64, bool) {
	this.rcond.L.Lock()
	defer this.rcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch {
		case readIndex >= writeIndex:
			this.rcond.Wait()
		case writeIndex-readIndex > this.bufferSize:
			this.rcond.Wait()
		default:
			if readIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.readIndex, int64(0))
			} else {
				atomic.AddInt64(&this.readIndex, int64(1))
			}

			index := readIndex & this.mask

			p_ := this.ringBuffer[index]
			//this.ringBuffer[index] = nil
			if p_ == nil {
				return nil, -1, false
			}
			p := p_.bArray

			if p == nil {
				return nil, -1, false
			}

			return p, index, true
		}
	}

}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in []byte) bool {
	this.wcond.L.Lock()
	defer this.wcond.L.Unlock()

	for {
		readIndex := atomic.LoadInt64(&this.readIndex)
		writeIndex := atomic.LoadInt64(&this.writeIndex)
		switch {
		case writeIndex-readIndex < 0:
			this.wcond.Wait()
		default:
			index := writeIndex & this.mask
			if writeIndex == math.MaxInt64 {
				atomic.StoreInt64(&this.writeIndex, int64(0))
			} else {
				atomic.AddInt64(&this.writeIndex, int64(1))
			}
			if this.ringBuffer[index] == nil {
				this.ringBuffer[index] = &ByteArray{bArray: in}
				this.rcond.Broadcast()
				return true
			} else {
				this.rcond.Broadcast()
				return false
			}
		}
	}

}

/**
2016.03.03 修改
完成
*/
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	atomic.StoreInt64(&this.done, 1)

	this.wcond.L.Lock()
	this.rcond.Broadcast()
	this.wcond.L.Unlock()

	this.rcond.L.Lock()
	this.wcond.Broadcast()
	this.rcond.L.Unlock()

	return nil
}

/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 0, 6)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			Log.Infoc(func() string {
				return fmt.Sprintf("sendrecv/peekMessageSize: %d=========", cnt)
			})
			// Peek cnt bytes from the input buffer.

			_, err := r.Read(b[cnt:(cnt + 1)])
			fmt.Println(b)
			if err != nil {
				return 0, err
			}
			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, m := binary.Uvarint(b[1 : cnt+1])
		Log.Infoc(func() string {
			return fmt.Sprintf("b[cnt:(cnt + 1)]==end")
		})
		// Total message length is remlen + 1 (msg type) + m (remlen bytes)

		remlen_ := int64(remlen)
		total = remlen_ + int64(1) + int64(m)
		Log.Infoc(func() string {
			return fmt.Sprintf("remlen===n===totle: %d===%d===%d", remlen, m, total)
		})
		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, 0, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		Log.Infoc(func() string {
			return fmt.Sprintf("len(b):%d", len(b))
		})
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		if this.WriteBuffer(b) {
			return total, err
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, index, ok := this.ReadBuffer()
		defer this.ReadCommit(index)
		if !ok {
			return total, io.EOF
		}

		Log.Debugc(func() string {
			return fmt.Sprintf("defer this.ReadCommit(%s)", index)
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(p) > 0 {
			n, err := w.Write(p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<|endoftext|>


<commit_before>package generate

const modelTemplate = `package {{.package_name}}

import (
	{{range $i := .model.Imports -}}
	"{{$i}}"
	{{end}}
)

type {{.model_name}} struct {
	{{range $a := .model.Attributes -}}
	{{$a}}
	{{end -}}
}

// String is not required by pop and may be deleted
func ({{.char}} {{.model_name}}) String() string {
	j{{.char}}, _ := json.Marshal({{.char}})
	return string(j{{.char}})
}

// {{.plural_model_name}} is not required by pop and may be deleted
type {{.plural_model_name}} []{{.model_name}}

// String is not required by pop and may be deleted
func ({{.char}} {{.plural_model_name}}) String() string {
	j{{.char}}, _ := json.Marshal({{.char}})
	return string(j{{.char}})
}

// Validate gets run everytime you call a "pop.Validate" method.
// This method is not required and may be deleted.
func ({{.char}} *{{.model_name}}) Validate(tx *pop.Connection) (*validate.Errors, error) {
	{{ if .model.ValidatableAttributes -}}
	return validate.Validate(
		{{ range $a := .model.ValidatableAttributes -}}
		&validators.{{capitalize $a.GoType}}IsPresent{Field: {{$.char}}.{{$a.Names.Proper}}, Name: "{{$a.Names.Proper}}"},
		{{end -}}
	), nil
	{{ else -}}
		return validate.NewErrors(), nil
	{{ end -}}
}

// ValidateSave gets run everytime you call "pop.ValidateSave" method.
// This method is not required and may be deleted.
func ({{.char}} *{{.model_name}}) ValidateSave(tx *pop.Connection) (*validate.Errors, error) {
	return validate.NewErrors(), nil
}

// ValidateUpdate gets run everytime you call "pop.ValidateUpdate" method.
// This method is not required and may be deleted.
func ({{.char}} *{{.model_name}}) ValidateUpdate(tx *pop.Connection) (*validate.Errors, error) {
	return validate.NewErrors(), nil
}
`

const modelTestTemplate = `package {{.package_name}}_test

import "testing"

func Test_{{.model_name}}(t *testing.T) {
	t.Fatal("This test needs to be implemented!")
}
`
<commit_msg>model generation fails to compile fixes #32<commit_after>package generate

const modelTemplate = `package {{.package_name}}

import (
	{{ range $i := .model.Imports -}}
	"{{$i}}"
	{{ end -}}
	{{ if .model.ValidatableAttributes -}}
	"github.com/markbates/validate/validators"
	{{ end -}}
)

type {{.model_name}} struct {
	{{range $a := .model.Attributes -}}
	{{$a}}
	{{end -}}
}

// String is not required by pop and may be deleted
func ({{.char}} {{.model_name}}) String() string {
	j{{.char}}, _ := json.Marshal({{.char}})
	return string(j{{.char}})
}

// {{.plural_model_name}} is not required by pop and may be deleted
type {{.plural_model_name}} []{{.model_name}}

// String is not required by pop and may be deleted
func ({{.char}} {{.plural_model_name}}) String() string {
	j{{.char}}, _ := json.Marshal({{.char}})
	return string(j{{.char}})
}

// Validate gets run everytime you call a "pop.Validate" method.
// This method is not required and may be deleted.
func ({{.char}} *{{.model_name}}) Validate(tx *pop.Connection) (*validate.Errors, error) {
	{{ if .model.ValidatableAttributes -}}
	return validate.Validate(
		{{ range $a := .model.ValidatableAttributes -}}
		&validators.{{capitalize $a.GoType}}IsPresent{Field: {{$.char}}.{{$a.Names.Proper}}, Name: "{{$a.Names.Proper}}"},
		{{end -}}
	), nil
	{{ else -}}
		return validate.NewErrors(), nil
	{{ end -}}
}

// ValidateSave gets run everytime you call "pop.ValidateSave" method.
// This method is not required and may be deleted.
func ({{.char}} *{{.model_name}}) ValidateSave(tx *pop.Connection) (*validate.Errors, error) {
	return validate.NewErrors(), nil
}

// ValidateUpdate gets run everytime you call "pop.ValidateUpdate" method.
// This method is not required and may be deleted.
func ({{.char}} *{{.model_name}}) ValidateUpdate(tx *pop.Connection) (*validate.Errors, error) {
	return validate.NewErrors(), nil
}
`

const modelTestTemplate = `package {{.package_name}}_test

import "testing"

func Test_{{.model_name}}(t *testing.T) {
	t.Fatal("This test needs to be implemented!")
}
`
<|endoftext|>


<commit_before>package gorums

import "fmt"

type ConfigOption interface{}

type NodeListOption interface {
	ConfigOption
	newConfig(*Manager) (Configuration, error)
}

type nodeIDMap struct {
	idMap map[string]uint32
}

func (o nodeIDMap) newConfig(mgr *Manager) (nodes Configuration, err error) {
	if len(o.idMap) == 0 {
		return nil, ConfigCreationError(fmt.Errorf("node-to-ID map required: WithNodeMap"))
	}
	nodes = make(Configuration, 0)
	for naddr, id := range o.idMap {
		node, found := mgr.Node(id)
		if !found {
			node, err = NewNodeWithID(naddr, id)
			if err != nil {
				return nil, ConfigCreationError(err)
			}
			err = mgr.AddNode(node)
			if err != nil {
				return nil, ConfigCreationError(err)
			}
		}
		nodes = append(nodes, node)
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, nil
}

// WithNodeMap returns a ConfigOption containing the provided
// mapping from node addresses to application-specific IDs.
func WithNodeMap(idMap map[string]uint32) NodeListOption {
	return &nodeIDMap{idMap: idMap}
}

type nodeList struct {
	addrsList []string
}

func (o nodeList) newConfig(mgr *Manager) (Configuration, error) {
	if len(o.addrsList) == 0 {
		return nil, ConfigCreationError(fmt.Errorf("node addresses required: WithNodeList"))
	}
	nodes := make(Configuration, 0)
	for _, naddr := range o.addrsList {
		node, err := NewNode(naddr)
		if err != nil {
			return nil, ConfigCreationError(err)
		}
		if n, found := mgr.Node(node.ID()); !found {
			err = mgr.AddNode(node)
			if err != nil {
				return nil, ConfigCreationError(err)
			}
		} else {
			node = n
		}
		nodes = append(nodes, node)
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, nil
}

// WithNodeList returns a ConfigOption containing the provided list of node addresses.
// With this option, node IDs are generated by the Manager.
func WithNodeList(addrsList []string) NodeListOption {
	return &nodeList{addrsList: addrsList}
}

type nodeIDs struct {
	nodeIDs []uint32
}

func (o nodeIDs) newConfig(mgr *Manager) (Configuration, error) {
	if len(o.nodeIDs) == 0 {
		return nil, ConfigCreationError(fmt.Errorf("node IDs required: WithNodeIDs"))
	}
	nodes := make(Configuration, 0)
	for _, id := range o.nodeIDs {
		node, found := mgr.Node(id)
		if !found {
			// Node IDs must have been registered previously
			return nil, ConfigCreationError(fmt.Errorf("node ID %d not found", id))
		}
		nodes = append(nodes, node)
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, nil
}

// WithNodeIDs returns a ConfigOption containing a list of node IDs.
// This assumes that the provided node IDs have already been registered with the manager.
func WithNodeIDs(ids []uint32) NodeListOption {
	return &nodeIDs{nodeIDs: ids}
}

type addNodes struct {
	old Configuration
	new NodeListOption
}

func (o addNodes) newConfig(mgr *Manager) (nodes Configuration, err error) {
	newNodes, err := o.new.newConfig(mgr)
	if err != nil {
		return nil, err
	}
	ac := &addConfig{old: o.old, add: newNodes}
	return ac.newConfig(mgr)
}

func (c Configuration) AddNodes(new NodeListOption) NodeListOption {
	return &addNodes{old: c, new: new}
}

type addConfig struct {
	old Configuration
	add Configuration
}

func (o addConfig) newConfig(mgr *Manager) (nodes Configuration, err error) {
	nodes = make(Configuration, 0)
	m := make(map[uint32]bool)
	for _, n := range append(o.old, o.add...) {
		if !m[n.id] {
			m[n.id] = true
			nodes = append(nodes, n)
		}
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, err
}

func (c Configuration) Add(add Configuration) NodeListOption {
	return &addConfig{old: c, add: add}
}
<commit_msg>Added doc comments for Add and AddNodes<commit_after>package gorums

import "fmt"

type ConfigOption interface{}

type NodeListOption interface {
	ConfigOption
	newConfig(*Manager) (Configuration, error)
}

type nodeIDMap struct {
	idMap map[string]uint32
}

func (o nodeIDMap) newConfig(mgr *Manager) (nodes Configuration, err error) {
	if len(o.idMap) == 0 {
		return nil, ConfigCreationError(fmt.Errorf("node-to-ID map required: WithNodeMap"))
	}
	nodes = make(Configuration, 0)
	for naddr, id := range o.idMap {
		node, found := mgr.Node(id)
		if !found {
			node, err = NewNodeWithID(naddr, id)
			if err != nil {
				return nil, ConfigCreationError(err)
			}
			err = mgr.AddNode(node)
			if err != nil {
				return nil, ConfigCreationError(err)
			}
		}
		nodes = append(nodes, node)
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, nil
}

// WithNodeMap returns a ConfigOption containing the provided
// mapping from node addresses to application-specific IDs.
func WithNodeMap(idMap map[string]uint32) NodeListOption {
	return &nodeIDMap{idMap: idMap}
}

type nodeList struct {
	addrsList []string
}

func (o nodeList) newConfig(mgr *Manager) (Configuration, error) {
	if len(o.addrsList) == 0 {
		return nil, ConfigCreationError(fmt.Errorf("node addresses required: WithNodeList"))
	}
	nodes := make(Configuration, 0)
	for _, naddr := range o.addrsList {
		node, err := NewNode(naddr)
		if err != nil {
			return nil, ConfigCreationError(err)
		}
		if n, found := mgr.Node(node.ID()); !found {
			err = mgr.AddNode(node)
			if err != nil {
				return nil, ConfigCreationError(err)
			}
		} else {
			node = n
		}
		nodes = append(nodes, node)
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, nil
}

// WithNodeList returns a ConfigOption containing the provided list of node addresses.
// With this option, node IDs are generated by the Manager.
func WithNodeList(addrsList []string) NodeListOption {
	return &nodeList{addrsList: addrsList}
}

type nodeIDs struct {
	nodeIDs []uint32
}

func (o nodeIDs) newConfig(mgr *Manager) (Configuration, error) {
	if len(o.nodeIDs) == 0 {
		return nil, ConfigCreationError(fmt.Errorf("node IDs required: WithNodeIDs"))
	}
	nodes := make(Configuration, 0)
	for _, id := range o.nodeIDs {
		node, found := mgr.Node(id)
		if !found {
			// Node IDs must have been registered previously
			return nil, ConfigCreationError(fmt.Errorf("node ID %d not found", id))
		}
		nodes = append(nodes, node)
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, nil
}

// WithNodeIDs returns a ConfigOption containing a list of node IDs.
// This assumes that the provided node IDs have already been registered with the manager.
func WithNodeIDs(ids []uint32) NodeListOption {
	return &nodeIDs{nodeIDs: ids}
}

type addNodes struct {
	old Configuration
	new NodeListOption
}

func (o addNodes) newConfig(mgr *Manager) (nodes Configuration, err error) {
	newNodes, err := o.new.newConfig(mgr)
	if err != nil {
		return nil, err
	}
	ac := &addConfig{old: o.old, add: newNodes}
	return ac.newConfig(mgr)
}

// Add returns a NodeListOption that can be used to create a new configuration combining c
// and the nodes provided via the NodeListOption new.
func (c Configuration) AddNodes(new NodeListOption) NodeListOption {
	return &addNodes{old: c, new: new}
}

type addConfig struct {
	old Configuration
	add Configuration
}

func (o addConfig) newConfig(mgr *Manager) (nodes Configuration, err error) {
	nodes = make(Configuration, 0)
	m := make(map[uint32]bool)
	for _, n := range append(o.old, o.add...) {
		if !m[n.id] {
			m[n.id] = true
			nodes = append(nodes, n)
		}
	}
	// Sort nodes to ensure deterministic iteration.
	OrderedBy(ID).Sort(mgr.nodes)
	OrderedBy(ID).Sort(nodes)
	return nodes, err
}

// Add returns a NodeListOption that can be used to create a new configuration combining c and add.
func (c Configuration) Add(add Configuration) NodeListOption {
	return &addConfig{old: c, add: add}
}
<|endoftext|>